diff --git a/include/tins/tcp_ip/stream.h b/include/tins/tcp_ip/stream.h index b6363e4..4ce49b1 100644 --- a/include/tins/tcp_ip/stream.h +++ b/include/tins/tcp_ip/stream.h @@ -391,9 +391,41 @@ public: #endif // TINS_HAVE_TCP_STREAM_CUSTOM_DATA /** - * Indicates whether this is a stream that we attached to after it had actually started + * Indicates whether this is a partial stream that we attached to after it had actually started */ - bool is_attached() const; + bool is_partial_stream() const; + + /** + * \brief Enables recovery mode on this stream. + * + * Recovery mode can be used when either a stream is having high packet loss or on partial + * streams. On the latter case, if a stream starts with out of order packets, then the holes + * left by them might never be filled. Enabling recovery mode right after attaching to + * a stream allows automatic recovery so the stream will skip the out of order packets + * and continue tracking the stream by ignoring those holes. + * + * The way recovery mode is, given a recovery window size, it will skip all out of order + * packets that arrive anywhere within the window given by the sequence number at the time of + * enabling recovery mode + the recovery window size. This is, given a stream for which the + * client sequence number is X and a recovery window of size Y, then enabling recovery mode + * at that point will ignore any out of order packets having sequence numbers in the range + * (X, X+Y]. "Ignoring" here means that the actual sequence number of the corresponding Flow + * (the client one in this case) will be set to the out of order packet's sequece number. + * This means that if an out of order packet is captured having a sequence number X + 5 right + * after enabling recovery mode, then the Flow's sequence number will be set to X + 5. + * + * Note that enabling recovery mode will override the Stream's out of order callbacks, so + * you if you set a callback and then call this method, your callback will be lost. + */ + void enable_recovery_mode(uint32_t recovery_window); + + /** + * \brief Returns true iff recovery mode is enabled + * + * Note that the recovery mode flag will be cleaned only after capturing an out of order + * packet that is outside of the recovery window. + */ + bool is_recovery_mode_enabled() const; private: static Flow extract_client_flow(const PDU& packet); static Flow extract_server_flow(const PDU& packet); @@ -406,6 +438,14 @@ private: void on_server_out_of_order(const Flow& flow, uint32_t seq, const payload_type& payload); + static void client_recovery_mode_handler(Stream& stream, uint32_t sequence_number, + const payload_type& payload, + uint32_t recovery_sequence_number_end); + static void server_recovery_mode_handler(Stream& stream, uint32_t sequence_number, + const payload_type& payload, + uint32_t recovery_sequence_number_end); + static bool recovery_mode_handler(Flow& flow, uint32_t sequence_number, + uint32_t recovery_sequence_number_end); Flow client_flow_; Flow server_flow_; @@ -420,7 +460,8 @@ private: timestamp_type last_seen_; bool auto_cleanup_client_; bool auto_cleanup_server_; - bool is_attached_; + bool is_partial_stream_; + unsigned directions_recovery_mode_enabled_; #ifdef TINS_HAVE_TCP_STREAM_CUSTOM_DATA boost::any user_data_; diff --git a/include/tins/tcp_ip/stream_follower.h b/include/tins/tcp_ip/stream_follower.h index 34cca40..6648329 100644 --- a/include/tins/tcp_ip/stream_follower.h +++ b/include/tins/tcp_ip/stream_follower.h @@ -186,7 +186,19 @@ public: /** * \brief Indicates whether partial streams should be followed. * + * Following partial streams allows capturing packets in the middle of a stream (e.g. + * not capturing the three way handshake) and still reassembling them. * + * This can cause some issues if the first packet captured is out of order, as that would + * create a hole in the sequence number range that might never be filled. In order to + * allow recovering successfully, there's 2 choices: + * + * - Skipping those holes manually by using Flow::advance_sequence. + * - Using Stream::enable_recovery_mode. This is the easiest mechanism and can be used + * on the new stream callback (make sure to only enable it for stream for which + * Stream::is_partial_stream is true). + * + * \param value Whether following partial stream is allowed. */ void follow_partial_streams(bool value); private: diff --git a/src/tcp_ip/stream.cpp b/src/tcp_ip/stream.cpp index a0bd5ff..6851d9c 100644 --- a/src/tcp_ip/stream.cpp +++ b/src/tcp_ip/stream.cpp @@ -62,15 +62,15 @@ Stream::Stream(PDU& packet, const timestamp_type& ts) : client_flow_(extract_client_flow(packet)), server_flow_(extract_server_flow(packet)), create_time_(ts), last_seen_(ts), auto_cleanup_client_(true), auto_cleanup_server_(true), - is_attached_(false) { + is_partial_stream_(false), directions_recovery_mode_enabled_(0) { const EthernetII* eth = packet.find_pdu(); if (eth) { client_hw_addr_ = eth->src_addr(); server_hw_addr_ = eth->dst_addr(); } const TCP& tcp = packet.rfind_pdu(); - // If this is not the first packet of a stream (SYN), then we've attached to it - is_attached_ = tcp.flags() != TCP::SYN; + // If this is not the first packet of a stream (SYN), then it's a partial stream + is_partial_stream_ = tcp.flags() != TCP::SYN; } void Stream::process_packet(PDU& packet, const timestamp_type& ts) { @@ -272,8 +272,21 @@ bool Stream::ack_tracking_enabled() const { return client_flow().ack_tracking_enabled() && server_flow().ack_tracking_enabled(); } -bool Stream::is_attached() const { - return is_attached_; +bool Stream::is_partial_stream() const { + return is_partial_stream_; +} + +void Stream::enable_recovery_mode(uint32_t recovery_window) { + using namespace std::placeholders; + client_out_of_order_callback(bind(&Stream::client_recovery_mode_handler, _1, _2, _3, + client_flow_.sequence_number() + recovery_window)); + server_out_of_order_callback(bind(&Stream::server_recovery_mode_handler, _1, _2, _3, + server_flow_.sequence_number() + recovery_window)); + directions_recovery_mode_enabled_ = 2; +} + +bool Stream::is_recovery_mode_enabled() const { + return directions_recovery_mode_enabled_ > 0; } void Stream::on_client_flow_data(const Flow& /*flow*/) { @@ -294,22 +307,49 @@ void Stream::on_server_flow_data(const Flow& /*flow*/) { } } -void Stream::on_client_out_of_order(const Flow& flow, - uint32_t seq, - const payload_type& payload) { +void Stream::on_client_out_of_order(const Flow& flow, uint32_t seq, const payload_type& payload) { if (on_client_out_of_order_callback_) { on_client_out_of_order_callback_(*this, seq, payload); } } -void Stream::on_server_out_of_order(const Flow& flow, - uint32_t seq, - const payload_type& payload) { +void Stream::on_server_out_of_order(const Flow& flow, uint32_t seq, const payload_type& payload) { if (on_server_out_of_order_callback_) { on_server_out_of_order_callback_(*this, seq, payload); } } +void Stream::client_recovery_mode_handler(Stream& stream, uint32_t sequence_number, + const payload_type& /*payload*/, + uint32_t recovery_sequence_number_end) { + if (!recovery_mode_handler(stream.client_flow(), sequence_number, + recovery_sequence_number_end)) { + stream.client_out_of_order_callback(stream_packet_callback_type()); + stream.directions_recovery_mode_enabled_--; + } +} + +void Stream::server_recovery_mode_handler(Stream& stream, uint32_t sequence_number, + const payload_type& /*payload*/, + uint32_t recovery_sequence_number_end) { + if (!recovery_mode_handler(stream.server_flow(), sequence_number, + recovery_sequence_number_end)) { + stream.server_out_of_order_callback(stream_packet_callback_type()); + stream.directions_recovery_mode_enabled_--; + } +} + +bool Stream::recovery_mode_handler(Flow& flow, uint32_t sequence_number, + uint32_t recovery_sequence_number_end) { + // If this packet comes after our sequence number (would create a hole), skip it + if (sequence_number > flow.sequence_number() && + sequence_number <= recovery_sequence_number_end) { + flow.advance_sequence(sequence_number); + } + // Return true iff we need to keep being in recovery mode + return recovery_sequence_number_end > sequence_number; +} + } // TCPIP } // Tins diff --git a/tests/src/tcp_ip_test.cpp b/tests/src/tcp_ip_test.cpp index 718d57e..84648a7 100644 --- a/tests/src/tcp_ip_test.cpp +++ b/tests/src/tcp_ip_test.cpp @@ -376,7 +376,7 @@ TEST_F(FlowTest, StreamFollower_ThreeWayHandshake) { EXPECT_EQ(Flow::ESTABLISHED, stream.server_flow().state()); EXPECT_EQ(61U, stream.server_flow().sequence_number()); - EXPECT_FALSE(stream.is_attached()); + EXPECT_FALSE(stream.is_partial_stream()); } TEST_F(FlowTest, StreamFollower_TCPOptions) { @@ -531,7 +531,7 @@ TEST_F(FlowTest, StreamFollower_AttachToStreams) { EXPECT_EQ(payload, merge_chunks(stream_client_payload_chunks)); Stream& stream = follower.find_stream(IPv4Address("1.2.3.4"), 22, IPv4Address("4.3.2.1"), 25); - EXPECT_TRUE(stream.is_attached()); + EXPECT_TRUE(stream.is_partial_stream()); } TEST_F(FlowTest, StreamFollower_AttachToStreams_PacketsInBothDirections) { @@ -572,6 +572,7 @@ TEST_F(FlowTest, StreamFollower_AttachToStreams_SecondPacketLost) { packets.erase(packets.begin() + 1); // Erase the 5-10th bytes trimmed_payload.erase(5, 5); + set_endpoints(packets, "1.2.3.4", 22, "4.3.2.1", 25); StreamFollower follower; follower.follow_partial_streams(true); @@ -589,6 +590,34 @@ TEST_F(FlowTest, StreamFollower_AttachToStreams_SecondPacketLost) { EXPECT_EQ(trimmed_payload, merge_chunks(stream_client_payload_chunks)); } +TEST_F(FlowTest, StreamFollower_AttachToStreams_RecoveryMode) { + using std::placeholders::_1; + + ordering_info_type chunks = split_payload(payload, 5); + vector packets = chunks_to_packets(30 /*initial_seq*/, chunks, payload); + string trimmed_payload = payload; + // Erase the 15-20th and 5-10th bytes + trimmed_payload.erase(15, 5); + trimmed_payload.erase(5, 5); + // Erase the second packet + packets.erase(packets.begin() + 3); + packets.erase(packets.begin() + 1); + + set_endpoints(packets, "1.2.3.4", 22, "4.3.2.1", 25); + StreamFollower follower; + follower.follow_partial_streams(true); + follower.new_stream_callback([&](Stream& stream) { + on_new_stream(stream); + stream.enable_recovery_mode(20 /*recovery window size*/); + }); + for (size_t i = 0; i < packets.size(); ++i) { + follower.process_packet(packets[i]); + } + EXPECT_EQ(packets.size(), stream_client_payload_chunks.size()); + EXPECT_EQ(trimmed_payload, merge_chunks(stream_client_payload_chunks)); + EXPECT_EQ(trimmed_payload, merge_chunks(stream_client_payload_chunks)); +} + #ifdef TINS_HAVE_ACK_TRACKER using namespace boost;