Add initial code for new TCP reassembly mechanism

include/tins/tcp_ip.h

@@ -0,0 +1,210 @@
+/*
+ * Copyright (c) 2016, Matias Fontanini
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ * 
+ * * Redistributions of source code must retain the above copyright
+ *   notice, this list of conditions and the following disclaimer.
+ * * Redistributions in binary form must reproduce the above
+ *   copyright notice, this list of conditions and the following disclaimer
+ *   in the documentation and/or other materials provided with the
+ *   distribution.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ */
+
+#ifndef TINS_TCP_IP_H
+#define TINS_TCP_IP_H
+
+#include "cxxstd.h"
+
+// This classes use C++11 features
+#if TINS_IS_CXX11
+
+#include <vector>
+#include <array>
+#include <map>
+#include <functional>
+#include <stdint.h>
+#include "macros.h"
+
+namespace Tins {
+
+class PDU;
+class TCP;
+class IPv4Address;
+class IPv6Address;
+
+namespace TCPIP {
+
+class TINS_API TCPFlow {
+public:
+    enum State {
+        UNKNOWN,
+        SYN_SENT,
+        ESTABLISHED,
+        FIN_SENT,
+        RST_SENT
+    };
+
+    typedef std::vector<uint8_t> payload_type;
+    typedef std::map<uint32_t, payload_type> buffered_payload_type;
+    typedef std::function<void(TCPFlow&)> event_callback;
+
+    TCPFlow(const IPv4Address& dest_address, uint16_t dest_port,
+            uint32_t sequence_number);
+
+    TCPFlow(const IPv6Address& dest_address, uint16_t dest_port,
+            uint32_t sequence_number);
+
+    void data_callback(const event_callback& callback);
+    void buffering_callback(const event_callback& callback);
+
+    void process_packet(PDU& pdu);
+
+    bool is_v6() const;
+    bool is_finished() const;
+    bool packet_belongs(const PDU& packet) const;
+
+    IPv4Address dst_addr_v4() const;
+    IPv6Address dst_addr_v6() const;
+    uint16_t dport() const;
+    const payload_type& payload() const;
+    payload_type& payload();
+
+    void state(State new_state);
+    State state() const;
+    uint32_t sequence_number() const;
+private:
+    void store_payload(uint32_t seq, const payload_type& payload);
+    buffered_payload_type::iterator erase_iterator(buffered_payload_type::iterator iter);
+    void update_state(const TCP& tcp);
+
+    payload_type payload_;
+    buffered_payload_type buffered_payload_;
+    uint32_t seq_number_;
+    std::array<uint8_t, 16> dest_address_;
+    uint16_t dest_port_;
+    event_callback on_data_callback_;
+    event_callback on_buffering_callback_;
+    bool is_v6_;
+    State state_;
+};
+
+class TINS_API TCPStream {
+public:
+    enum State {
+        SYN_SENT,
+        SYN_RCVD,
+        ESTABLISHED,
+        CLOSE_WAIT,
+        FIN_WAIT_1,
+        FIN_WAIT_2,
+        TIME_WAIT,
+        CLOSED
+    };
+
+    typedef std::function<void(TCPStream&)> stream_callback;
+
+    TCPStream(const PDU& initial_packet);
+    TCPStream(const TCPFlow& client_flow, const TCPFlow& server_flow);
+
+    void process_packet(PDU& packet);
+
+    TCPFlow& client_flow();
+    const TCPFlow& client_flow() const;
+    TCPFlow& server_flow();
+    const TCPFlow& server_flow() const;
+
+    void client_data_callback(const stream_callback& callback);
+    void server_data_callback(const stream_callback& callback);
+    void client_buffering_callback(const stream_callback& callback);
+    void server_buffering_callback(const stream_callback& callback);
+
+    void setup_flows_callbacks();
+private:
+    static TCPFlow extract_client_flow(const PDU& packet);
+    static TCPFlow extract_server_flow(const PDU& packet);
+
+
+    void on_client_flow_data(const TCPFlow& flow);
+    void on_server_flow_data(const TCPFlow& flow);
+    void on_client_buffering(const TCPFlow& flow);
+    void on_server_buffering(const TCPFlow& flow);
+
+    TCPFlow client_flow_;
+    TCPFlow server_flow_;
+    stream_callback on_client_data_callback_;
+    stream_callback on_server_data_callback_;
+    stream_callback on_client_buffering_callback_;
+    stream_callback on_server_buffering_callback_;
+    State state_;
+};
+
+class TINS_API TCPStreamFollower {
+public:
+    typedef TCPStream::stream_callback stream_callback;
+
+    TCPStreamFollower();
+
+    void process_packet(PDU& packet);
+
+    void client_data_callback(const stream_callback& callback);
+    void server_data_callback(const stream_callback& callback);
+    void client_buffering_callback(const stream_callback& callback);
+    void server_buffering_callback(const stream_callback& callback);
+
+    TCPStream& find_stream(IPv4Address client_addr, uint16_t client_port,
+                           IPv4Address server_addr, uint16_t server_port);
+private:
+    typedef std::array<uint8_t, 16> address_type;
+
+    struct stream_id {
+        stream_id(const address_type& client_addr, uint16_t client_port,
+                  const address_type& server_addr, uint16_t server_port);
+
+        address_type min_address;
+        address_type max_address;
+        uint16_t min_address_port;
+        uint16_t max_address_port;
+
+        bool operator<(const stream_id& rhs) const;
+
+        static size_t hash(const stream_id& id);
+    };
+
+    typedef std::map<stream_id, TCPStream> streams_type;
+
+    stream_id make_stream_id(const PDU& packet);
+    TCPStream make_stream(const PDU& packet);
+    static address_type serialize(IPv4Address address);
+    static address_type serialize(const IPv6Address& address);
+
+    streams_type streams_;
+    stream_callback on_client_data_callback_;
+    stream_callback on_server_data_callback_;
+    stream_callback on_client_buffering_callback_;
+    stream_callback on_server_buffering_callback_;
+    bool attach_to_flows_;
+};
+
+} // TCPIP
+} // Tins
+
+#endif // TINS_IS_CXX11 
+
+#endif // TINS_TCP_IP_H

