wireforge_core/
tcp.rs

1//! TCP packet parser and builder with option parsing.
2
3use alloc::vec;
4use alloc::vec::Vec;
5use core::net::Ipv4Addr;
6
7use crate::types::IpProtocol;
8use crate::util::{pseudo_header_checksum, seal_transport_checksum, read_u16be, read_u32be, write_u16be, write_u32be};
9
10pub const TCP_MIN_HEADER_LEN: usize = 20;
11
12// TCP flag bit positions
13const FLAG_FIN: u8 = 0x01;
14const FLAG_SYN: u8 = 0x02;
15const FLAG_RST: u8 = 0x04;
16const FLAG_PSH: u8 = 0x08;
17const FLAG_ACK: u8 = 0x10;
18const FLAG_URG: u8 = 0x20;
19
20/// Default TCP data offset (5 × 4 = 20-byte header, no options).
21const DEFAULT_DATA_OFFSET: u8 = 5;
22
23// TCP option kinds
24const OPT_END: u8 = 0;
25const OPT_NOOP: u8 = 1;
26const OPT_MSS: u8 = 2;
27const OPT_WINDOW_SCALE: u8 = 3;
28const OPT_SACK_PERMITTED: u8 = 4;
29const OPT_SACK: u8 = 5;
30const OPT_TIMESTAMP: u8 = 8;
31
32/// Parsed TCP option.
33#[derive(Debug, Clone)]
34pub enum TcpOption {
35    EndOfList,
36    NoOp,
37    MaximumSegmentSize(u16),
38    WindowScale(u8),
39    SackPermitted,
40    Sack(Vec<(u32, u32)>),
41    Timestamp { ts_val: u32, ts_ecr: u32 },
42    Unknown { kind: u8, data: Vec<u8> },
43}
44
45/// Iterator over TCP options in the header.
46pub struct TcpOptionsIter<'a> {
47    data: &'a [u8],
48}
49
50impl<'a> TcpOptionsIter<'a> {
51    fn new(data: &'a [u8]) -> Self {
52        Self { data }
53    }
54}
55
56impl<'a> Iterator for TcpOptionsIter<'a> {
57    type Item = TcpOption;
58
59    fn next(&mut self) -> Option<Self::Item> {
60        if self.data.is_empty() {
61            return None;
62        }
63        let kind = self.data[0];
64        match kind {
65            OPT_END => {
66                self.data = &[];
67                Some(TcpOption::EndOfList)
68            }
69            OPT_NOOP => {
70                self.data = &self.data[1..];
71                Some(TcpOption::NoOp)
72            }
73            OPT_MSS => {
74                if self.data.len() < 4 {
75                    self.data = &[];
76                    return Some(TcpOption::Unknown { kind, data: self.data.to_vec() });
77                }
78                let mss = read_u16be(&self.data[2..4]);
79                self.data = &self.data[4..];
80                Some(TcpOption::MaximumSegmentSize(mss))
81            }
82            OPT_WINDOW_SCALE => {
83                if self.data.len() < 3 {
84                    self.data = &[];
85                    return Some(TcpOption::Unknown { kind, data: self.data.to_vec() });
86                }
87                let shift = self.data[2];
88                self.data = &self.data[3..];
89                Some(TcpOption::WindowScale(shift))
90            }
91            OPT_SACK_PERMITTED => {
92                if self.data.len() < 2 {
93                    self.data = &[];
94                    return Some(TcpOption::Unknown { kind, data: self.data.to_vec() });
95                }
96                self.data = &self.data[2..];
97                Some(TcpOption::SackPermitted)
98            }
99            OPT_SACK => {
100                let len = self.data.get(1).copied().unwrap_or(0) as usize;
101                if len < 2 || self.data.len() < len {
102                    self.data = &[];
103                    return Some(TcpOption::Unknown { kind, data: self.data.to_vec() });
104                }
105                let blocks_data = &self.data[2..len];
106                let mut blocks = Vec::new();
107                for chunk in blocks_data.chunks(8) {
108                    if chunk.len() == 8 {
109                        blocks.push((read_u32be(&chunk[..4]), read_u32be(&chunk[4..8])));
110                    }
111                }
112                self.data = &self.data[len..];
113                Some(TcpOption::Sack(blocks))
114            }
115            OPT_TIMESTAMP => {
116                if self.data.len() < 10 {
117                    self.data = &[];
118                    return Some(TcpOption::Unknown { kind, data: self.data.to_vec() });
119                }
120                let ts_val = read_u32be(&self.data[2..6]);
121                let ts_ecr = read_u32be(&self.data[6..10]);
122                self.data = &self.data[10..];
123                Some(TcpOption::Timestamp { ts_val, ts_ecr })
124            }
125            _ => {
126                let len = self.data.get(1).copied().unwrap_or(0) as usize;
127                if kind > 1 && len >= 2 && self.data.len() >= len {
128                    let data = self.data[..len].to_vec();
129                    self.data = &self.data[len..];
130                    Some(TcpOption::Unknown { kind, data })
131                } else {
132                    let remaining = self.data.to_vec();
133                    self.data = &[];
134                    Some(TcpOption::Unknown { kind, data: remaining })
135                }
136            }
137        }
138    }
139}
140
141/// Zero-copy TCP packet parser.
142#[derive(Debug, Clone)]
143pub struct TcpPacket<'a> {
144    buf: &'a [u8],
145}
146
147impl<'a> TcpPacket<'a> {
148    pub fn new(buf: &'a [u8]) -> Option<Self> {
149        if buf.len() < TCP_MIN_HEADER_LEN {
150            return None;
151        }
152        let data_offset = (buf[12] >> 4) as usize * 4;
153        if data_offset < TCP_MIN_HEADER_LEN || buf.len() < data_offset {
154            return None;
155        }
156        Some(Self { buf })
157    }
158
159    /// Raw bytes backing this packet.
160    #[inline]
161    pub fn as_bytes(&self) -> &'a [u8] { self.buf }
162
163    #[inline]
164    pub fn source_port(&self) -> u16 {
165        read_u16be(&self.buf[..2])
166    }
167
168    #[inline]
169    pub fn destination_port(&self) -> u16 {
170        read_u16be(&self.buf[2..4])
171    }
172
173    #[inline]
174    pub fn sequence(&self) -> u32 {
175        read_u32be(&self.buf[4..8])
176    }
177
178    #[inline]
179    pub fn ack_number(&self) -> u32 {
180        read_u32be(&self.buf[8..12])
181    }
182
183    #[inline]
184    pub fn data_offset(&self) -> u8 {
185        self.buf[12] >> 4
186    }
187
188    #[inline]
189    pub fn flags(&self) -> u8 {
190        self.buf[13] & 0x3F
191    }
192
193    #[inline]
194    pub fn fin(&self) -> bool {
195        self.flags() & FLAG_FIN != 0
196    }
197
198    #[inline]
199    pub fn syn(&self) -> bool {
200        self.flags() & FLAG_SYN != 0
201    }
202
203    #[inline]
204    pub fn rst(&self) -> bool {
205        self.flags() & FLAG_RST != 0
206    }
207
208    #[inline]
209    pub fn psh(&self) -> bool {
210        self.flags() & FLAG_PSH != 0
211    }
212
213    #[inline]
214    pub fn ack(&self) -> bool {
215        self.flags() & FLAG_ACK != 0
216    }
217
218    #[inline]
219    pub fn urg(&self) -> bool {
220        self.flags() & FLAG_URG != 0
221    }
222
223    #[inline]
224    pub fn window(&self) -> u16 {
225        read_u16be(&self.buf[14..16])
226    }
227
228    #[inline]
229    pub fn checksum(&self) -> u16 {
230        read_u16be(&self.buf[16..18])
231    }
232
233    #[inline]
234    pub fn urgent_pointer(&self) -> u16 {
235        read_u16be(&self.buf[18..20])
236    }
237
238    /// TCP options (if data_offset > 5).
239    pub fn options(&self) -> TcpOptionsIter<'a> {
240        let hdr_len = self.header_length();
241        if hdr_len > TCP_MIN_HEADER_LEN {
242            TcpOptionsIter::new(&self.buf[TCP_MIN_HEADER_LEN..hdr_len])
243        } else {
244            TcpOptionsIter::new(&[])
245        }
246    }
247
248    /// Payload after the TCP header.
249    #[inline]
250    pub fn payload(&self) -> &'a [u8] {
251        &self.buf[self.header_length()..]
252    }
253
254    #[inline]
255    pub fn header_length(&self) -> usize {
256        self.data_offset() as usize * 4
257    }
258
259    /// Verify the TCP checksum (includes IPv4 pseudo-header).
260    pub fn verify_checksum(&self, src: Ipv4Addr, dst: Ipv4Addr) -> bool {
261        pseudo_header_checksum(
262            &src.octets(),
263            &dst.octets(),
264            IpProtocol::Tcp.into(),
265            self.buf,
266        ) == 0
267    }
268}
269
270// ---------------------------------------------------------------------------
271// Builder
272// ---------------------------------------------------------------------------
273
274pub struct TcpPacketBuilder {
275    buf: Vec<u8>,
276    options: Vec<u8>,
277    payload: Option<Vec<u8>>,
278}
279
280impl Default for TcpPacketBuilder {
281    fn default() -> Self {
282        Self::new()
283    }
284}
285
286impl TcpPacketBuilder {
287    pub fn new() -> Self {
288        let mut buf = vec![0u8; TCP_MIN_HEADER_LEN];
289        buf[12] = DEFAULT_DATA_OFFSET << 4;
290        Self { buf, options: Vec::new(), payload: None }
291    }
292
293    pub fn source_port(mut self, port: u16) -> Self {
294        write_u16be(&mut self.buf[..2], port);
295        self
296    }
297
298    pub fn destination_port(mut self, port: u16) -> Self {
299        write_u16be(&mut self.buf[2..4], port);
300        self
301    }
302
303    pub fn sequence(mut self, seq: u32) -> Self {
304        write_u32be(&mut self.buf[4..8], seq);
305        self
306    }
307
308    pub fn ack_number(mut self, ack: u32) -> Self {
309        write_u32be(&mut self.buf[8..12], ack);
310        self
311    }
312
313    pub fn window(mut self, w: u16) -> Self {
314        write_u16be(&mut self.buf[14..16], w);
315        self
316    }
317
318    pub fn urgent_pointer(mut self, up: u16) -> Self {
319        write_u16be(&mut self.buf[18..20], up);
320        self
321    }
322
323    pub fn syn(mut self, on: bool) -> Self {
324        if on { self.buf[13] |= FLAG_SYN; } else { self.buf[13] &= !FLAG_SYN; }
325        self
326    }
327
328    pub fn ack(mut self, on: bool) -> Self {
329        if on { self.buf[13] |= FLAG_ACK; } else { self.buf[13] &= !FLAG_ACK; }
330        self
331    }
332
333    pub fn fin(mut self, on: bool) -> Self {
334        if on { self.buf[13] |= FLAG_FIN; } else { self.buf[13] &= !FLAG_FIN; }
335        self
336    }
337
338    pub fn rst(mut self, on: bool) -> Self {
339        if on { self.buf[13] |= FLAG_RST; } else { self.buf[13] &= !FLAG_RST; }
340        self
341    }
342
343    pub fn psh(mut self, on: bool) -> Self {
344        if on { self.buf[13] |= FLAG_PSH; } else { self.buf[13] &= !FLAG_PSH; }
345        self
346    }
347
348    pub fn urg(mut self, on: bool) -> Self {
349        if on { self.buf[13] |= FLAG_URG; } else { self.buf[13] &= !FLAG_URG; }
350        self
351    }
352
353    /// Add a TCP option as raw bytes (kind + length + data).
354    pub fn add_option(mut self, data: &[u8]) -> Self {
355        self.options.extend_from_slice(data);
356        self
357    }
358
359    /// Convenience: add MSS option.
360    pub fn mss(mut self, mss: u16) -> Self {
361        self.options.push(OPT_MSS);
362        self.options.push(4);
363        self.options.extend_from_slice(&mss.to_be_bytes());
364        self
365    }
366
367    /// Convenience: add Window Scale option.
368    pub fn window_scale(mut self, shift: u8) -> Self {
369        self.options.push(OPT_WINDOW_SCALE);
370        self.options.push(3);
371        self.options.push(shift);
372        self
373    }
374
375    /// Convenience: add SACK Permitted option.
376    pub fn sack_permitted(mut self) -> Self {
377        self.options.push(OPT_SACK_PERMITTED);
378        self.options.push(2);
379        self
380    }
381
382    /// Convenience: add Timestamp option.
383    pub fn timestamp(mut self, ts_val: u32, ts_ecr: u32) -> Self {
384        self.options.push(OPT_TIMESTAMP);
385        self.options.push(10);
386        self.options.extend_from_slice(&ts_val.to_be_bytes());
387        self.options.extend_from_slice(&ts_ecr.to_be_bytes());
388        self
389    }
390
391    pub fn payload(mut self, data: &[u8]) -> Self {
392        self.payload = Some(data.to_vec());
393        self
394    }
395
396    /// Build the TCP packet, computing the checksum (IPv4 pseudo-header).
397    pub fn build(mut self, src: Ipv4Addr, dst: Ipv4Addr) -> Vec<u8> {
398        // Pad options to 4-byte boundary with NoOps
399        while !self.options.len().is_multiple_of(4) {
400            self.options.push(OPT_NOOP);
401        }
402        // Update data_offset
403        let total_hdr_words = (TCP_MIN_HEADER_LEN + self.options.len()) / 4;
404        self.buf[12] = (total_hdr_words as u8) << 4;
405
406        let mut packet = self.buf;
407        if !self.options.is_empty() {
408            packet.extend_from_slice(&self.options);
409        }
410        if let Some(ref p) = self.payload {
411            packet.extend_from_slice(p);
412        }
413        seal_transport_checksum(&mut packet, 16, &src.octets(), &dst.octets(), IpProtocol::Tcp.into());
414        packet
415    }
416}
417
418// ---------------------------------------------------------------------------
419// Tests
420// ---------------------------------------------------------------------------
421
422#[cfg(test)]
423mod tests {
424    use super::*;
425
426    // Minimal TCP header: src=1234, dst=80, seq=1, ack=0, data_offset=5, flags=SYN
427    fn sample_tcp_syn() -> Vec<u8> {
428        let mut data = vec![0u8; TCP_MIN_HEADER_LEN];
429        write_u16be(&mut data[..2], 1234);
430        write_u16be(&mut data[2..4], 80);
431        write_u32be(&mut data[4..8], 1);      // seq
432        write_u32be(&mut data[8..12], 0);      // ack
433        data[12] = DEFAULT_DATA_OFFSET << 4;
434        data[13] = FLAG_SYN;
435        write_u16be(&mut data[14..16], 65535); // window
436        data
437    }
438
439    #[test]
440    fn parse_tcp_no_options() {
441        let data = sample_tcp_syn();
442        let pkt = TcpPacket::new(&data).unwrap();
443        assert_eq!(pkt.source_port(), 1234);
444        assert_eq!(pkt.destination_port(), 80);
445        assert_eq!(pkt.sequence(), 1);
446        assert_eq!(pkt.ack_number(), 0);
447        assert_eq!(pkt.data_offset(), 5);
448        assert_eq!(pkt.header_length(), 20);
449        assert!(pkt.syn());
450        assert!(!pkt.ack());
451        assert!(!pkt.fin());
452        assert!(!pkt.rst());
453        assert_eq!(pkt.window(), 65535);
454    }
455
456    #[test]
457    fn parse_tcp_too_short() {
458        assert!(TcpPacket::new(&[]).is_none());
459        assert!(TcpPacket::new(&[0u8; 19]).is_none());
460        // 20 bytes with data_offset=5 is valid
461        let mut d = [0u8; 20];
462        d[12] = DEFAULT_DATA_OFFSET << 4;
463        assert!(TcpPacket::new(&d).is_some());
464    }
465
466    #[test]
467    fn parse_tcp_data_offset_exceeds_buffer() {
468        let mut data = vec![0u8; 22];
469        data[12] = 0x60; // data_offset=6 (24 bytes)
470        assert!(TcpPacket::new(&data).is_none());
471    }
472
473    #[test]
474    fn tcp_options_mss_and_window_scale() {
475        let mut data = sample_tcp_syn();
476        // MSS: kind=2, len=4, mss=1460
477        data.extend_from_slice(&[OPT_MSS, 4, 0x05, 0xB4]);
478        // Window Scale: kind=3, len=3, shift=7
479        data.extend_from_slice(&[OPT_WINDOW_SCALE, 3, 7]);
480        // Pad to 4-byte boundary with NoOps (1 byte padding)
481        data.push(OPT_NOOP);
482        // data_offset: (20 + 8) / 4 = 7
483        data[12] = 0x70;
484
485        let pkt = TcpPacket::new(&data).unwrap();
486        assert_eq!(pkt.data_offset(), 7);
487        assert_eq!(pkt.header_length(), 28);
488
489        let opts: Vec<_> = pkt.options().collect();
490        assert_eq!(opts.len(), 3); // MSS, WindowScale, NoOp
491        match &opts[0] {
492            TcpOption::MaximumSegmentSize(mss) => assert_eq!(*mss, 1460),
493            _ => panic!("expected MSS"),
494        }
495        match &opts[1] {
496            TcpOption::WindowScale(shift) => assert_eq!(*shift, 7),
497            _ => panic!("expected WindowScale"),
498        }
499    }
500
501    #[test]
502    fn tcp_options_timestamp() {
503        let mut data = sample_tcp_syn();
504        data.extend_from_slice(&[OPT_TIMESTAMP, 10,
505            0x00, 0x00, 0x00, 0x01,  // ts_val
506            0x00, 0x00, 0x00, 0x02,  // ts_ecr
507        ]);
508        // Pad to 4-byte boundary (10 + 2 = 12)
509        data.push(OPT_NOOP);
510        data.push(OPT_NOOP);
511        data[12] = 0x80; // data_offset=8 (32 bytes)
512
513        let pkt = TcpPacket::new(&data).unwrap();
514        let opts: Vec<_> = pkt.options().collect();
515        match &opts[0] {
516            TcpOption::Timestamp { ts_val, ts_ecr } => {
517                assert_eq!(*ts_val, 1);
518                assert_eq!(*ts_ecr, 2);
519            }
520            _ => panic!("expected Timestamp"),
521        }
522    }
523
524    #[test]
525    fn tcp_options_sack() {
526        let mut data = sample_tcp_syn();
527        // SACK: kind=5, len=18, 2 blocks (each 8 bytes: left_edge + right_edge)
528        data.extend_from_slice(&[OPT_SACK, 18]);
529        // Block 1: left=100, right=200
530        data.extend_from_slice(&100u32.to_be_bytes());
531        data.extend_from_slice(&200u32.to_be_bytes());
532        // Block 2: left=300, right=400
533        data.extend_from_slice(&300u32.to_be_bytes());
534        data.extend_from_slice(&400u32.to_be_bytes());
535        // Pad to 4-byte boundary (18 + 2 = 20)
536        data.push(OPT_NOOP);
537        data.push(OPT_NOOP);
538        data[12] = 0xA0; // data_offset=10 (40 bytes)
539
540        let pkt = TcpPacket::new(&data).unwrap();
541        let opts: Vec<_> = pkt.options().collect();
542        match &opts[0] {
543            TcpOption::Sack(blocks) => {
544                assert_eq!(blocks.len(), 2);
545                assert_eq!(blocks[0], (100, 200));
546                assert_eq!(blocks[1], (300, 400));
547            }
548            _ => panic!("expected SACK"),
549        }
550    }
551
552    #[test]
553    fn tcp_sack_permitted() {
554        let mut data = sample_tcp_syn();
555        data.extend_from_slice(&[OPT_SACK_PERMITTED, 2]);
556        data.push(OPT_NOOP);
557        data.push(OPT_NOOP);
558        data[12] = 0x60; // data_offset=6
559
560        let pkt = TcpPacket::new(&data).unwrap();
561        let opts: Vec<_> = pkt.options().collect();
562        assert!(matches!(opts[0], TcpOption::SackPermitted));
563    }
564
565    #[test]
566    fn tcp_build_and_verify_checksum() {
567        let src = Ipv4Addr::new(10, 0, 0, 1);
568        let dst = Ipv4Addr::new(10, 0, 0, 2);
569        let pkt_bytes = TcpPacketBuilder::new()
570            .source_port(12345)
571            .destination_port(80)
572            .sequence(1000)
573            .syn(true)
574            .window(65535)
575            .mss(1460)
576            .build(src, dst);
577
578        let pkt = TcpPacket::new(&pkt_bytes).unwrap();
579        assert!(pkt.syn());
580        assert_eq!(pkt.source_port(), 12345);
581        assert_eq!(pkt.destination_port(), 80);
582        assert_eq!(pkt.sequence(), 1000);
583        assert!(pkt.verify_checksum(src, dst));
584    }
585
586    #[test]
587    fn tcp_build_roundtrip_with_options() {
588        let src = Ipv4Addr::new(192, 168, 1, 1);
589        let dst = Ipv4Addr::new(192, 168, 1, 2);
590        let pkt_bytes = TcpPacketBuilder::new()
591            .source_port(54321)
592            .destination_port(443)
593            .sequence(0x12345678)
594            .ack_number(0x87654321)
595            .syn(true)
596            .ack(true)
597            .window(8192)
598            .mss(1460)
599            .window_scale(7)
600            .sack_permitted()
601            .timestamp(0x100, 0x200)
602            .payload(&[0x01, 0x02, 0x03])
603            .build(src, dst);
604
605        let pkt = TcpPacket::new(&pkt_bytes).unwrap();
606        assert_eq!(pkt.source_port(), 54321);
607        assert_eq!(pkt.destination_port(), 443);
608        assert_eq!(pkt.sequence(), 0x12345678);
609        assert_eq!(pkt.ack_number(), 0x87654321);
610        assert!(pkt.syn());
611        assert!(pkt.ack());
612        assert_eq!(pkt.window(), 8192);
613        assert_eq!(pkt.payload(), &[0x01, 0x02, 0x03]);
614        assert!(pkt.verify_checksum(src, dst));
615
616        // Check options
617        let opts: Vec<_> = pkt.options().collect();
618        assert!(opts.iter().any(|o| matches!(o, TcpOption::MaximumSegmentSize(1460))));
619        assert!(opts.iter().any(|o| matches!(o, TcpOption::WindowScale(7))));
620        assert!(opts.iter().any(|o| matches!(o, TcpOption::SackPermitted)));
621        assert!(opts.iter().any(|o| matches!(o, TcpOption::Timestamp { ts_val: 0x100, ts_ecr: 0x200 })));
622    }
623
624    #[test]
625    fn tcp_end_of_list_option() {
626        let mut data = sample_tcp_syn();
627        data.push(OPT_END); // End of list
628        data.extend_from_slice(&[OPT_MSS, 4, 0x05, 0xB4]); // should not be parsed
629        data.push(OPT_NOOP);
630        data.push(OPT_NOOP);
631        data[12] = 0x60;
632
633        let pkt = TcpPacket::new(&data).unwrap();
634        let opts: Vec<_> = pkt.options().collect();
635        assert_eq!(opts.len(), 1);
636        assert!(matches!(opts[0], TcpOption::EndOfList));
637    }
638}
wireforge_core/tcp.rs

wireforge_core/
tcp.rs
