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::types::IpProtocol;
8use crate::util::{pseudo_header_checksum, seal_transport_checksum, read_u16be, write_u16be};
9
10pub const UDP_HEADER_LEN: usize = 8;
11
12/// Zero-copy UDP packet parser.
13#[derive(Debug, Clone)]
14pub struct UdpPacket<'a> {
15    buf: &'a [u8],
16}
17
18impl<'a> UdpPacket<'a> {
19    crate::parser_new!(UDP_HEADER_LEN);
20
21    /// Raw bytes backing this packet.
22    #[inline]
23    pub fn as_bytes(&self) -> &'a [u8] { self.buf }
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            IpProtocol::Udp.into(),
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(self, src: Ipv4Addr, dst: Ipv4Addr) -> Vec<u8> {
102        let mut packet = self.buf;
103        if let Some(ref p) = self.payload {
104            packet.extend_from_slice(p);
105        }
106        let total_len = packet.len();
107        write_u16be(&mut packet[4..6], total_len as u16);
108        seal_transport_checksum(&mut packet, 6, &src.octets(), &dst.octets(), IpProtocol::Udp.into());
109        packet
110    }
111}
112
113// ---------------------------------------------------------------------------
114// Tests
115// ---------------------------------------------------------------------------
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn parse_udp() {
123        let data: &[u8] = &[
124            0x12, 0x34, // src port = 4660
125            0x00, 0x50, // dst port = 80
126            0x00, 0x0C, // length = 12
127            0x00, 0x00, // checksum = 0
128            0xAA, 0xBB, 0xCC, 0xDD, // payload
129        ];
130        let pkt = UdpPacket::new(data).unwrap();
131        assert_eq!(pkt.source_port(), 4660);
132        assert_eq!(pkt.destination_port(), 80);
133        assert_eq!(pkt.length(), 12);
134        assert_eq!(pkt.checksum(), 0);
135        assert_eq!(pkt.payload(), &[0xAA, 0xBB, 0xCC, 0xDD]);
136    }
137
138    #[test]
139    fn parse_udp_too_short() {
140        assert!(UdpPacket::new(&[]).is_none());
141        assert!(UdpPacket::new(&[0u8; 7]).is_none());
142        assert!(UdpPacket::new(&[0u8; 8]).is_some());
143    }
144
145    #[test]
146    fn udp_zero_checksum() {
147        let data: &[u8] = &[
148            0x00, 0x00, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00,
149        ];
150        let pkt = UdpPacket::new(data).unwrap();
151        assert!(pkt.verify_checksum(
152            Ipv4Addr::new(10, 0, 0, 1),
153            Ipv4Addr::new(10, 0, 0, 2),
154        ));
155    }
156
157    #[test]
158    fn udp_build_and_verify_checksum() {
159        let src = Ipv4Addr::new(192, 168, 1, 1);
160        let dst = Ipv4Addr::new(192, 168, 1, 2);
161        let pkt_bytes = UdpPacketBuilder::new()
162            .source_port(12345)
163            .destination_port(53)
164            .payload(&[0x01, 0x02, 0x03, 0x04])
165            .build(src, dst);
166
167        let pkt = UdpPacket::new(&pkt_bytes).unwrap();
168        assert_eq!(pkt.source_port(), 12345);
169        assert_eq!(pkt.destination_port(), 53);
170        assert_eq!(pkt.length(), 12);
171        assert_eq!(pkt.payload(), &[0x01, 0x02, 0x03, 0x04]);
172        assert!(pkt.verify_checksum(src, dst));
173    }
174}