Document new TCP stream classes

cmake/appveyor.yml

@@ -37,7 +37,7 @@ after_build:
 - 7z a libtins-%platform%-%Configuration%.zip libtins
 test_script:
 - cd c:\projects\libtins\build
-- ctest -C %Configuration%
+- ctest -C %Configuration% -V
 deploy_script:
 - ps: Push-AppveyorArtifact "install\libtins-$env:Platform-$env:Configuration.zip"
 skip_commits:

examples/http_dump.cpp

@@ -89,10 +89,6 @@ void on_client_data(Stream& stream) {
     // Now print it, prepending some information about the stream
     cout << client_endpoint(stream) << " >> " 
          << server_endpoint(stream) << ": " << endl << data << endl;
-    // Now erase the stored data, as we've already processed it. This is important,
-    // since if we don't do this, the connection will keep buffering data until
-    // the stream is closed
-    stream.client_payload().clear();
 }
 
 // Whenever there's new server data on the stream, this callback is executed.
@@ -101,7 +97,6 @@ void on_server_data(Stream& stream) {
     string data(stream.server_payload().begin(), stream.server_payload().end());
     cout << server_endpoint(stream) << " >> " 
          << client_endpoint(stream) << ": " << endl << data << endl;
-    stream.server_payload().clear();
 }
 
 // When a connection is closed, this callback is executed.

include/tins/exceptions.h

@@ -273,6 +273,16 @@ public:
     }
 };
 
