Skip to main content

wireforge_core/
udp.rs

1//! UDP packet parser and builder.
2
3use alloc::vec;
4use alloc::vec::Vec;
5use core::net::Ipv4Addr;
6
7use crate::util::{pseudo_header_checksum, read_u16be, write_u16be};
8
9pub const UDP_HEADER_LEN: usize = 8;
10
11/// Zero-copy UDP packet parser.
12#[derive(Debug, Clone)]
13pub struct UdpPacket<'a> {
14    buf: &'a [u8],
15}
16
17impl<'a> UdpPacket<'a> {
18    pub fn new(buf: &'a [u8]) -> Option<Self> {
19        if buf.len() < UDP_HEADER_LEN {
20            return None;
21        }
22        Some(Self { buf })
23    }
24
25    #[inline]
26    pub fn source_port(&self) -> u16 {
27        read_u16be(&self.buf[..2])
28    }
29
30    #[inline]
31    pub fn destination_port(&self) -> u16 {
32        read_u16be(&self.buf[2..4])
33    }
34
35    #[inline]
36    pub fn length(&self) -> u16 {
37        read_u16be(&self.buf[4..6])
38    }
39
40    #[inline]
41    pub fn checksum(&self) -> u16 {
42        read_u16be(&self.buf[6..8])
43    }
44
45    #[inline]
46    pub fn payload(&self) -> &'a [u8] {
47        &self.buf[UDP_HEADER_LEN..]
48    }
49
50    /// Verify the UDP checksum (includes IPv4 pseudo-header).
51    /// Returns true for zero checksum (allowed in IPv4 UDP).
52    pub fn verify_checksum(&self, src: Ipv4Addr, dst: Ipv4Addr) -> bool {
53        if self.checksum() == 0 {
54            return true; // zero checksum means "not computed" per RFC 768
55        }
56        pseudo_header_checksum(
57            &src.octets(),
58            &dst.octets(),
59            17, // UDP protocol number
60            self.buf,
61        ) == 0
62    }
63}
64
65// ---------------------------------------------------------------------------
66// Builder
67// ---------------------------------------------------------------------------
68
69pub struct UdpPacketBuilder {
70    buf: Vec<u8>,
71    payload: Option<Vec<u8>>,
72}
73
74impl Default for UdpPacketBuilder {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80impl UdpPacketBuilder {
81    pub fn new() -> Self {
82        Self { buf: vec![0u8; UDP_HEADER_LEN], payload: None }
83    }
84
85    pub fn source_port(mut self, port: u16) -> Self {
86        write_u16be(&mut self.buf[..2], port);
87        self
88    }
89
90    pub fn destination_port(mut self, port: u16) -> Self {
91        write_u16be(&mut self.buf[2..4], port);
92        self
93    }
94
95    pub fn payload(mut self, data: &[u8]) -> Self {
96        self.payload = Some(data.to_vec());
97        self
98    }
99
100    /// Build the UDP packet, computing the checksum (IPv4 pseudo-header).
101    pub fn build(mut self, src: Ipv4Addr, dst: Ipv4Addr) -> Vec<u8> {
102        self.buf[6] = 0;
103        self.buf[7] = 0;
104
105        let mut packet = self.buf;
106        if let Some(ref p) = self.payload {
107            packet.extend_from_slice(p);
108        }
109        let total_len = packet.len();
110        write_u16be(&mut packet[4..6], total_len as u16);
111
112        let csum = pseudo_header_checksum(
113            &src.octets(),
114            &dst.octets(),
115            17, // UDP
116            &packet,
117        );
118        write_u16be(&mut packet[6..8], csum);
119        packet
120    }
121}
122
123// ---------------------------------------------------------------------------
124// Tests
125// ---------------------------------------------------------------------------
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn parse_udp() {
133        let data: &[u8] = &[
134            0x12, 0x34, // src port = 4660
135            0x00, 0x50, // dst port = 80
136            0x00, 0x0C, // length = 12
137            0x00, 0x00, // checksum = 0
138            0xAA, 0xBB, 0xCC, 0xDD, // payload
139        ];
140        let pkt = UdpPacket::new(data).unwrap();
141        assert_eq!(pkt.source_port(), 4660);
142        assert_eq!(pkt.destination_port(), 80);
143        assert_eq!(pkt.length(), 12);
144        assert_eq!(pkt.checksum(), 0);
145        assert_eq!(pkt.payload(), &[0xAA, 0xBB, 0xCC, 0xDD]);
146    }
147
148    #[test]
149    fn parse_udp_too_short() {
150        assert!(UdpPacket::new(&[]).is_none());
151        assert!(UdpPacket::new(&[0u8; 7]).is_none());
152        assert!(UdpPacket::new(&[0u8; 8]).is_some());
153    }
154
155    #[test]
156    fn udp_zero_checksum() {
157        let data: &[u8] = &[
158            0x00, 0x00, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00,
159        ];
160        let pkt = UdpPacket::new(data).unwrap();
161        assert!(pkt.verify_checksum(
162            Ipv4Addr::new(10, 0, 0, 1),
163            Ipv4Addr::new(10, 0, 0, 2),
164        ));
165    }
166
167    #[test]
168    fn udp_build_and_verify_checksum() {
169        let src = Ipv4Addr::new(192, 168, 1, 1);
170        let dst = Ipv4Addr::new(192, 168, 1, 2);
171        let pkt_bytes = UdpPacketBuilder::new()
172            .source_port(12345)
173            .destination_port(53)
174            .payload(&[0x01, 0x02, 0x03, 0x04])
175            .build(src, dst);
176
177        let pkt = UdpPacket::new(&pkt_bytes).unwrap();
178        assert_eq!(pkt.source_port(), 12345);
179        assert_eq!(pkt.destination_port(), 53);
180        assert_eq!(pkt.length(), 12);
181        assert_eq!(pkt.payload(), &[0x01, 0x02, 0x03, 0x04]);
182        assert!(pkt.verify_checksum(src, dst));
183    }
184}