ippacket/
udp.rs

1use std::fmt;
2use std::io;
3
4use byteorder::NetworkEndian;
5
6use crate::bytes::{Bytes, Checksum};
7use crate::ip::IpHeader;
8
9#[derive(Clone)]
10pub struct UdpHeader(Bytes);
11
12impl UdpHeader {
13    pub fn with_bytes(bytes: Bytes) -> io::Result<(UdpHeader, Bytes)> {
14        let mut header = UdpHeader(bytes);
15        let mut payload = try_split!(header.0, UdpHeader::len());
16        if let Some(data_len) = header.data_len() {
17            try_split!(payload, data_len);
18        }
19        Ok((header, payload))
20    }
21
22    pub fn len() -> usize {
23        8
24    }
25
26    pub fn src(&self) -> u16 {
27        self.0.read_u16::<NetworkEndian>(0).unwrap()
28    }
29
30    pub fn dest(&self) -> u16 {
31        self.0.read_u16::<NetworkEndian>(2).unwrap()
32    }
33
34    pub fn udp_len(&self) -> usize {
35        self.0.read_u16::<NetworkEndian>(4).unwrap() as usize
36    }
37
38    pub fn data_len(&self) -> Option<usize> {
39        let udp_len = self.udp_len();
40        if udp_len > UdpHeader::len() {
41            Some(udp_len - UdpHeader::len())
42        } else {
43            None
44        }
45    }
46
47    fn checksum(&self) -> u16 {
48        self.0.read_u16::<NetworkEndian>(6).unwrap()
49    }
50
51    pub fn checksum_valid<V: Iterator<Item = u16>>(&self, header: &IpHeader, data: V) -> bool {
52        self.checksum() == self.calculated_checksum(header, data)
53    }
54
55    pub fn calculated_checksum<V: Iterator<Item = u16>>(&self, header: &IpHeader, data: V) -> u16 {
56        let pseudo = header.pseudo_iter(self.udp_len());
57        self.0
58            .slice(0, 6)
59            .pair_iter()
60            .chain(pseudo)
61            .chain(data)
62            .checksum()
63    }
64
65    pub fn set_src(&mut self, src: u16) {
66        self.0.write_u16::<NetworkEndian>(0, src).unwrap();
67    }
68
69    pub fn set_dest(&mut self, dest: u16) {
70        self.0.write_u16::<NetworkEndian>(2, dest).unwrap();
71    }
72
73    pub fn set_udp_len(&mut self, len: usize) {
74        self.0.write_u16::<NetworkEndian>(4, len as u16).unwrap();
75    }
76
77    pub fn set_data_len(&mut self, len: usize) {
78        self.set_udp_len(len + UdpHeader::len());
79    }
80
81    fn set_checksum(&mut self, checksum: u16) {
82        self.0.write_u16::<NetworkEndian>(6, checksum).unwrap();
83    }
84
85    pub fn calculate_checksum<V: Iterator<Item = u16>>(&mut self, header: &IpHeader, data: V) {
86        let checksum = self.calculated_checksum(header, data);
87        self.set_checksum(checksum);
88    }
89}
90
91impl fmt::Debug for UdpHeader {
92    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
93        fmt.debug_struct("UdpHeader")
94            .field("src", &self.src())
95            .field("dest", &self.dest())
96            .field("udp_len", &self.udp_len())
97            .finish()
98    }
99}
100
101#[derive(Default, Debug, Clone)]
102pub struct UdpHeaderBuilder {
103    src: Option<u16>,
104    dest: Option<u16>,
105}
106
107impl UdpHeaderBuilder {
108    pub fn len() -> usize {
109        UdpHeader::len()
110    }
111
112    pub fn src(mut self, src: u16) -> UdpHeaderBuilder {
113        self.src = Some(src);
114        self
115    }
116
117    pub fn dest(mut self, dest: u16) -> UdpHeaderBuilder {
118        self.dest = Some(dest);
119        self
120    }
121
122    pub fn build(self, bytes: Bytes) -> (UdpHeader, Bytes) {
123        let src = self.src.unwrap_or_else(|| unimplemented!());
124        let dest = self.dest.unwrap_or_else(|| unimplemented!());
125
126        let (mut header, remaining) = UdpHeader::with_bytes(bytes).unwrap();
127        header.set_src(src);
128        header.set_dest(dest);
129        (header, remaining)
130    }
131}