src/CMakeLists.txt

@@ -53,6 +53,7 @@ ADD_LIBRARY(
     snap.cpp
     sniffer.cpp
     tcp.cpp
+    tcp_ip.cpp
     tcp_stream.cpp
     udp.cpp
     utils.cpp

src/tcp_ip.cpp

@@ -0,0 +1,529 @@
+/*
+ * Copyright (c) 2016, Matias Fontanini
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ * 
+ * * Redistributions of source code must retain the above copyright
+ *   notice, this list of conditions and the following disclaimer.
+ * * Redistributions in binary form must reproduce the above
+ *   copyright notice, this list of conditions and the following disclaimer
+ *   in the documentation and/or other materials provided with the
+ *   distribution.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ */
+
+#include <limits>
+#include <algorithm>
+#include "tcp_ip.h"
+#include "memory.h"
+#include "ip_address.h"
+#include "ipv6_address.h"
+#include "tcp.h"
+#include "ip.h"
+#include "ipv6.h"
+#include "rawpdu.h"
+#include "exceptions.h"
+#include "memory_helpers.h"
+
+using std::make_pair;
+using std::runtime_error;
+using std::numeric_limits;
+using std::max;
+using std::swap;
+
+using Tins::Memory::OutputMemoryStream;
+using Tins::Memory::InputMemoryStream;
+
+namespace Tins {
+namespace TCPIP {
+
+// As defined by RFC 1982 - 2 ^ (SERIAL_BITS - 1)
+static const uint32_t seq_number_diff = 2147483648U;
+
+// Compares sequence numbers as defined by RFC 1982.
+int seq_compare(uint32_t seq1, uint32_t seq2) {
+    if (seq1 == seq2) {
+        return 0;
+    }
+    if (seq1 < seq2) {
+        return (seq2 - seq1 < seq_number_diff) ? -1 : 1;
+    }
+    else {
+        return (seq1 - seq2 > seq_number_diff) ? -1 : 1;
+    }
+}
+
+// TCPFlow
+
+TCPFlow::TCPFlow(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) {
+    OutputMemoryStream output(dest_address_.data(), dest_address_.size());
+    output.write(dest_address);
+}
+
+TCPFlow::TCPFlow(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) {
+    OutputMemoryStream output(dest_address_.data(), dest_address_.size());
+    output.write(dest_address);
+}
+
+void TCPFlow::data_callback(const event_callback& callback) {
+    on_data_callback_ = callback;
+}
+
+void TCPFlow::buffering_callback(const event_callback& callback) {
+    on_buffering_callback_= callback;
+}
+
+void TCPFlow::process_packet(PDU& pdu) {
+    TCP* tcp = pdu.find_pdu<TCP>();
+    RawPDU* raw = pdu.find_pdu<RawPDU>(); 
+    // If we sent a packet with RST or FIN on, this flow is done
+    if (tcp) {
+        update_state(*tcp);
+    }
+    if (!tcp || !raw) {
+        return;
+    }
+    const uint32_t chunk_end = tcp->seq() + raw->payload_size();
+    // If the end of the chunk ends after our current sequence number, process it.
+    if (seq_compare(chunk_end, seq_number_) >= 0) {
+        bool added_some = false;
+        uint32_t seq = tcp->seq();
+        // If it starts before our sequence number, slice it
+        if (seq_compare(seq, seq_number_) < 0) {
+            const uint32_t diff = seq_number_ - seq;
+            raw->payload().erase(
+                raw->payload().begin(),
+                raw->payload().begin() + diff
+            );
+            seq = seq_number_;
+        }
+        // Store this payload
+        store_payload(seq, 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) {
+            // Does this fragment start before our sequence number?
+            if (seq_compare(iter->first, seq_number_) < 0) {
+                uint32_t fragment_end = iter->first + iter->second.size();
+                int comparison = seq_compare(fragment_end, seq_number_);
+                // Does it end after our sequence number? 
+                if (comparison > 0) {
+                    // Then slice it
+                    payload_type& payload = iter->second;
+                    payload.erase(
+                        payload.begin(),
+                        payload.begin() + (seq_number_ - iter->first)
+                    );
+                    store_payload(seq_number_, iter->second);
+                    iter = erase_iterator(iter);
+                }
+                else {
+                    // Otherwise, we've seen this part of the payload. Erase it.
+                    iter = erase_iterator(iter);
+                }
+            }
+            else {
+                // They're equal. Add this payload.
+                payload_.insert(
+                    payload_.end(),
+                    iter->second.begin(), 
+                    iter->second.end()
+                );
+                seq_number_ += iter->second.size();
+                iter = erase_iterator(iter);
+                added_some = true;
+                // If we don't have any other payload, we're done
+                if (buffered_payload_.empty()) {
+                    break;
+                }
+            }
+        }
+        if (added_some) {
+            if (on_data_callback_) {
+                on_data_callback_(*this);
+            }
+        }
+        else {
+            if (on_buffering_callback_) {
+                on_buffering_callback_(*this);
+            }
+        }
+    }
+}
+
+void TCPFlow::store_payload(uint32_t seq, const 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));
+    }
+    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;
+    }
+}
+
+TCPFlow::buffered_payload_type::iterator TCPFlow::erase_iterator(buffered_payload_type::iterator iter) {
+    buffered_payload_type::iterator output = iter;
+    ++output;
+    buffered_payload_.erase(iter);
+    if (output == buffered_payload_.end()) {
+        output = buffered_payload_.begin();
+    }
+    return output;
+}
+
+void TCPFlow::update_state(const TCP& tcp) {
+    if ((tcp.flags() & TCP::FIN) != 0) {
+        state_ = FIN_SENT;
+    }
+    else if ((tcp.flags() & TCP::RST) != 0) {
+        state_ = RST_SENT;
+    }
+    else if (state_ == SYN_SENT && (tcp.flags() & TCP::ACK) != 0) {
+        state_ = ESTABLISHED;
+        seq_number_++;
+    }
+    else if (state_ == UNKNOWN && (tcp.flags() & TCP::SYN) != 0) {
+        state_ = SYN_SENT;
+        seq_number_ = tcp.seq();
+    }
+}
+
+bool TCPFlow::is_v6() const {
+    return is_v6_;
+}
+
+bool TCPFlow::is_finished() const {
+    return state_ == FIN_SENT || state_ == RST_SENT;
+}
+
+bool TCPFlow::packet_belongs(const PDU& packet) const {
+    if (is_v6()) {
+        const IPv6* ip = packet.find_pdu<IPv6>();
+        if (!ip || ip->dst_addr() != dst_addr_v6()) {
+            return false;
+        }
+    }
+    else {
+        const IP* ip = packet.find_pdu<IP>();
+        if (!ip || ip->dst_addr() != dst_addr_v4()) {
+            return false;
+        }
+    }
+    const TCP* tcp = packet.find_pdu<TCP>();
+    return tcp && tcp->dport() == dport();
+}
+
+IPv4Address TCPFlow::dst_addr_v4() const {
+    InputMemoryStream stream(dest_address_.data(), dest_address_.size());
+    return stream.read<IPv4Address>();
+}
+
+IPv6Address TCPFlow::dst_addr_v6() const {
+    InputMemoryStream stream(dest_address_.data(), dest_address_.size());
+    return stream.read<IPv6Address>();
+}
+
+uint16_t TCPFlow::dport() const {
+    return dest_port_;
+}
+
+const TCPFlow::payload_type& TCPFlow::payload() const {
+    return payload_;
+}
+
+TCPFlow::payload_type& TCPFlow::payload() {
+    return payload_;
+}
+
+void TCPFlow::state(State new_state) {
+    state_ = new_state;
+}
+
+TCPFlow::State TCPFlow::state() const {
+    return state_;
+}
+
+uint32_t TCPFlow::sequence_number() const {
+    return seq_number_;
+}
+
+// TCPStream
+
+TCPStream::TCPStream(const PDU& packet) 
+: client_flow_(extract_client_flow(packet)), 
+server_flow_(extract_server_flow(packet)) {
+
+}
+
+TCPStream::TCPStream(const TCPFlow& client_flow, const TCPFlow& server_flow) 
+: client_flow_(client_flow), server_flow_(server_flow) {
+
+}
+
+void TCPStream::process_packet(PDU& packet) {
+    if (client_flow_.packet_belongs(packet)) {
+        client_flow_.process_packet(packet);
+    }
+    else if (server_flow_.packet_belongs(packet)) {
+        server_flow_.process_packet(packet);
+    }
+}
+
+TCPFlow& TCPStream::client_flow() {
+    return client_flow_;
+}
+
+const TCPFlow& TCPStream::client_flow() const {
+    return client_flow_;
+}
+
+TCPFlow& TCPStream::server_flow() {
+    return server_flow_;
+}
+
+const TCPFlow& TCPStream::server_flow() const {
+    return server_flow_;
+}
+
+void TCPStream::client_data_callback(const stream_callback& callback) {
+    on_client_data_callback_ = callback;
+}
+
+void TCPStream::server_data_callback(const stream_callback& callback) {
+    on_server_data_callback_ = callback;
+}
+
+void TCPStream::client_buffering_callback(const stream_callback& callback) {
+    on_client_buffering_callback_ = callback;
+}
+
+void TCPStream::server_buffering_callback(const stream_callback& callback) {
+    on_server_buffering_callback_ = callback;
+}
+
+TCPFlow TCPStream::extract_client_flow(const PDU& packet) {
+    const TCP* tcp = packet.find_pdu<TCP>();
+    if (!tcp) {
+        // TODO: define proper exception
+        throw runtime_error("No TCP");
+    }
+    if (const IP* ip = packet.find_pdu<IP>()) {
+        return TCPFlow(ip->dst_addr(), tcp->dport(), tcp->seq());
+    }
+    else if (const IPv6* ip = packet.find_pdu<IPv6>()) {
+        return TCPFlow(ip->dst_addr(), tcp->dport(), tcp->seq());
+    }
+    else {
+        // TODO: define proper exception
+        throw runtime_error("No valid layer 3");
+    }
+}
+
+TCPFlow TCPStream::extract_server_flow(const PDU& packet) {
+    const TCP* tcp = packet.find_pdu<TCP>();
+    if (!tcp) {
+        // TODO: define proper exception
+        throw runtime_error("No TCP");
+    }
+    if (const IP* ip = packet.find_pdu<IP>()) {
+        return TCPFlow(ip->src_addr(), tcp->sport(), tcp->ack_seq());
+    }
+    else if (const IPv6* ip = packet.find_pdu<IPv6>()) {
+        return TCPFlow(ip->src_addr(), tcp->sport(), tcp->ack_seq());
+    }
+    else {
+        // TODO: define proper exception
+        throw runtime_error("No valid layer 3");
+    }
+}
+
+void TCPStream::setup_flows_callbacks() {
+    using std::placeholders::_1;
+    client_flow_.data_callback(bind(&TCPStream::on_client_flow_data, this, _1));
+    server_flow_.data_callback(bind(&TCPStream::on_server_flow_data, this, _1));
+    client_flow_.buffering_callback(bind(&TCPStream::on_client_buffering, this, _1));
+    server_flow_.buffering_callback(bind(&TCPStream::on_server_buffering, this, _1));
+}
+
+void TCPStream::on_client_flow_data(const TCPFlow& flow) {
+    if (on_client_data_callback_) {
+        on_client_data_callback_(*this);
+    }
+}
+
+void TCPStream::on_server_flow_data(const TCPFlow& flow) {
+    if (on_server_data_callback_) {
+        on_server_data_callback_(*this);
+    }
+}
+
+void TCPStream::on_client_buffering(const TCPFlow& flow) {
+    if (on_client_buffering_callback_) {
+        on_client_buffering_callback_(*this);
+    }
+}
+
+void TCPStream::on_server_buffering(const TCPFlow& flow) {
+    if (on_server_buffering_callback_) {
+        on_server_buffering_callback_(*this);
+    }
+}
+
+// TCPStreamFollower
+
+TCPStreamFollower::TCPStreamFollower() 
+: attach_to_flows_(false) {
+
+}
+
+void TCPStreamFollower::process_packet(PDU& packet) {
+    stream_id identifier = make_stream_id(packet);
+    streams_type::iterator iter = streams_.find(identifier);
+    bool process = true;
+    if (iter == streams_.end()) {
+        const TCP& tcp = packet.rfind_pdu<TCP>();
+        // Start tracking if they're either SYNs or they contain data (attach
+        // to an already running flow).
+        if (tcp.flags() == TCP::SYN || (attach_to_flows_ && tcp.find_pdu<RawPDU>() != 0)) {
+            iter = streams_.insert(make_pair(identifier, make_stream(packet))).first;
+            iter->second.setup_flows_callbacks();
+            if (tcp.flags() == TCP::SYN) {
+                // If it's a SYN, set the proper state
+                iter->second.client_flow().state(TCPFlow::SYN_SENT);
+                process = false;
+            }
+            else {
+                // Otherwise, assume the connection is established
+                iter->second.client_flow().state(TCPFlow::ESTABLISHED);
+                iter->second.server_flow().state(TCPFlow::ESTABLISHED);
+            }
+        }
+        else {
+            process = false;
+        }
+    }
+    // We'll process it if we had already seen this stream or if we just attached to
+    // it and it contains payload
+    if (process) {
+        iter->second.process_packet(packet);
+    }
+}
+
+void TCPStreamFollower::client_data_callback(const stream_callback& callback) {
+    on_client_data_callback_ = callback;
+}
+
+void TCPStreamFollower::server_data_callback(const stream_callback& callback) {
+    on_server_data_callback_ = callback;
+}
+
+void TCPStreamFollower::client_buffering_callback(const stream_callback& callback) {
+    on_client_buffering_callback_ = callback;
+}
+
+void TCPStreamFollower::server_buffering_callback(const stream_callback& callback) {
+    on_server_buffering_callback_ = callback;
+}
+
+TCPStream& TCPStreamFollower::find_stream(IPv4Address client_addr, uint16_t client_port,
+                                          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;
+    }
+}
+
+TCPStreamFollower::stream_id TCPStreamFollower::make_stream_id(const PDU& packet) {
+    const TCP* tcp = packet.find_pdu<TCP>();
+    if (!tcp) {
+        // TODO: define proper exception
+        throw runtime_error("No TCP");
+    }
+    if (const IP* ip = packet.find_pdu<IP>()) {
+        return stream_id(serialize(ip->src_addr()), tcp->sport(),
+                         serialize(ip->dst_addr()), tcp->dport());
+    }
+    else if (const IPv6* ip = packet.find_pdu<IPv6>()) {
+        return stream_id(serialize(ip->src_addr()), tcp->sport(),
+                         serialize(ip->dst_addr()), tcp->dport());
+    }
+    else {
+        // TODO: define proper exception
+        throw runtime_error("No layer 3");
+    }
+}
+
+TCPStream TCPStreamFollower::make_stream(const PDU& packet) {
+    TCPStream stream(packet);
+    stream.client_data_callback(on_client_data_callback_);
+    stream.server_data_callback(on_server_data_callback_);
+    stream.client_buffering_callback(on_client_buffering_callback_);
+    stream.server_buffering_callback(on_server_buffering_callback_);
+    return stream;
+}
+
+TCPStreamFollower::address_type TCPStreamFollower::serialize(IPv4Address address) {
+    address_type addr;
+    OutputMemoryStream output(addr.data(), addr.size());
+    output.write(address);
+    return addr; 
+}
+
+TCPStreamFollower::address_type TCPStreamFollower::serialize(const IPv6Address& address) {
+    address_type addr;
+    OutputMemoryStream output(addr.data(), addr.size());
+    output.write(address);
+    return addr;
+}
+
+// stream_id
+
+TCPStreamFollower::stream_id::stream_id(const address_type& client_addr,
+                                        uint16_t client_port,
+                                        const address_type& server_addr,
+                                        uint16_t server_port) 
+: min_address(client_addr), max_address(server_addr), min_address_port(client_port),
+max_address_port(server_port) {
+    if (min_address > max_address) {
+        swap(min_address, max_address);
+        swap(min_address_port, max_address_port);
+    }
+}
+
+bool TCPStreamFollower::stream_id::operator<(const stream_id& rhs) const {
+    return tie(min_address, min_address_port, max_address, max_address_port) <
+        tie(rhs.min_address, rhs.min_address_port, rhs.max_address, rhs.max_address_port);
+}
+
+} // TCPIP
+} // Tins

