Skip to main content

wireforge_core/
ipv4.rs

1//! IPv4 packet parser and builder.
2
3use alloc::vec;
4use alloc::vec::Vec;
5use core::net::Ipv4Addr;
6
7use crate::types::IpProtocol;
8use crate::util::{internet_checksum, seal_internet_checksum, read_u16be, read_u32be, write_u16be, write_u32be, BuildError, PacketBuilder};
9
10/// Minimum IPv4 header length in bytes (without options).
11pub const IPV4_MIN_HEADER_LEN: usize = 20;
12
13/// Zero-copy IPv4 packet parser.
14#[derive(Debug, Clone)]
15pub struct Ipv4Packet<'a> {
16    buf: &'a [u8],
17}
18
19impl<'a> Ipv4Packet<'a> {
20    /// Attempt to parse an IPv4 packet. Returns `None` if `buf` is shorter than
21    /// the minimum header length, or if the IHL-declared header length exceeds `buf`.
22    pub fn new(buf: &'a [u8]) -> Option<Self> {
23        if buf.len() < IPV4_MIN_HEADER_LEN {
24            return None;
25        }
26        let ihl = (buf[0] & 0x0F) as usize * 4;
27        if ihl < IPV4_MIN_HEADER_LEN || buf.len() < ihl {
28            return None;
29        }
30        Some(Self { buf })
31    }
32
33    /// Raw bytes backing this packet.
34    #[inline]
35    pub fn as_bytes(&self) -> &'a [u8] { self.buf }
36
37    #[inline]
38    pub fn version(&self) -> u8 {
39        self.buf[0] >> 4
40    }
41
42    #[inline]
43    pub fn ihl(&self) -> u8 {
44        self.buf[0] & 0x0F
45    }
46
47    #[inline]
48    pub fn dscp(&self) -> u8 {
49        self.buf[1] >> 2
50    }
51
52    #[inline]
53    pub fn ecn(&self) -> u8 {
54        self.buf[1] & 0x03
55    }
56
57    #[inline]
58    pub fn total_length(&self) -> u16 {
59        read_u16be(&self.buf[2..4])
60    }
61
62    #[inline]
63    pub fn identification(&self) -> u16 {
64        read_u16be(&self.buf[4..6])
65    }
66
67    #[inline]
68    pub fn flags(&self) -> u8 {
69        self.buf[6] >> 5
70    }
71
72    #[inline]
73    pub fn dont_fragment(&self) -> bool {
74        self.flags() & 0x02 != 0
75    }
76
77    #[inline]
78    pub fn more_fragments(&self) -> bool {
79        self.flags() & 0x01 != 0
80    }
81
82    #[inline]
83    pub fn fragment_offset(&self) -> u16 {
84        read_u16be(&self.buf[6..8]) & 0x1FFF
85    }
86
87    #[inline]
88    pub fn ttl(&self) -> u8 {
89        self.buf[8]
90    }
91
92    #[inline]
93    pub fn protocol(&self) -> IpProtocol {
94        IpProtocol::from(self.buf[9])
95    }
96
97    #[inline]
98    pub fn checksum(&self) -> u16 {
99        read_u16be(&self.buf[10..12])
100    }
101
102    #[inline]
103    pub fn source(&self) -> Ipv4Addr {
104        Ipv4Addr::from_bits(read_u32be(&self.buf[12..16]))
105    }
106
107    #[inline]
108    pub fn destination(&self) -> Ipv4Addr {
109        Ipv4Addr::from_bits(read_u32be(&self.buf[16..20]))
110    }
111
112    /// IP options bytes (if any).
113    #[inline]
114    pub fn options(&self) -> &'a [u8] {
115        let hdr_len = self.ihl() as usize * 4;
116        if hdr_len > IPV4_MIN_HEADER_LEN {
117            &self.buf[IPV4_MIN_HEADER_LEN..hdr_len]
118        } else {
119            &[]
120        }
121    }
122
123    /// Payload following the IPv4 header.
124    #[inline]
125    pub fn payload(&self) -> &'a [u8] {
126        &self.buf[self.ihl() as usize * 4..]
127    }
128
129    /// Total header length in bytes (IHL × 4).
130    #[inline]
131    pub fn header_length(&self) -> usize {
132        self.ihl() as usize * 4
133    }
134
135    /// Verify the IPv4 header checksum.
136    pub fn verify_checksum(&self) -> bool {
137        internet_checksum(&self.buf[..self.header_length()]) == 0
138    }
139}
140
141// ---------------------------------------------------------------------------
142// Builder
143// ---------------------------------------------------------------------------
144
145/// Builder for IPv4 packets.
146#[derive(Debug, Clone)]
147pub struct Ipv4PacketBuilder {
148    buf: Vec<u8>,
149    payload: Option<Vec<u8>>,
150}
151
152impl Default for Ipv4PacketBuilder {
153    fn default() -> Self {
154        Self::new()
155    }
156}
157
158impl Ipv4PacketBuilder {
159    /// Create a new builder with a default IPv4 header (version=4, IHL=5, TTL=64).
160    pub fn new() -> Self {
161        let mut buf = vec![0u8; IPV4_MIN_HEADER_LEN];
162        buf[0] = 0x45; // version=4, IHL=5
163        buf[8] = 64; // TTL=64
164        Self { buf, payload: None }
165    }
166
167    pub fn dscp(mut self, dscp: u8) -> Self {
168        self.buf[1] = (self.buf[1] & 0x03) | (dscp << 2);
169        self
170    }
171
172    pub fn ecn(mut self, ecn: u8) -> Self {
173        self.buf[1] = (self.buf[1] & 0xFC) | (ecn & 0x03);
174        self
175    }
176
177    pub fn identification(mut self, id: u16) -> Self {
178        write_u16be(&mut self.buf[4..6], id);
179        self
180    }
181
182    pub fn dont_fragment(mut self, df: bool) -> Self {
183        if df {
184            self.buf[6] |= 0x40;
185        } else {
186            self.buf[6] &= !0x40;
187        }
188        self
189    }
190
191    pub fn more_fragments(mut self, mf: bool) -> Self {
192        if mf {
193            self.buf[6] |= 0x20;
194        } else {
195            self.buf[6] &= !0x20;
196        }
197        self
198    }
199
200    pub fn fragment_offset(mut self, offset: u16) -> Self {
201        let current = read_u16be(&self.buf[6..8]);
202        let new = (current & 0xE000) | (offset & 0x1FFF);
203        write_u16be(&mut self.buf[6..8], new);
204        self
205    }
206
207    pub fn ttl(mut self, ttl: u8) -> Self {
208        self.buf[8] = ttl;
209        self
210    }
211
212    pub fn protocol(mut self, proto: IpProtocol) -> Self {
213        self.buf[9] = proto.into();
214        self
215    }
216
217    pub fn source(mut self, addr: Ipv4Addr) -> Self {
218        write_u32be(&mut self.buf[12..16], addr.to_bits());
219        self
220    }
221
222    pub fn destination(mut self, addr: Ipv4Addr) -> Self {
223        write_u32be(&mut self.buf[16..20], addr.to_bits());
224        self
225    }
226
227    pub fn payload(mut self, data: &[u8]) -> Result<Self, BuildError> {
228        if data.len() > 65535 - self.buf.len() {
229            return Err(BuildError::PayloadTooLarge {
230                max: 65535 - self.buf.len(),
231                actual: data.len(),
232            });
233        }
234        self.payload = Some(data.to_vec());
235        Ok(self)
236    }
237
238    /// Build the IPv4 packet, computing total_length and header checksum automatically.
239    pub fn build(self) -> Vec<u8> {
240        <Self as PacketBuilder>::build(self)
241    }
242}
243
244impl PacketBuilder for Ipv4PacketBuilder {
245    type Output = Vec<u8>;
246
247    fn build(mut self) -> Vec<u8> {
248        let total_len = self.buf.len() + self.payload.as_ref().map_or(0, |p| p.len());
249        write_u16be(&mut self.buf[2..4], total_len as u16);
250
251        seal_internet_checksum(&mut self.buf, 10);
252
253        let mut packet = self.buf;
254        if let Some(p) = self.payload {
255            packet.extend_from_slice(&p);
256        }
257        packet
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    // Real IPv4 header: 192.168.10.99 → 172.16.10.12, TCP, TTL=64
266    const SAMPLE_IPV4: &[u8] = &[
267        0x45, 0x00, 0x00, 0x3C, 0x1C, 0x46, 0x40, 0x00,
268        0x40, 0x06, 0x00, 0x00, 0xAC, 0x10, 0x0A, 0x63,
269        0xAC, 0x10, 0x0A, 0x0C,
270    ];
271
272    #[test]
273    fn parse_ipv4() {
274        let pkt = Ipv4Packet::new(SAMPLE_IPV4).unwrap();
275        assert_eq!(pkt.version(), 4);
276        assert_eq!(pkt.ihl(), 5);
277        assert_eq!(pkt.header_length(), 20);
278        assert_eq!(pkt.protocol(), IpProtocol::Tcp);
279        assert_eq!(pkt.ttl(), 64);
280        assert_eq!(pkt.source(), Ipv4Addr::new(172, 16, 10, 99));
281        assert_eq!(pkt.destination(), Ipv4Addr::new(172, 16, 10, 12));
282        assert!(pkt.dont_fragment());
283        assert!(!pkt.more_fragments());
284    }
285
286    #[test]
287    fn parse_too_short() {
288        assert!(Ipv4Packet::new(&[0u8; 19]).is_none());
289        assert!(Ipv4Packet::new(&[]).is_none());
290    }
291
292    #[test]
293    fn parse_ihl_exceeds_buffer() {
294        // IHL=6 (24 bytes) but only 22 bytes provided
295        let mut data = [0u8; 22];
296        data[0] = 0x46; // version=4, IHL=6
297        assert!(Ipv4Packet::new(&data).is_none());
298    }
299
300    #[test]
301    fn verify_checksum() {
302        // Build a packet and verify its checksum
303        let pkt_bytes = Ipv4PacketBuilder::new()
304            .source(Ipv4Addr::new(10, 0, 0, 1))
305            .destination(Ipv4Addr::new(10, 0, 0, 2))
306            .protocol(IpProtocol::Tcp)
307            .ttl(64)
308            .build();
309
310        let pkt = Ipv4Packet::new(&pkt_bytes).unwrap();
311        assert!(pkt.verify_checksum());
312    }
313
314    #[test]
315    fn build_and_parse_roundtrip() {
316        let pkt_bytes = Ipv4PacketBuilder::new()
317            .source(Ipv4Addr::new(192, 168, 1, 1))
318            .destination(Ipv4Addr::new(192, 168, 1, 2))
319            .protocol(IpProtocol::Udp)
320            .ttl(128)
321            .dont_fragment(true)
322            .identification(0x1234)
323            .build();
324
325        let pkt = Ipv4Packet::new(&pkt_bytes).unwrap();
326        assert_eq!(pkt.source(), Ipv4Addr::new(192, 168, 1, 1));
327        assert_eq!(pkt.destination(), Ipv4Addr::new(192, 168, 1, 2));
328        assert_eq!(pkt.protocol(), IpProtocol::Udp);
329        assert_eq!(pkt.ttl(), 128);
330        assert!(pkt.dont_fragment());
331        assert_eq!(pkt.identification(), 0x1234);
332        assert!(pkt.verify_checksum());
333    }
334
335    #[test]
336    fn parse_with_options() {
337        // IPv4 header with 4 bytes of options (IHL=6, total header=24)
338        let mut data = vec![0x46u8, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00,
339            0x40, 0x06, 0x00, 0x00,
340            0xC0, 0xA8, 0x01, 0x01, 0xC0, 0xA8, 0x01, 0x02,
341            0x01, 0x02, 0x03, 0x04, // options
342            0xAA, 0xBB, // payload
343        ];
344        // Compute checksum
345        let csum = internet_checksum(&data[..24]);
346        data[10..12].copy_from_slice(&csum.to_be_bytes());
347
348        let pkt = Ipv4Packet::new(&data).unwrap();
349        assert_eq!(pkt.ihl(), 6);
350        assert_eq!(pkt.header_length(), 24);
351        assert_eq!(pkt.options(), &[0x01, 0x02, 0x03, 0x04]);
352        assert_eq!(pkt.payload(), &[0xAA, 0xBB]);
353        assert!(pkt.verify_checksum());
354    }
355}