1use 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#[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 pub fn verify_checksum(&self, src: Ipv4Addr, dst: Ipv4Addr) -> bool {
53 if self.checksum() == 0 {
54 return true; }
56 pseudo_header_checksum(
57 &src.octets(),
58 &dst.octets(),
59 17, self.buf,
61 ) == 0
62 }
63}
64
65pub 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 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, &packet,
117 );
118 write_u16be(&mut packet[6..8], csum);
119 packet
120 }
121}
122
123#[cfg(test)]
128mod tests {
129 use super::*;
130
131 #[test]
132 fn parse_udp() {
133 let data: &[u8] = &[
134 0x12, 0x34, 0x00, 0x50, 0x00, 0x0C, 0x00, 0x00, 0xAA, 0xBB, 0xCC, 0xDD, ];
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}