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}