1use 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#[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 #[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 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 IpProtocol::Udp.into(),
60 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(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#[cfg(test)]
118mod tests {
119 use super::*;
120
121 #[test]
122 fn parse_udp() {
123 let data: &[u8] = &[
124 0x12, 0x34, 0x00, 0x50, 0x00, 0x0C, 0x00, 0x00, 0xAA, 0xBB, 0xCC, 0xDD, ];
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}