tests/src/CMakeLists.txt

@@ -83,6 +83,7 @@ ADD_CUSTOM_TARGET(
     SLLTest
     SNAPTest
     STPTest
+    TCPIPTest
     TCPTest
     TCPStreamTest
     UDPTest
@@ -128,6 +129,7 @@ ADD_EXECUTABLE(SLLTest EXCLUDE_FROM_ALL sll.cpp)
 ADD_EXECUTABLE(SNAPTest EXCLUDE_FROM_ALL snap.cpp)
 ADD_EXECUTABLE(STPTest EXCLUDE_FROM_ALL stp.cpp)
 ADD_EXECUTABLE(TCPTest EXCLUDE_FROM_ALL tcp.cpp)
+ADD_EXECUTABLE(TCPIPTest EXCLUDE_FROM_ALL tcp_ip.cpp)
 ADD_EXECUTABLE(TCPStreamTest EXCLUDE_FROM_ALL tcp_stream.cpp)
 ADD_EXECUTABLE(UDPTest EXCLUDE_FROM_ALL udp.cpp)
 ADD_EXECUTABLE(UtilsTest EXCLUDE_FROM_ALL utils.cpp)
@@ -208,6 +210,7 @@ ADD_TEST(SLL SLLTest)
 ADD_TEST(SNAP SNAPTest)
 ADD_TEST(STP STPTest)
 ADD_TEST(TCP TCPTest)
+ADD_TEST(TCPIP TCPIPTest)
 ADD_TEST(TCPStream TCPStreamTest)
 ADD_TEST(UDP UDPTest)
 ADD_TEST(Utils UtilsTest)

tests/src/tcp_ip.cpp

@@ -0,0 +1,319 @@
+#include "cxxstd.h"
+
+#if TINS_IS_CXX11
+
+#include <gtest/gtest.h>
+#include <iostream>
+#include <algorithm>
+#include <string>
+#include <limits>
+#include <cassert>
+#include "tcp_ip.h"
+#include "tcp.h"
+#include "ip.h"
+#include "ip_address.h"
+#include "ipv6_address.h"
+#include "ethernetII.h"
+#include "rawpdu.h"
+#include "utils.h"
+
+using namespace std;
+using namespace Tins;
+using namespace Tins::TCPIP;
+
+class TCPFlowTest : public testing::Test {
+public:
+    struct order_element {
+        order_element(size_t payload_index, uint32_t payload_size) 
+        : payload_index(payload_index),payload_size(payload_size) {
+
+        }
+
+        size_t payload_index;
+        uint32_t payload_size;
+    };
+
+    static const size_t num_packets = 20;
+    static EthernetII packets[], overlapped_packets1[], 
+                      overlapped_packets2[], overlapped_packets3[],
+                      overlapped_packets4[], overlapped_packets5[];
+    static const string payload;
+    typedef vector<order_element> ordering_info_type;
+    
+    void cumulative_flow_data_handler(TCPFlow& flow);
+    void cumulative_stream_client_data_handler(TCPStream& stream);
+    void cumulative_stream_server_data_handler(TCPStream& stream);
+    void buffered_payload_handle(TCPFlow& session);
+    void run_test(uint32_t initial_seq, const ordering_info_type& chunks, 
+                  const string& payload);
+    void run_test(uint32_t initial_seq, const ordering_info_type& chunks);
+    void run_tests(const ordering_info_type& chunks, const string& payload);
+    void run_tests(const ordering_info_type& chunks);
+    ordering_info_type split_payload(const string& payload, uint32_t chunk_size);
+    string merge_chunks(const vector<TCPFlow::payload_type>& chunks);
+    vector<EthernetII> chunks_to_packets(uint32_t initial_seq,
+                                         const ordering_info_type& chunks, 
+                                         const string& payload);
+    vector<EthernetII> three_way_handshake(uint32_t client_seq, uint32_t server_seq,
+                                           IPv4Address client_addr, uint16_t client_port,
+                                           IPv4Address server_addr, uint16_t server_port);
+    void set_endpoints(vector<EthernetII>& packets, IPv4Address src_addr,
+                       uint16_t src_port, IPv4Address dst_addr,
+                       uint16_t dst_port);
+    
+    vector<TCPFlow::payload_type> flow_payload_chunks;
+    vector<TCPFlow::payload_type> stream_client_payload_chunks;
+    vector<TCPFlow::payload_type> stream_server_payload_chunks;
+};
+
+const string TCPFlowTest::payload = "Lorem ipsum dolor sit amet, consectetur adipiscing elit. "
+                                    "Sed at aliquam arcu. Sed at iaculis magna. Nam ut dolor "
+                                    "eget velit mattis posuere ut non dui. Aliquam faucibus "
+                                    "erat pretium ligula tincidunt eget tristique justo placerat. "
+                                    "Phasellus turpis tellus, ornare ultricies egestas vitae, "
+                                    "mollis sed neque. Sed et libero in nunc pharetra auctor ut "
+                                    "a eros. Mauris quis faucibus nibh. \nLorem ipsum dolor sit "
+                                    "amet, consectetur adipiscing elit. Sed at aliquam arcu. "
+                                    "Sed at iaculis magna. Nam ut dolor eget velit mattis "
+                                    "posuere ut non dui. Aliquam faucibus erat pretium ligula "
+                                    "tincidunt eget tristique justo placerat. Phasellus turpis "
+                                    "tellus, ornare ultricies egestas vitae, mollis sed neque. "
+                                    "Sed et libero in nunc pharetra auctor ut a eros. Mauris "
+                                    "quis faucibus nibh. \n\n\nCurabitur sem erat, bibendum "
+                                    "quis condimentum ut, imperdiet at est. Duis sagittis rhoncus "
+                                    "felis at ultricies. In libero urna, dignissim eu elementum "
+                                    "quis, consectetur a neque. Praesent leo sem, cursus sed lobortis "
+                                    "sit amet, ornare ac augue. Mauris tristique semper ipsum at "
+                                    "consequat. Sed fringilla dolor ut lacus sagittis quis ultricies "
+                                    "leo vulputate. Maecenas dignissim imperdiet justo. Cras libero "
+                                    "odio, vehicula et adipiscing quis, luctus vel ante. \nAliquam "
+                                    "imperdiet est quis nunc malesuada eget convallis tellus "
+                                    "ullamcorper. Vivamus ullamcorper eros sit amet odio sollicitudin "
+                                    "rutrum. Donec pellentesque faucibus nulla, ut fringilla risus "
+                                    "aliquam eget. Sed et ante mi. Morbi a turpis et tellus dapibus "
+                                    "iaculis. Etiam faucibus tellus sed metus consequat rutrum. "
+                                    "Fusce sit amet nulla massa, tempus vulputate sem. Cras tincidunt "
+                                    "quam in libero rutrum interdum. Aliquam quam sapien, facilisis "
+                                    "at vestibulum et, venenatis id mauris. Morbi rutrum gravida "
+                                    "ultricies. \nAenean et justo ut libero euismod sollicitudin. "
+                                    "Nullam enim dui, iaculis vitae bibendum et, commodo in tellus. "
+                                    "Nullam eget purus mi, a ullamcorper lorem. Suspendisse potenti. "
+                                    "Duis ac justo ut leo euismod gravida sit amet at lectus. Lorem "
+                                    "ipsum dolor sit amet, consectetur adipiscing elit. Maecenas sed "
+                                    "arcu vitae nisi sollicitudin gravida. Nulla facilisis nibh turpis. "
+                                    "Maecenas quis imperdiet arcu. Sed sit amet nulla urna, at "
+                                    "vestibulum mauris. Suspendisse quis elit dui. Class aptent taciti "
+                                    "sociosqu ad litora torquent per conubia nostra, per inceptos "
+                                    "himenaeos. \n";
+
+void TCPFlowTest::cumulative_flow_data_handler(TCPFlow& flow) {
+    flow_payload_chunks.push_back(flow.payload());
+    flow.payload().clear();
+}
+
+void TCPFlowTest::cumulative_stream_client_data_handler(TCPStream& stream) {
+    stream_client_payload_chunks.push_back(stream.client_flow().payload());
+    stream.client_flow().payload().clear();
+}
+
+void TCPFlowTest::cumulative_stream_server_data_handler(TCPStream& stream) {
+    stream_server_payload_chunks.push_back(stream.server_flow().payload());
+    stream.server_flow().payload().clear();
+}
+
+void TCPFlowTest::buffered_payload_handle(TCPFlow& session) {
+
+}
+
+void TCPFlowTest::run_test(uint32_t initial_seq, const ordering_info_type& chunks, 
+                           const string& payload) {
+    using std::placeholders::_1;
+    flow_payload_chunks.clear();
+
+    TCPFlow flow(IPv4Address("1.2.3.4"), 22, initial_seq);
+    flow.data_callback(bind(&TCPFlowTest::cumulative_flow_data_handler, this, _1));
+    vector<EthernetII> packets = chunks_to_packets(initial_seq, chunks, payload);
+    for (size_t i = 0; i < packets.size(); ++i) {
+        flow.process_packet(packets[i]);
+    }
+    string flow_payload = merge_chunks(flow_payload_chunks);
+    EXPECT_EQ(payload, string(flow_payload.begin(), flow_payload.end()));
+}
+
+void TCPFlowTest::run_test(uint32_t initial_seq, const ordering_info_type& chunks) {
+    run_test(initial_seq, chunks, payload);
+}
+
+void TCPFlowTest::run_tests(const ordering_info_type& chunks, const string& payload) {
+    run_test(0, chunks, payload);
+    run_test(20, chunks, payload);
+    run_test(numeric_limits<uint32_t>::max() / 2, chunks, payload);
+    run_test(numeric_limits<uint32_t>::max() - 2, chunks, payload);
+    run_test(numeric_limits<uint32_t>::max() - 5, chunks, payload);
+    run_test(numeric_limits<uint32_t>::max() - 10, chunks, payload);
+    run_test(numeric_limits<uint32_t>::max() - 34, chunks, payload);
+    run_test(numeric_limits<uint32_t>::max() - 31, chunks, payload);
+}
+
+void TCPFlowTest::run_tests(const ordering_info_type& chunks) {
+    run_tests(chunks, payload);
+}
+
+TCPFlowTest::ordering_info_type TCPFlowTest::split_payload(const string& payload,
+                                                           uint32_t chunk_size) {
+    ordering_info_type output;
+    uint32_t chunk_count = payload.size() / chunk_size;
+    for (uint32_t i = 0; i < chunk_count; ++i) {
+        output.push_back(order_element(i * chunk_size, chunk_size));
+    }
+    if (chunk_count * chunk_size < payload.size()) {
+        uint32_t index = chunk_count * chunk_size;
+        output.push_back(order_element(index, payload.size() - index));
+    }
+    return output;
+}
+
+string TCPFlowTest::merge_chunks(const vector<TCPFlow::payload_type>& chunks) {
+    string output;
+    for (size_t i = 0; i < chunks.size(); ++i) {
+        TCPFlow::payload_type this_chunk = chunks[i];
+        output += string(this_chunk.begin(), this_chunk.end());
+    }
+    return output;
+}
+
+vector<EthernetII> TCPFlowTest::chunks_to_packets(uint32_t initial_seq,
+                                                  const ordering_info_type& chunks, 
+                                                  const string& payload) {
+    vector<EthernetII> output;
+    for (size_t i = 0; i < chunks.size(); ++i) {
+        const order_element& element = chunks[i];
+        assert(element.payload_index + element.payload_size <= payload.size());
+        TCP tcp;
+        RawPDU raw(payload.begin() + element.payload_index, 
+                   payload.begin() + element.payload_index + element.payload_size);
+        tcp.seq(initial_seq + element.payload_index);
+        output.push_back(EthernetII() / IP() / tcp / raw);
+    }
+    return output;
+}
+
+vector<EthernetII> TCPFlowTest::three_way_handshake(uint32_t client_seq, uint32_t server_seq,
+                                                    IPv4Address client_addr, uint16_t client_port,
+                                                    IPv4Address server_addr, uint16_t server_port) {
+    vector<EthernetII> output;
+    output.push_back(EthernetII() / IP(server_addr, client_addr) / TCP(server_port, client_port));
+    output.push_back(EthernetII() / IP(client_addr, server_addr) / TCP(client_port, server_port));
+    output.push_back(EthernetII() / IP(server_addr, client_addr) / TCP(server_port, client_port));
+    output[0].rfind_pdu<TCP>().flags(TCP::SYN);
+    output[0].rfind_pdu<TCP>().seq(client_seq);
+    output[1].rfind_pdu<TCP>().flags(TCP::SYN | TCP::ACK);
+    output[1].rfind_pdu<TCP>().seq(server_seq);
+    output[1].rfind_pdu<TCP>().ack_seq(client_seq + 1);
+    output[2].rfind_pdu<TCP>().flags(TCP::ACK);
+    output[2].rfind_pdu<TCP>().seq(client_seq + 1);
+    output[2].rfind_pdu<TCP>().ack_seq(server_seq + 1);
+    return output;
+}
+
+void TCPFlowTest::set_endpoints(vector<EthernetII>& packets, IPv4Address src_addr,
+                                uint16_t src_port, IPv4Address dst_addr,
+                                uint16_t dst_port) {
+    for (size_t i = 0; i < packets.size(); ++i) {
+        packets[i].rfind_pdu<IP>().src_addr(src_addr);
+        packets[i].rfind_pdu<IP>().dst_addr(dst_addr);
+        packets[i].rfind_pdu<TCP>().sport(src_port);
+        packets[i].rfind_pdu<TCP>().dport(dst_port);
+    }
+}
+
+TEST_F(TCPFlowTest, ReassembleStreamPlain) {
+    ordering_info_type chunks = split_payload(payload, 5);
+    run_tests(chunks);
+}
+
+TEST_F(TCPFlowTest, ReassembleStreamReordering) {
+    ordering_info_type chunks = split_payload(payload, 5);
+    // e.g. input [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+    // after this it's [2, 1, 0, 3, 6, 5, 4, 7, 10, 9, 8]
+    for (size_t i = 0; i < chunks.size(); i += 4) {
+        if (i + 2 < chunks.size()) {
+            swap(chunks[i], chunks[i + 2]);
+        }
+    }
+    run_tests(chunks);
+}
+
+TEST_F(TCPFlowTest, ReassembleStreamReversed) {
+    ordering_info_type chunks = split_payload(payload, 5);
+    reverse(chunks.begin(), chunks.end());
+    run_tests(chunks);
+}
+
+TEST_F(TCPFlowTest, Overlapping) {
+    string payload = "Hello world. This is a payload";
+    ordering_info_type chunks;
+    // "Hello "
+    chunks.push_back(order_element(0, 6));
+    // "ello Wo"
+    chunks.push_back(order_element(1, 7));
+    // "lo World"
+    chunks.push_back(order_element(3, 8));
+    chunks.push_back(order_element(10, payload.size() - 10));
+    chunks.push_back(order_element(9, 1));
+    run_tests(chunks, payload);
+
+    reverse(chunks.begin(), chunks.end());
+    run_tests(chunks, payload);
+
+    swap(chunks[2], chunks[4]);
+    run_tests(chunks, payload);
+}
+
+TEST_F(TCPFlowTest, TCPStreamFollower_ThreeWayHandshake) {
+    using std::placeholders::_1;
+
+    vector<EthernetII> packets = three_way_handshake(29, 60, "1.2.3.4", 22, "4.3.2.1", 25);
+    TCPStreamFollower follower;
+    follower.client_data_callback(bind(&TCPFlowTest::cumulative_stream_client_data_handler, 
+                                       this, _1));
+    for (size_t i = 0; i < packets.size(); ++i) {
+        follower.process_packet(packets[i]);
+    }
+    TCPStream& stream = follower.find_stream("1.2.3.4", 22, "4.3.2.1", 25);
+    EXPECT_EQ(TCPFlow::ESTABLISHED, stream.client_flow().state());
+    EXPECT_EQ(TCPFlow::SYN_SENT, stream.server_flow().state());
+    EXPECT_EQ(30, stream.client_flow().sequence_number());
+    EXPECT_EQ(60, stream.server_flow().sequence_number());
+    EXPECT_EQ(IPv4Address("4.3.2.1"), stream.client_flow().dst_addr_v4());
+    EXPECT_EQ(25, stream.client_flow().dport());
+    EXPECT_EQ(IPv4Address("1.2.3.4"), stream.server_flow().dst_addr_v4());
+    EXPECT_EQ(22, stream.server_flow().dport());
+
+    IP server_packet = IP("1.2.3.4", "4.3.2.1") / TCP(22, 25);
+    server_packet.rfind_pdu<TCP>().flags(TCP::ACK);
+    follower.process_packet(server_packet);
+
+    EXPECT_EQ(TCPFlow::ESTABLISHED, stream.server_flow().state());
+    EXPECT_EQ(61, stream.server_flow().sequence_number());
+}
+
+TEST_F(TCPFlowTest, TCPStreamFollower_FollowStream) {
+    using std::placeholders::_1;
+
+    vector<EthernetII> packets = three_way_handshake(29, 60, "1.2.3.4", 22, "4.3.2.1", 25);
+    ordering_info_type chunks = split_payload(payload, 5);
+    vector<EthernetII> chunk_packets = chunks_to_packets(30 /*initial_seq*/, chunks, payload);
+    set_endpoints(chunk_packets, "1.2.3.4", 22, "4.3.2.1", 25);
+    packets.insert(packets.end(), chunk_packets.begin(), chunk_packets.end());
+    TCPStreamFollower follower;
+    follower.client_data_callback(bind(&TCPFlowTest::cumulative_stream_client_data_handler, 
+                                       this, _1));
+    for (size_t i = 0; i < packets.size(); ++i) {
+        follower.process_packet(packets[i]);
+    }
+    EXPECT_EQ(chunk_packets.size(), stream_client_payload_chunks.size());
+    EXPECT_EQ(payload, merge_chunks(stream_client_payload_chunks));
+}
+
+#endif // TINS_IS_CXX11