+/**
+ * \brief Exception thrown when a stream is not found
+ */
+class stream_not_found : public exception_base {
+public:
+    const char* what() const throw() {
+        return "Stream not found";
+    }
+};
+
 namespace Crypto {
 namespace WPA2 {
     /**

include/tins/tcp_ip.h

@@ -51,8 +51,29 @@ class IPv6Address;
 
 namespace TCPIP {
 
+/**
+ * \brief Represents an unidirectional TCP flow between 2 endpoints
+ *
+ * This class will keep the state for all the traffic sent by
+ * one of the peers in a TCP connection. This contains the sequence number,
+ * payload ready to be read and buffered payload, along with some other
+ * properties of the flow.
+ *
+ * A TCP stream (see class Stream) is made out of 2 Flows, so you should 
+ * probably have a look at that class first.
+ *
+ * You shouldn't normally need to interact with this class. Stream already
+ * provides proxys to most of its Flow's attributes.
+ */
 class TINS_API Flow {
 public:
+    /**
+     * \brief Enum that indicates the state of this flow.
+     * 
+     * Note that although similar, this is not mapped to a TCP state-machine 
+     * state. This is mostly used internally to know which packets the flow is
+     * expecting and to know when it's done sending data.
+     */
     enum State {
         UNKNOWN,
         SYN_SENT,
@@ -61,35 +82,154 @@ public:
         RST_SENT
     };
 
+    /** 
+     * The type used to store the payload
+     */
     typedef std::vector<uint8_t> payload_type;
+
+    /**
+     * The type used to store the buffered payload
+     */
     typedef std::map<uint32_t, payload_type> buffered_payload_type;
+
+    /**
+     * The type used to store the callbacks that this class triggers
+     */
     typedef std::function<void(Flow&)> event_callback;
 
-    Flow(const IPv4Address& dest_address, uint16_t dest_port,
+    /** 
+     * Construct a Flow from an IPv4 address
+     *
+     * \param dst_address This flow's destination address
+     * \param dst_port This flow's destination port
+     * \param sequence_number The initial sequence number to be used 
+     */
+    Flow(const IPv4Address& dst_address, uint16_t dst_port,
          uint32_t sequence_number);
     
-    Flow(const IPv6Address& dest_address, uint16_t dest_port,
+    /** 
+     * Construct a Flow from an IPv6 address
+     *
+     * \param dst_address This flow's destination address
+     * \param dst_port This flow's destination port
+     * \param sequence_number The initial sequence number to be used 
+     */
+    Flow(const IPv6Address& dst_address, uint16_t dst_port,
          uint32_t sequence_number);
 
+    /**
+     * \brief Sets the callback that will be executed when data is readable
+     *
+     * Whenever this flow has readable data, this callback will be executed.
+     * By readable, this means that there's non-out-of-order data captured.
+     *
+     * \param callback The callback to be executed   
+     */
     void data_callback(const event_callback& callback);
+
+    /**
+     * \brief Sets the callback that will be executed when data is buffered.
+     *
+     * Whenever this flow receives out-of-order data, this callback will be
+     * executed.
+     * 
+     * \param callback The callback to be executed
+     */
     void buffering_callback(const event_callback& callback);
 
+    /**
+     * \brief Processes a packet.
+     *
+     * If this packet contains data and starts or overlaps with the current
+     * sequence number, then the data will be appended to this flow's payload
+     * and the data_callback will be executed.
+     *
+     * If this packet contains out-of-order data, it will be buffered and the
+     * buffering_callback will be executed.
+     *
+     * \param pdu The packet to be processed
+     * \sa Flow::data_callback
+     * \sa Flow::buffering_callback
+     */
     void process_packet(PDU& pdu);
 
+    /**
+     * Indicates whether this flow uses IPv6 addresses
+     */
     bool is_v6() const;
+
+    /**
+     * \brief Indicates whether this flow is finished
+     *
+     * A finished is considered to be finished if either it sent a
+     * packet with the FIN or RST flags on. 
+     */
     bool is_finished() const;
+
+    /**
+     * \brief Indicates whether a packet belongs to this flow
+     *
+     * Since Flow represents a unidirectional stream, this will only check
+     * the destination endpoint and not the source one.
+     *
+     * \param packet The packet to be checked
+     */
     bool packet_belongs(const PDU& packet) const;
 
+    /**
+     * \brief Getter for the IPv4 destination address
+     *
+     * Note that it's only safe to execute this method if is_v6() == false
+     */
     IPv4Address dst_addr_v4() const;
+
+    /**
+     * \brief Getter for the IPv6 destination address
+     *
+     * Note that it's only safe to execute this method if is_v6() == true
+     */
     IPv6Address dst_addr_v6() const;
+
+    /** 
+     * Getter for this flow's destination port
+     */
     uint16_t dport() const;
+
+    /** 
+     * Getter for this flow's payload (const)
+     */
     const payload_type& payload() const;
+
+    /** 
+     * Getter for this flow's destination port
+     */
     payload_type& payload();
+
+    /** 
+     * Getter for this flow's state
+     */
     State state() const;
+
+    /** 
+     * Getter for this flow's sequence number
+     */
     uint32_t sequence_number() const;
+
+    /** 
+     * Getter for this flow's buffered payload (const)
+     */
     const buffered_payload_type& buffered_payload() const;
+
+    /** 
+     * Getter for this flow's buffered payload
+     */
     buffered_payload_type& buffered_payload();
 
+    /**
+     * Sets the state of this flow
+     *
+     * \param new_state The new state of this flow
+     */
     void state(State new_state);
 private:
     void store_payload(uint32_t seq, const payload_type& payload);
@@ -103,57 +243,208 @@ private:
     uint16_t dest_port_;
     event_callback on_data_callback_;
     event_callback on_buffering_callback_;
-    bool is_v6_;
     State state_;
+    bool is_v6_;
 };
 
+/** 
+ * \brief Represents a TCP stream
+ *
+ * A TCP stream is made out of 2 Flows, one in each direction, plus 
+ * some other attributes and callbacks.
+ *
+ * This class works using callbacks. Whenever the stream is created, you should
+ * set at least the client/server callbacks so you are notified whenever the
+ * client/server has sent data. Note that setting these is not mandatory, so
+ * you can subscribe to just the callbacks you need.
+ *
+ * \sa Stream::auto_cleanup_payloads
+ */
 class TINS_API Stream {
 public:
-    enum State {
-        SYN_SENT,
-        SYN_RCVD,
-        ESTABLISHED,
-        CLOSE_WAIT,
-        FIN_WAIT_1,
-        FIN_WAIT_2,
-        TIME_WAIT,
-        CLOSED
-    };
-
+    /**
+     * The type used for callbacks
+     */
     typedef std::function<void(Stream&)> stream_callback;
+
+    /**
+     * The type used to store payloads
+     */
     typedef Flow::payload_type payload_type;
 
+    /**
+     * \brief Constructs a TCP stream using the provided packet.
+     */
     Stream(const PDU& initial_packet);
-    Stream(const Flow& client_flow, const Flow& server_flow);
 
+    /**
+     * \brief Processes this packet.
+     *
+     * This will forward the packet appropriately to the client
+     * or server flow.
+     */
     void process_packet(PDU& packet);
 
+    /**
+     * Getter for the client flow
+     */
     Flow& client_flow();
+
+    /**
+     * Getter for the client flow (const)
+     */
     const Flow& client_flow() const;
+
+    /**
+     * Getter for the server flow
+     */
     Flow& server_flow();
+
+    /**
+     * Getter for the server flow (const)
+     */
     const Flow& server_flow() const;
 
+    /**
+     * \brief Indicates whether this stream is finished.
+     *
+     * This stream is finished if either peer sent a packet with
+     * the RST flag on, or both peers sent a FIN.
+     */
     bool is_finished() const;
+
+    /**
+     * Indicates whether this packet uses IPv6 addresses
+     */
     bool is_v6() const;
 
+    /**
+     * \brief Retrieves the client's IPv4 address
+     *
+     * Note that it's only valid to call this method if is_v6() == false
+     */
     IPv4Address client_addr_v4() const;
+
+    /**
+     * \brief Retrieves the client's IPv6 address
+     *
+     * Note that it's only valid to call this method if is_v6() == true
+     */
     IPv6Address client_addr_v6() const;
+
+    /**
+     * \brief Retrieves the server's IPv4 address
+     *
+     * Note that it's only valid to call this method if is_v6() == false
+     */
     IPv4Address server_addr_v4() const;
+
+    /**
+     * \brief Retrieves the server's IPv6 address
+     *
+     * Note that it's only valid to call this method if is_v6() == true
+     */
     IPv6Address server_addr_v6() const;
+
+    /**
+     * Getter for the client's port
+     */
     uint16_t client_port() const;
+
+    /**
+     * Getter for the server's port
+     */
     uint16_t server_port() const;
+
+    /**
+     * Getter for the client's payload (const)
+     */
     const payload_type& client_payload() const;
+
+    /**
+     * Getter for the client's payload
+     */
     payload_type& client_payload();
+
+    /**
+     * Getter for the server's payload (const)
+     */
     const payload_type& server_payload() const;
+
+    /**
+     * Getter for the server's payload
+     */
     payload_type& server_payload();
 
+    /**
+     * \brief Sets the callback to be executed when the stream is closed
+     *
+     * \param callback The callback to be set
+     */
     void stream_closed_callback(const stream_callback& callback);
+
+    /**
+     * \brief Sets the callback to be executed when there's client data
+     *
+     * \sa Flow::data_callback
+     * \param callback The callback to be set
+     */
     void client_data_callback(const stream_callback& callback);
+
+    /**
+     * \brief Sets the callback to be executed when there's server data
+     *
+     * \sa Flow::data_callback
+     * \param callback The callback to be set
+     */
     void server_data_callback(const stream_callback& callback);
+
+    /**
+     * \brief Sets the callback to be executed when there's new buffered 
+     * client data
+     *
+     * \sa Flow::buffering_callback
+     * \param callback The callback to be set
+     */
     void client_buffering_callback(const stream_callback& callback);
+
+    /**
+     * \brief Sets the callback to be executed when there's new buffered 
+     * client data
+     *
+     * \sa Flow::buffering_callback
+     * \param callback The callback to be set
+     */
     void server_buffering_callback(const stream_callback& callback);
 
+    /**
+     * \brief Sets the internal callbacks. 
+     *
+     * This shouldn't normally need to be called except if you're constructing
+     * this object and then moving it around before persisting it somewhere.  
+     */
     void setup_flows_callbacks();
+
+    /**
+     * \brief Indicates whether each flow's payloads should be automatically
+     * erased.
+     *
+     * If this property is true, then whenever there's new data for a stream,
+     * the appropriate callback will be executed and then the payload will be 
+     * erased. 
+     *
+     * If this property is false, then the payload <b>will not</b> be erased
+     * and the user is responsible for clearing the payload vector. 
+     *
+     * Setting this property to false is useful if it's desired to hold all 
+     * of the data sent on the stream before processing it. Note that this
+     * can lead to the memory growing a lot.
+     *
+     * This property is true by default. 
+     *
+     * \param value The value to be set for this property
+     */
+    void auto_cleanup_payloads(bool value);
 private:
     static Flow extract_client_flow(const PDU& packet);
     static Flow extract_server_flow(const PDU& packet);
@@ -170,20 +461,91 @@ private:
     stream_callback on_server_data_callback_;
     stream_callback on_client_buffering_callback_;
     stream_callback on_server_buffering_callback_;
-    State state_;
+    bool auto_cleanup_;
 };
 
+/**
+ * \brief Represents a class that follows TCP and reassembles streams
+ *
+ * This class processes packets and whenever it detects a new connection
+ * being open, it starts tracking it. This will follow all data sent by
+ * each peer and make it available to the user in a simple way.
+ *
+ * In order to use this class, just create an instance and set the 
+ * new stream callback to some function that you want:
+ *
+ * \code
+ * void on_new_stream(TCPStream& stream) {
+ *     // Do something with it.
+ *     // This is the perfect time to set the stream's client/server 
+ *     // write callbacks so you are notified whenever there's new
+ *     // data on the stream
+ * }
+ *
+ * // Create it 
+ * StreamFollower follower;
+ * // Set the callback
+ * follower.new_stream_callback(&on_new_stream);
+ * \endcode
+ */
 class TINS_API StreamFollower {
 public:
+    /**
+     * \brief The type used for callbacks
+     */
     typedef Stream::stream_callback stream_callback;
 
+    /** 
+     * Default constructor
+     */
     StreamFollower();
 
+    /** 
+     * \brief Processes a packet
+     *
+     * This will detect if this packet belongs to an existing stream 
+     * and process it, or if it belongs to a new one, in which case it
+     * starts tracking it.
+     *
+     * This method always returns true so it can be easily plugged as
+     * the argument to Sniffer::sniff_loop.
+     *
+     * \param packet The packet to be processed
+     * \return Always true
+     */
     bool process_packet(PDU& packet);
+
+    /**
+     * \brief Sets the callback to be executed when a new stream is captured.
+     *
+     * Whenever a new stream is captured, the provided callback will be 
+     * executed.
+     *
+     * \param callback The callback to be set
+     */
     void new_stream_callback(const stream_callback& callback);
 
+    /**
+     * Finds the stream identified by the provided arguments.
+     *
+     * \param client_addr The client's address
+     * \param client_port The client's port
+     * \param server_addr The server's address
+     * \param server_addr The server's port
+     */
     Stream& find_stream(IPv4Address client_addr, uint16_t client_port,
                         IPv4Address server_addr, uint16_t server_port);
+
+    /**
+     * Finds the stream identified by the provided arguments.
+     *
+     * \param client_addr The client's address
+     * \param client_port The client's port
+     * \param server_addr The server's address
+     * \param server_addr The server's port
+     */
+    Stream& find_stream(IPv6Address client_addr, uint16_t client_port,
+                        IPv6Address server_addr, uint16_t server_port);
 private:
     static const size_t DEFAULT_MAX_BUFFERED_CHUNKS;
     typedef std::array<uint8_t, 16> address_type;
@@ -205,6 +567,7 @@ private:
     typedef std::map<stream_id, Stream> streams_type;
 
     stream_id make_stream_id(const PDU& packet);
+    Stream& find_stream(const stream_id& id);
     static address_type serialize(IPv4Address address);
     static address_type serialize(const IPv6Address& address);
 

src/tcp_ip.cpp

@@ -72,16 +72,16 @@ int seq_compare(uint32_t seq1, uint32_t seq2) {
 
 Flow::Flow(const IPv4Address& dest_address, uint16_t dest_port,
            uint32_t sequence_number) 
-: seq_number_(sequence_number), dest_port_(dest_port), is_v6_(false),
-  state_(UNKNOWN) {
+: seq_number_(sequence_number), dest_port_(dest_port), state_(UNKNOWN),
+  is_v6_(false) {
     OutputMemoryStream output(dest_address_.data(), dest_address_.size());
     output.write(dest_address);
 }
 
 Flow::Flow(const IPv6Address& dest_address, uint16_t dest_port,
            uint32_t sequence_number) 
-: seq_number_(sequence_number), dest_port_(dest_port), is_v6_(true),
-  state_(UNKNOWN) {
+: seq_number_(sequence_number), dest_port_(dest_port), state_(UNKNOWN),
+  is_v6_(true) {
     OutputMemoryStream output(dest_address_.data(), dest_address_.size());
     output.write(dest_address);
 }
@@ -283,13 +283,12 @@ void Flow::state(State new_state) {
 
 Stream::Stream(const PDU& packet) 
 : client_flow_(extract_client_flow(packet)),
-server_flow_(extract_server_flow(packet)) {
-
+  server_flow_(extract_server_flow(packet)), auto_cleanup_(true) {
+    const TCP& tcp = packet.rfind_pdu<TCP>();
+    // If it's a SYN, set the proper state
+    if (tcp.flags() == TCP::SYN) {
+        client_flow().state(Flow::SYN_SENT);
     }
-
-Stream::Stream(const Flow& client_flow, const Flow& server_flow) 
-: client_flow_(client_flow), server_flow_(server_flow) {
-
 }
 
 void Stream::process_packet(PDU& packet) {
@@ -441,17 +440,26 @@ void Stream::setup_flows_callbacks() {
     server_flow_.buffering_callback(bind(&Stream::on_server_buffering, this, _1));
 }
 
+void Stream::auto_cleanup_payloads(bool value) {
+    auto_cleanup_ = value;
+}
 
-void Stream::on_client_flow_data(const Flow& flow) {
+void Stream::on_client_flow_data(const Flow& /*flow*/) {
     if (on_client_data_callback_) {
         on_client_data_callback_(*this);
     }
+    if (auto_cleanup_) {
+        client_payload().clear();
+    }
 }
 
-void Stream::on_server_flow_data(const Flow& flow) {
+void Stream::on_server_flow_data(const Flow& /*flow*/) {
     if (on_server_data_callback_) {
         on_server_data_callback_(*this);
     }
+    if (auto_cleanup_) {
+        server_payload().clear();
+    }
 }
 
 void Stream::on_client_buffering(const Flow& flow) {
@@ -494,8 +502,6 @@ bool StreamFollower::process_packet(PDU& packet) {
                 throw runtime_error("No new connection callback set");
             }
             if (tcp.flags() == TCP::SYN) {
-                // If it's a SYN, set the proper state
-                iter->second.client_flow().state(Flow::SYN_SENT);
                 process = false;
             }
             else {
@@ -530,14 +536,14 @@ Stream& StreamFollower::find_stream(IPv4Address client_addr, uint16_t client_por
                                     IPv4Address server_addr, uint16_t server_port) {
     stream_id identifier(serialize(client_addr), client_port,
                          serialize(server_addr), server_port);
-    streams_type::iterator iter = streams_.find(identifier);
-    if (iter == streams_.end()) {
-        // TODO: define proper exception
-        throw runtime_error("Stream not found");
-    }
-    else {
-        return iter->second;
+    return find_stream(identifier);
 }
+
+Stream& StreamFollower::find_stream(IPv6Address client_addr, uint16_t client_port,
+                                    IPv6Address server_addr, uint16_t server_port) {
+    stream_id identifier(serialize(client_addr), client_port,
+                         serialize(server_addr), server_port);
+    return find_stream(identifier);
 }
 
 StreamFollower::stream_id StreamFollower::make_stream_id(const PDU& packet) {
@@ -560,6 +566,16 @@ StreamFollower::stream_id StreamFollower::make_stream_id(const PDU& packet) {
     }
 }
 
+Stream& StreamFollower::find_stream(const stream_id& id) {
+    streams_type::iterator iter = streams_.find(id);
+    if (iter == streams_.end()) {
+        throw stream_not_found();
+    }
+    else {
+        return iter->second;
+    }
+}
+
 StreamFollower::address_type StreamFollower::serialize(IPv4Address address) {
     address_type addr;
     OutputMemoryStream output(addr.data(), addr.size());

tests/src/tcp_ip.cpp

@@ -13,6 +13,7 @@
 #include "ip.h"
 #include "ip_address.h"
 #include "ipv6_address.h"
+#include "exceptions.h"
 #include "ethernetII.h"
 #include "rawpdu.h"
 #include "utils.h"
@@ -120,12 +121,10 @@ void FlowTest::on_new_stream(Stream& stream) {
 
 void FlowTest::cumulative_stream_client_data_handler(Stream& stream) {
     stream_client_payload_chunks.push_back(stream.client_flow().payload());
-    stream.client_flow().payload().clear();
 }
 
 void FlowTest::cumulative_stream_server_data_handler(Stream& stream) {
     stream_server_payload_chunks.push_back(stream.server_flow().payload());
-    stream.server_flow().payload().clear();
 }
 
 void FlowTest::buffered_payload_handle(Flow& session) {
@@ -286,7 +285,8 @@ TEST_F(FlowTest, StreamFollower_ThreeWayHandshake) {
     for (size_t i = 0; i < packets.size(); ++i) {
         follower.process_packet(packets[i]);
     }
-    Stream& stream = follower.find_stream("1.2.3.4", 22, "4.3.2.1", 25);
+    Stream& stream = follower.find_stream(IPv4Address("1.2.3.4"), 22,
+                                          IPv4Address("4.3.2.1"), 25);
     EXPECT_EQ(Flow::ESTABLISHED, stream.client_flow().state());
     EXPECT_EQ(Flow::SYN_SENT, stream.server_flow().state());
     EXPECT_EQ(30, stream.client_flow().sequence_number());
@@ -317,7 +317,8 @@ TEST_F(FlowTest, StreamFollower_RSTClosesStream) {
     for (size_t i = 0; i < packets.size(); ++i) {
         follower.process_packet(packets[i]);
     }
-    Stream stream = follower.find_stream("1.2.3.4", 22, "4.3.2.1", 25);
+    Stream stream = follower.find_stream(IPv4Address("1.2.3.4"), 22,
+                                         IPv4Address("4.3.2.1"), 25);
 
     IP server_packet = IP("1.2.3.4", "4.3.2.1") / TCP(22, 25);
     server_packet.rfind_pdu<TCP>().flags(TCP::RST);
@@ -336,7 +337,8 @@ TEST_F(FlowTest, StreamFollower_FINClosesStream) {
     for (size_t i = 0; i < packets.size(); ++i) {
         follower.process_packet(packets[i]);
     }
-    Stream stream = follower.find_stream("1.2.3.4", 22, "4.3.2.1", 25);
+    Stream stream = follower.find_stream(IPv4Address("1.2.3.4"), 22,
+                                         IPv4Address("4.3.2.1"), 25);
 
     IP server_packet = IP("1.2.3.4", "4.3.2.1") / TCP(22, 25);
     server_packet.rfind_pdu<TCP>().flags(TCP::FIN | TCP::ACK);
@@ -367,7 +369,10 @@ TEST_F(FlowTest, StreamFollower_StreamIsRemovedWhenFinished) {
     follower.process_packet(server_packet);
 
     // We shouldn't be able to find it
-    EXPECT_THROW(follower.find_stream("1.2.3.4", 22, "4.3.2.1", 25), runtime_error);
+    EXPECT_THROW(
+        follower.find_stream(IPv4Address("1.2.3.4"), 22, IPv4Address("4.3.2.1"), 25), 
+        stream_not_found
+    );
 }
 
 TEST_F(FlowTest, StreamFollower_FollowStream) {