Line data    Source code

       1             : /*
       2             :  * Copyright (c) 2013 Juniper Networks, Inc. All rights reserved.
       3             :  */
       4             : 
       5             : #ifndef SRC_IO_TCP_SESSION_H_
       6             : #define SRC_IO_TCP_SESSION_H_
       7             : 
       8             : #include <deque>
       9             : #include <list>
      10             : #include <string>
      11             : #include <mutex>
      12             : #include <atomic>
      13             : 
      14             : #include <boost/asio/buffer.hpp>
      15             : #include <boost/asio/io_service.hpp>
      16             : #include <boost/asio/ip/tcp.hpp>
      17             : #include <boost/asio/strand.hpp>
      18             : #include <boost/intrusive_ptr.hpp>
      19             : #include <boost/function.hpp>
      20             : #include <boost/scoped_ptr.hpp>
      21             : 
      22             : #include "base/util.h"
      23             : #include "base/task.h"
      24             : #include "io/tcp_server.h"
      25             : 
      26             : #define SSL_SHORT_READ_ERROR 335544539
      27             : 
      28             : class EventManager;
      29             : class TcpServer;
      30             : class TcpSession;
      31             : class TcpMessageWriter;
      32             : 
      33             : // TcpSession
      34             : //
      35             : // Concurrency: the session is created by the event manager thread, which
      36             : // also invokes the AsyncHandlers. ReleaseBuffer and Send will typically be
      37             : // invoked by a different thread.
      38             : class TcpSession {
      39             : public:
      40             :     static const int kDefaultBufferSize = 16 * 1024;
      41             :     static const int kDefaultWriteBufferSize = 32 * 1024;
      42             : 
      43             :     enum Event {
      44             :         EVENT_NONE,
      45             :         ACCEPT,
      46             :         CONNECT_COMPLETE,
      47             :         CONNECT_FAILED,
      48             :         CLOSE
      49             :     };
      50             : 
      51             :     enum Direction {
      52             :         ACTIVE,
      53             :         PASSIVE
      54             :     };
      55             : 
      56             :     typedef boost::asio::ip::tcp::socket Socket;
      57             :     typedef boost::asio::ip::tcp::socket::native_handle_type NativeSocketType;
      58             :     typedef boost::asio::ip::tcp::endpoint Endpoint;
      59             :     typedef boost::function<void(TcpSession *, Event)> EventObserver;
      60             :     typedef boost::asio::const_buffer Buffer;
      61             : 
      62             :     // TcpSession constructor takes ownership of socket.
      63             :     TcpSession(TcpServer *server, Socket *socket,
      64             :                bool async_read_ready = true,
      65             :                size_t buffer_send_size = TcpSession::kDefaultWriteBufferSize);
      66             :     // Performs a non-blocking send operation.
      67             :     virtual bool Send(const uint8_t *data, size_t size, size_t *sent);
      68             : 
      69             :     // Called by TcpServer to trigger async read.
      70             :     virtual bool Connected(Endpoint remote);
      71             : 
      72             :     // Called by TcpServer to trigger async read.
      73             :     virtual void Accepted();
      74             : 
      75             :     void ConnectFailed();
      76             : 
      77             :     void Close();
      78             : 
      79      116583 :     virtual std::string ToString() const { return name_; }
      80             : 
      81             :     // Getters and setters
      82      959627 :     virtual Socket *socket() const { return socket_.get(); }
      83             :     NativeSocketType sock_descriptor() { return socket_->native_handle(); }
      84      226718 :     TcpServer *server() { return server_.get(); }
      85             :     int32_t local_port() const;
      86             :     int32_t remote_port() const;
      87             : 
      88             :     // Concurrency: changing the observer guarantees mutual exclusion with
      89             :     // the observer invocation. e.g. if the caller sets the observer to NULL
      90             :     // it is guaranteed that the observer will not get invoked after this
      91             :     // method returns.
      92             :     void set_observer(EventObserver observer);
      93             : 
      94             :     // Buffers must be freed in arrival order.
      95             :     virtual void ReleaseBuffer(Buffer buffer);
      96             : 
      97             :     // This function returns the instance to run SessionTask.
      98             :     // Returning Task::kTaskInstanceAny would allow multiple session tasks to
      99             :     // run in parallel.
     100             :     // Derived class may override implementation if it expects the all the
     101             :     // Tasks of the session to run in specific instance
     102             :     // Note: Two tasks of same task ID and task instance can't run
     103             :     //       at in parallel
     104             :     // E.g. BgpSession is created per BgpPeer and to ensure that
     105             :     //      there is one SessionTask per peer, PeerIndex is returned
     106             :     //      from this function
     107             :     virtual int GetSessionInstance() const;
     108             : 
     109     1572076 :     static const uint8_t *BufferData(const Buffer &buffer) {
     110     1572076 :         return boost::asio::buffer_cast<const uint8_t *>(buffer);
     111             :     }
     112     2714550 :     static size_t BufferSize(const Buffer &buffer) {
     113     2714550 :         return boost::asio::buffer_size(buffer);
     114             :     }
     115             : 
     116     1603469 :     bool IsEstablished() const {
     117     1603469 :         std::scoped_lock lock(mutex_);
     118     1603489 :         return established_;
     119     1603489 :     }
     120             : 
     121       10113 :     bool IsClosed() const {
     122       10113 :         std::scoped_lock lock(mutex_);
     123       10115 :         return closed_;
     124       10115 :     }
     125             : 
     126        2933 :     bool IsServerSession() {
     127        2933 :         if (direction_ == PASSIVE) return true;
     128        1477 :         return false;
     129             :     }
     130             : 
     131     8308997 :     Endpoint remote_endpoint() const {
     132     8308997 :         return remote_;
     133             :     }
     134             : 
     135     1675544 :     const std::string &remote_addr_string() const {
     136     1675544 :         return remote_addr_str_;
     137             :     }
     138             : 
     139             :     Endpoint local_endpoint() const;
     140             : 
     141             :     const boost::system::error_code &close_reason() const {
     142             :         return close_reason_;
     143             :     }
     144             : 
     145             :     virtual boost::system::error_code SetSocketOptions();
     146             :     static bool IsSocketErrorHard(const boost::system::error_code &ec);
     147       14326 :     void set_read_on_connect(bool read) { read_on_connect_ = read; }
     148             :     void SessionEstablished(Endpoint remote, Direction direction);
     149             : 
     150             :     virtual void AsyncReadStart();
     151             :     virtual void SetDeferReader(bool defer_reader);
     152             :     // Is the reader deferred ? If reader is deferred, SetDeferReader needs
     153             :     // to be called to undefer/restart reading.
     154     1319956 :     virtual bool IsReaderDeferred() const {
     155     1319956 :         return defer_reader_;
     156             :     }
     157             : 
     158        1058 :     const io::SocketStats &GetSocketStats() const { return stats_; }
     159             :     void GetRxSocketStats(SocketIOStats *socket_stats) const;
     160             :     void GetTxSocketStats(SocketIOStats *socket_stats) const;
     161             : 
     162           6 :     void GetRxSocketStats(SocketIOStats &socket_stats) const {
     163           6 :         GetRxSocketStats(&socket_stats);
     164           6 :     }
     165             : 
     166           6 :     void GetTxSocketStats(SocketIOStats &socket_stats) const {
     167           6 :         GetTxSocketStats(&socket_stats);
     168           6 :     }
     169             : 
     170             :     int SetMd5SocketOption(uint32_t peer_ip, const std::string &md5_password);
     171             :     int ClearMd5SocketOption(uint32_t peer_ip);
     172             :     int SetDscpSocketOption(uint8_t value);
     173             :     uint8_t GetDscpValue() const;
     174       23812 :     const std::string &ToUVEKey() const { return uve_key_str_; }
     175             :     boost::system::error_code SetTcpNoDelay();
     176             :     boost::system::error_code SetTcpSendBufSize(uint32_t size);
     177             :     boost::system::error_code SetTcpRecvBufSize(uint32_t size);
     178             : 
     179             : protected:
     180             :     typedef boost::intrusive_ptr<TcpSession> TcpSessionPtr;
     181             :     static void AsyncReadHandler(TcpSessionPtr session);
     182             :     static void AsyncWriteHandler(TcpSessionPtr session,
     183             :                                   const boost::system::error_code &error,
     184             :                                   std::size_t bytes_transferred);
     185             : 
     186             :     void AsyncReadStartInternal(TcpSessionPtr session);
     187             :     virtual Task* CreateReaderTask(boost::asio::mutable_buffer, size_t);
     188             : 
     189             :     virtual ~TcpSession();
     190             : 
     191             :     // Read handler. Called from a TBB task.
     192             :     virtual void OnRead(Buffer buffer) = 0;
     193             :     // Callback after socket is ready for write.
     194             :     virtual void WriteReady(const boost::system::error_code &error);
     195             : 
     196             :     void AsyncWriteInternal(TcpSessionPtr session);
     197             : 
     198             :     virtual void AsyncReadSome();
     199             :     virtual size_t GetReadBufferSize() const;
     200             :     virtual size_t ReadSome(boost::asio::mutable_buffer buffer,
     201             :                             boost::system::error_code *error);
     202             :     virtual void AsyncWrite(const uint8_t *data, std::size_t size);
     203             : 
     204     1316355 :     virtual int reader_task_id() const {
     205     1316355 :         return reader_task_id_;
     206             :     }
     207             : 
     208     1228282 :     bool established() const { return established_; }
     209             : 
     210         114 :     EventObserver observer() { return observer_; }
     211             :     boost::system::error_code SetSocketKeepaliveOptions(int keepalive_time,
     212             :             int keepalive_intvl, int keepalive_probes,
     213             :             int tcp_user_timeout_val = 0);
     214             : 
     215             :     void CloseInternal(const boost::system::error_code &ec,
     216             :                        bool call_observer, bool notify_server = true);
     217             : 
     218             :     void TriggerAsyncReadHandler();
     219             : 
     220             :     // Protects session state and buffer queue.
     221             :     mutable std::mutex mutex_;
     222             :     io::SocketStats stats_;
     223             : 
     224             : protected:
     225             :     typedef boost::asio::strand<boost::asio::io_context::executor_type> Strand;
     226             :     boost::scoped_ptr<Strand> io_strand_;
     227             : 
     228             : private:
     229             :     class Reader;
     230             :     friend class TcpServer;
     231             :     friend class TcpMessageWriter;
     232             :     friend void intrusive_ptr_add_ref(TcpSession *session);
     233             :     friend void intrusive_ptr_release(TcpSession *session);
     234             :     typedef std::list<boost::asio::mutable_buffer> BufferQueue;
     235             : 
     236             :     static void WriteReadyInternal(TcpSessionPtr session,
     237             :                                    const boost::system::error_code &error,
     238             :                                    uint64_t block_start_time);
     239             :     void ReleaseBufferLocked(Buffer buffer);
     240             :     void SetEstablished(Endpoint remote, Direction dir);
     241             : 
     242     3445430 :     bool IsClosedLocked() const {
     243     3445430 :         return closed_;
     244             :     }
     245             : 
     246     2161465 :     bool IsEstablishedLocked() const {
     247     2161465 :         return (established_ && !tcp_close_in_progress_);
     248             :     }
     249             : 
     250             :     void SetName();
     251             : 
     252             :     boost::asio::mutable_buffer AllocateBuffer(size_t buffer_size);
     253             :     void DeleteBuffer(boost::asio::mutable_buffer buffer);
     254             : 
     255             :     static int reader_task_id_;
     256             : 
     257             :     TcpServerPtr server_;
     258             :     boost::scoped_ptr<Socket> socket_;
     259             :     bool read_on_connect_;
     260             : 
     261             :     /**************** protected by mutex_ ****************/
     262             :     bool established_;             // In TCP ESTABLISHED state.
     263             :     bool closed_;                  // Close has been called.
     264             :     Endpoint remote_;              // Remote end-point
     265             :     std::string remote_addr_str_;  // Remote end-point address string
     266             :     Direction direction_;          // direction (active, passive)
     267             :     BufferQueue buffer_queue_;
     268             :     boost::system::error_code close_reason_;
     269             :     /**************** end protected by mutex_ ****************/
     270             : 
     271             :     // Protects observer manipulation and invocation. When this lock is
     272             :     // held the session mutex should not be held and vice-versa.
     273             :     std::mutex obs_mutex_;
     274             :     EventObserver observer_;
     275             : 
     276             :     boost::scoped_ptr<TcpMessageWriter> writer_;
     277             : 
     278             :     std::atomic<int> refcount_;
     279             :     std::string name_;
     280             :     std::atomic<bool> defer_reader_;
     281             :     std::string uve_key_str_;
     282             :     std::atomic<bool> write_blocked_;
     283             :     std::atomic<bool> tcp_close_in_progress_;
     284             : 
     285             :     DISALLOW_COPY_AND_ASSIGN(TcpSession);
     286             : };
     287             : 
     288    46380411 : inline void intrusive_ptr_add_ref(TcpSession *session) {
     289    46380411 :     session->refcount_.fetch_add(1);
     290    46380411 : }
     291             : 
     292    46373260 : inline void intrusive_ptr_release(TcpSession *session) {
     293    46373260 :     int prev = session->refcount_.fetch_sub(1);
     294    46373260 :     if (prev == 1) {
     295       30996 :         delete session;
     296             :     }
     297    46373258 : }
     298             : 
     299             : // TcpMessageReader
     300             : //
     301             : // Provides base implementation of OnRead() for TcpSession assuming
     302             : // fixed message header length
     303             : //
     304             : class TcpMessageReader {
     305             : public:
     306             :     typedef boost::asio::const_buffer Buffer;
     307             :     typedef boost::function<bool(const uint8_t *, size_t)> ReceiveCallback;
     308             : 
     309             :     TcpMessageReader(TcpSession *session, ReceiveCallback callback);
     310             :     virtual ~TcpMessageReader();
     311             :     virtual void OnRead(Buffer buffer);
     312             : 
     313             : protected:
     314             :     virtual int MsgLength(Buffer buffer, int offset) = 0;
     315             :     virtual const int GetHeaderLenSize() = 0;
     316             :     virtual const int GetMaxMessageSize() = 0;
     317             : 
     318             : private:
     319             :     typedef std::deque<Buffer> BufferQueue;
     320             : 
     321             :     // Copy the queue into one contiguous buffer.
     322             :     uint8_t *BufferConcat(uint8_t *data, Buffer buffer, int msglength);
     323             : 
     324             :     int QueueByteLength() const;
     325             : 
     326             :     Buffer PullUp(uint8_t *data, Buffer buffer, size_t size) const;
     327             : 
     328             :     int AllocBufferSize(int length);
     329             : 
     330             :     TcpSession *session_;
     331             :     ReceiveCallback callback_;
     332             :     BufferQueue queue_;
     333             :     int offset_;
     334             :     int remain_;
     335             : 
     336             :     DISALLOW_COPY_AND_ASSIGN(TcpMessageReader);
     337             : };
     338             : 
     339             : #endif  // SRC_IO_TCP_SESSION_H_