1 files changed, 125 insertions, 0 deletions
diff --git a/include/wireframe/net/tcp_reassembly.hpp b/include/wireframe/net/tcp_reassembly.hpp
new file mode 100644
index 0000000..90824a4
--- /dev/null
+++ b/include/wireframe/net/tcp_reassembly.hpp
@@ -0,0 +1,125 @@
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <optional>
+#include <span>
+#include <tuple>
+#include <vector>
+
+#include "wireframe/net/ipv4.hpp"
+
+// Minimal, in-order-only TCP stream reassembly: tracks each flow's two
+// directions separately, accumulating payload bytes as segments arrive
+// exactly in sequence order. Out-of-order segments and retransmissions
+// are dropped rather than buffered for later reordering - a real
+// limitation, but a reasonable one for a learning-focused reassembler
+// capturing directly on an endpoint (this project's demonstrated use
+// all session: lo, wlp1s0, tailscale0), where segments mostly do
+// arrive in order. A capture point far from either endpoint (e.g. a
+// middlebox) would need real out-of-order buffering this doesn't do.
+//
+// The point: HTTP's dissector (wireframe/l7/http.hpp) only ever sees
+// one segment at a time, so a request/response split across TCP
+// segments - a Host: header landing in the second packet of a
+// request, say - is invisible to it. Feeding the *reassembled* stream
+// back through the same parse_http() lets it see what single-segment
+// dissection structurally can't.
+namespace wireframe::net {
+
+struct FlowKey {
+    Ipv4Address ip_a;
+    std::uint16_t port_a;
+    Ipv4Address ip_b;
+    std::uint16_t port_b;
+
+    bool operator<(const FlowKey& other) const {
+        return std::tie(ip_a.bytes, port_a, ip_b.bytes, port_b) <
+               std::tie(other.ip_a.bytes, other.port_a, other.ip_b.bytes, other.port_b);
+    }
+};
+
+// Canonicalizes a (src, dst) pair into a direction-independent
+// FlowKey - both directions of the same connection map to the same
+// key - plus whether this segment's source was the "a" side.
+inline std::pair<FlowKey, bool> canonicalize_flow(const Ipv4Address& src_ip,
+                                                    std::uint16_t src_port,
+                                                    const Ipv4Address& dst_ip,
+                                                    std::uint16_t dst_port) {
+    bool src_is_a = std::tie(src_ip.bytes, src_port) < std::tie(dst_ip.bytes, dst_port);
+    FlowKey key = src_is_a ? FlowKey{src_ip, src_port, dst_ip, dst_port}
+                            : FlowKey{dst_ip, dst_port, src_ip, src_port};
+    return {key, src_is_a};
+}
+
+struct DirectionState {
+    bool syn_seen = false;
+    std::uint32_t next_seq = 0;
+    std::vector<unsigned char> buffer;
+};
+
+struct FlowState {
+    DirectionState a_to_b;
+    DirectionState b_to_a;
+};
+
+class TcpReassembler {
+public:
+    explicit TcpReassembler(std::size_t max_buffer_per_direction = 65536,
+                             std::size_t max_flows = 4096)
+        : max_buffer_(max_buffer_per_direction), max_flows_(max_flows) {}
+
+    // Feeds one TCP segment in. Returns a snapshot of the *sender's*
+    // accumulated stream so far if this segment extended it
+    // contiguously in order; nullopt if the segment was out of order,
+    // a retransmission, a control segment with no payload, or the flow
+    // table was full and this would be a brand new flow. Returned by
+    // value rather than by reference: the buffer this points at can
+    // grow/move on the next call, and bounded copies (max 64 KiB by
+    // default) are cheap enough that this isn't worth the lifetime risk.
+    std::optional<std::vector<unsigned char>> process_segment(
+        const Ipv4Address& src_ip, std::uint16_t src_port, const Ipv4Address& dst_ip,
+        std::uint16_t dst_port, std::uint32_t seq, std::uint8_t flags,
+        std::span<const unsigned char> payload) {
+        auto [key, src_is_a] = canonicalize_flow(src_ip, src_port, dst_ip, dst_port);
+
+        auto it = flows_.find(key);
+        if (it == flows_.end()) {
+            if (flows_.size() >= max_flows_) return std::nullopt;  // table full: drop new flows
+            it = flows_.emplace(key, FlowState{}).first;
+        }
+        DirectionState& dir = src_is_a ? it->second.a_to_b : it->second.b_to_a;
+
+        constexpr std::uint8_t kSyn = 0x02;
+        if (flags & kSyn) {
+            dir.syn_seen = true;
+            dir.next_seq = seq + 1;  // the SYN itself consumes one sequence number
+            return std::nullopt;
+        }
+
+        // seq != dir.next_seq covers both out-of-order segments and
+        // retransmissions (a retransmit repeats a seq already below
+        // next_seq) - unsigned wraparound makes plain equality correct
+        // even across a sequence-number wrap, no need for RFC 1982
+        // serial-number comparison for an exact-match check like this.
+        if (!dir.syn_seen || payload.empty() || seq != dir.next_seq) {
+            return std::nullopt;
+        }
+
+        if (dir.buffer.size() + payload.size() <= max_buffer_) {
+            dir.buffer.insert(dir.buffer.end(), payload.begin(), payload.end());
+        }
+        dir.next_seq = seq + static_cast<std::uint32_t>(payload.size());
+
+        return dir.buffer;
+    }
+
+    std::size_t flow_count() const { return flows_.size(); }
+
+private:
+    std::map<FlowKey, FlowState> flows_;
+    std::size_t max_buffer_;
+    std::size_t max_flows_;
+};
+
+}  // namespace wireframe::net