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

wireforge_core/
tcp.rs
