1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
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
|