1 use crate::*;
2 use alloc::vec::Vec;
3 
4 #[derive(Clone)]
5 pub(crate) struct TestPacket {
6     pub link: Option<LinkHeader>,
7     pub vlan: Option<VlanHeader>,
8     pub net: Option<NetHeaders>,
9     pub transport: Option<TransportHeader>,
10 }
11 
12 impl TestPacket {
len(&self, payload: &[u8]) -> usize13     pub fn len(&self, payload: &[u8]) -> usize {
14         self.link.as_ref().map_or(0, |x| x.header_len())
15             + self.vlan.as_ref().map_or(0, |x| x.header_len())
16             + self.net.as_ref().map_or(0, |x| x.header_len())
17             + self.transport.as_ref().map_or(0, |x| x.header_len())
18             + payload.len()
19     }
20 
to_vec(&self, payload: &[u8]) -> Vec<u8>21     pub fn to_vec(&self, payload: &[u8]) -> Vec<u8> {
22         let mut result = Vec::with_capacity(self.len(payload));
23         if let Some(link) = &self.link {
24             link.write(&mut result).unwrap();
25         }
26         if let Some(vlan) = &self.vlan {
27             vlan.write(&mut result).unwrap();
28         }
29         if let Some(ip) = &self.net {
30             match ip {
31                 NetHeaders::Ipv4(ipv4, exts) => {
32                     ipv4.write_raw(&mut result).unwrap();
33                     exts.write(&mut result, ipv4.protocol).unwrap();
34                 }
35                 NetHeaders::Ipv6(ipv6, exts) => {
36                     ipv6.write(&mut result).unwrap();
37                     exts.write(&mut result, ipv6.next_header).unwrap();
38                 }
39             }
40         }
41         if let Some(transport) = &self.transport {
42             transport.write(&mut result).unwrap();
43         }
44         result.extend_from_slice(payload);
45         result
46     }
47 
set_ether_type(&mut self, ether_type: EtherType)48     pub fn set_ether_type(&mut self, ether_type: EtherType) {
49         if let Some(vlan) = &mut self.vlan {
50             use VlanHeader::*;
51             match vlan {
52                 Single(single) => {
53                     single.ether_type = ether_type;
54                 }
55                 Double(double) => {
56                     double.inner.ether_type = ether_type;
57                 }
58             }
59         } else if let Some(link) = &mut self.link {
60             match link {
61                 LinkHeader::Ethernet2(ethernet) => ethernet.ether_type = ether_type,
62                 LinkHeader::LinuxSll(linux_sll) => {
63                     linux_sll.protocol_type.change_value(ether_type.0)
64                 }
65             }
66         }
67     }
68 
set_payload_len(&mut self, payload_len: usize)69     pub fn set_payload_len(&mut self, payload_len: usize) {
70         use NetHeaders::*;
71         match &mut self.net {
72             None => {}
73             Some(Ipv4(ref mut header, ref mut exts)) => {
74                 header
75                     .set_payload_len(
76                         exts.header_len()
77                             + self.transport.as_ref().map_or(0, |t| t.header_len())
78                             + payload_len,
79                     )
80                     .unwrap();
81             }
82             Some(Ipv6(ref mut header, ref mut exts)) => {
83                 header
84                     .set_payload_length(
85                         exts.header_len()
86                             + self.transport.as_ref().map_or(0, |t| t.header_len())
87                             + payload_len,
88                     )
89                     .unwrap();
90             }
91         }
92 
93         use TransportHeader::*;
94         match &mut self.transport {
95             None => {}
96             Some(Udp(ref mut udp)) => {
97                 udp.length = udp.header_len_u16() + payload_len as u16;
98             }
99             Some(Tcp(_)) => {}
100             Some(Icmpv4(_)) => {}
101             Some(Icmpv6(_)) => {}
102         }
103     }
104 
105     /// Set the length relative to the end of the ip headers.
set_payload_le_from_ip_on(&mut self, payload_len_from_ip_on: isize)106     pub fn set_payload_le_from_ip_on(&mut self, payload_len_from_ip_on: isize) {
107         use NetHeaders::*;
108         match self.net.as_mut().unwrap() {
109             Ipv4(ref mut header, ref mut exts) => {
110                 header
111                     .set_payload_len((exts.header_len() as isize + payload_len_from_ip_on) as usize)
112                     .unwrap();
113             }
114             Ipv6(ref mut header, ref mut exts) => {
115                 header
116                     .set_payload_length(
117                         (exts.header_len() as isize + payload_len_from_ip_on) as usize,
118                     )
119                     .unwrap();
120             }
121         }
122     }
123 
is_ip_payload_fragmented(&self) -> bool124     pub fn is_ip_payload_fragmented(&self) -> bool {
125         self.net.as_ref().map_or(false, |net| match net {
126             NetHeaders::Ipv4(h, _) => h.is_fragmenting_payload(),
127             NetHeaders::Ipv6(_, e) => e.is_fragmenting_payload(),
128         })
129     }
130 }
131