diff --git a/include/tins/tcp_ip.h b/include/tins/tcp_ip.h index bf69f18..0a4e26d 100644 --- a/include/tins/tcp_ip.h +++ b/include/tins/tcp_ip.h @@ -231,8 +231,16 @@ public: * \param new_state The new state of this flow */ void state(State new_state); + + /** + * \brief Sets whether this flow should ignore data packets + * + * If the data packets are ignored then the flow will just be + * followed to keep track of its state. + */ + void ignore_data_packets(); private: - void store_payload(uint32_t seq, const payload_type& payload); + void store_payload(uint32_t seq, payload_type payload); buffered_payload_type::iterator erase_iterator(buffered_payload_type::iterator iter); void update_state(const TCP& tcp); @@ -245,6 +253,7 @@ private: event_callback on_buffering_callback_; State state_; bool is_v6_; + bool ignore_data_packets_; }; /** @@ -417,6 +426,22 @@ public: */ void server_buffering_callback(const stream_callback& callback); + /** + * \brief Indicates that the data packets sent by the client should be + * ignored + * + * \sa Flow::ignore_data_packets + */ + void ignore_client_data(); + + /** + * \brief Indicates that the data packets sent by the server should be + * ignored + * + * \sa Flow::ignore_data_packets + */ + void ignore_server_data(); + /** * \brief Sets the internal callbacks. * diff --git a/src/tcp_ip.cpp b/src/tcp_ip.cpp index 6df9149..9a6e490 100644 --- a/src/tcp_ip.cpp +++ b/src/tcp_ip.cpp @@ -45,6 +45,7 @@ using std::make_pair; using std::bind; +using std::pair; using std::runtime_error; using std::numeric_limits; using std::max; @@ -77,7 +78,7 @@ 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), state_(UNKNOWN), - is_v6_(false) { + is_v6_(false), ignore_data_packets_(false) { OutputMemoryStream output(dest_address_.data(), dest_address_.size()); output.write(dest_address); } @@ -85,7 +86,7 @@ Flow::Flow(const IPv4Address& dest_address, uint16_t dest_port, Flow::Flow(const IPv6Address& dest_address, uint16_t dest_port, uint32_t sequence_number) : seq_number_(sequence_number), dest_port_(dest_port), state_(UNKNOWN), - is_v6_(true) { + is_v6_(true), ignore_data_packets_(false) { OutputMemoryStream output(dest_address_.data(), dest_address_.size()); output.write(dest_address); } @@ -105,6 +106,9 @@ void Flow::process_packet(PDU& pdu) { if (tcp) { update_state(*tcp); } + if (ignore_data_packets_) { + return; + } if (!tcp || !raw) { return; } @@ -123,7 +127,7 @@ void Flow::process_packet(PDU& pdu) { seq = seq_number_; } // Store this payload - store_payload(seq, raw->payload()); + store_payload(seq, move(raw->payload())); // Keep looping while the fragments seq is lower or equal to our seq buffered_payload_type::iterator iter = buffered_payload_.find(seq_number_); while (iter != buffered_payload_.end() && seq_compare(iter->first, seq_number_) <= 0) { @@ -139,7 +143,7 @@ void Flow::process_packet(PDU& pdu) { payload.begin(), payload.begin() + (seq_number_ - iter->first) ); - store_payload(seq_number_, iter->second); + store_payload(seq_number_, move(iter->second)); iter = erase_iterator(iter); } else { @@ -176,16 +180,16 @@ void Flow::process_packet(PDU& pdu) { } } -void Flow::store_payload(uint32_t seq, const payload_type& payload) { +void Flow::store_payload(uint32_t seq, payload_type payload) { buffered_payload_type::iterator iter = buffered_payload_.find(seq); // New segment, store it if (iter == buffered_payload_.end()) { - buffered_payload_.insert(make_pair(seq, payload)); + buffered_payload_.insert(make_pair(seq, move(payload))); } else if (iter->second.size() < payload.size()) { // If we already have payload on this position but it's a shorter // chunk than the new one, replace it - iter->second = payload; + iter->second = move(payload); } } @@ -283,6 +287,10 @@ void Flow::state(State new_state) { state_ = new_state; } +void Flow::ignore_data_packets() { + ignore_data_packets_ = true; +} + // Stream Stream::Stream(const PDU& packet) @@ -343,6 +351,14 @@ void Stream::server_buffering_callback(const stream_callback& callback) { on_server_buffering_callback_ = callback; } +void Stream::ignore_client_data() { + client_flow().ignore_data_packets(); +} + +void Stream::ignore_server_data() { + server_flow().ignore_data_packets(); +} + bool Stream::is_finished() const { const Flow::State client_state = client_flow_.state(); const Flow::State server_state = server_flow_.state(); diff --git a/tests/src/tcp_ip.cpp b/tests/src/tcp_ip.cpp index d4f2bd7..52aedf9 100644 --- a/tests/src/tcp_ip.cpp +++ b/tests/src/tcp_ip.cpp @@ -276,6 +276,22 @@ TEST_F(FlowTest, Overlapping) { run_tests(chunks, payload); } +TEST_F(FlowTest, IgnoreDataPackets) { + using std::placeholders::_1; + + ordering_info_type chunks = split_payload(payload, 5); + Flow flow(IPv4Address("1.2.3.4"), 22, 0); + flow.data_callback(bind(&FlowTest::cumulative_flow_data_handler, this, _1)); + flow.ignore_data_packets(); + vector packets = chunks_to_packets(0, chunks, payload); + for (size_t i = 0; i < packets.size(); ++i) { + flow.process_packet(packets[i]); + } + EXPECT_TRUE(flow_payload_chunks.empty()); +} + +// Stream follower tests + TEST_F(FlowTest, StreamFollower_ThreeWayHandshake) { using std::placeholders::_1;