1use alloc::vec;
4use alloc::vec::Vec;
5use core::net::Ipv4Addr;
6
7use crate::types::IpProtocol;
8use crate::util::{internet_checksum, seal_internet_checksum, read_u16be, read_u32be, write_u16be, write_u32be, BuildError, PacketBuilder};
9
10pub const IPV4_MIN_HEADER_LEN: usize = 20;
12
13#[derive(Debug, Clone)]
15pub struct Ipv4Packet<'a> {
16 buf: &'a [u8],
17}
18
19impl<'a> Ipv4Packet<'a> {
20 pub fn new(buf: &'a [u8]) -> Option<Self> {
23 if buf.len() < IPV4_MIN_HEADER_LEN {
24 return None;
25 }
26 let ihl = (buf[0] & 0x0F) as usize * 4;
27 if ihl < IPV4_MIN_HEADER_LEN || buf.len() < ihl {
28 return None;
29 }
30 Some(Self { buf })
31 }
32
33 #[inline]
35 pub fn as_bytes(&self) -> &'a [u8] { self.buf }
36
37 #[inline]
38 pub fn version(&self) -> u8 {
39 self.buf[0] >> 4
40 }
41
42 #[inline]
43 pub fn ihl(&self) -> u8 {
44 self.buf[0] & 0x0F
45 }
46
47 #[inline]
48 pub fn dscp(&self) -> u8 {
49 self.buf[1] >> 2
50 }
51
52 #[inline]
53 pub fn ecn(&self) -> u8 {
54 self.buf[1] & 0x03
55 }
56
57 #[inline]
58 pub fn total_length(&self) -> u16 {
59 read_u16be(&self.buf[2..4])
60 }
61
62 #[inline]
63 pub fn identification(&self) -> u16 {
64 read_u16be(&self.buf[4..6])
65 }
66
67 #[inline]
68 pub fn flags(&self) -> u8 {
69 self.buf[6] >> 5
70 }
71
72 #[inline]
73 pub fn dont_fragment(&self) -> bool {
74 self.flags() & 0x02 != 0
75 }
76
77 #[inline]
78 pub fn more_fragments(&self) -> bool {
79 self.flags() & 0x01 != 0
80 }
81
82 #[inline]
83 pub fn fragment_offset(&self) -> u16 {
84 read_u16be(&self.buf[6..8]) & 0x1FFF
85 }
86
87 #[inline]
88 pub fn ttl(&self) -> u8 {
89 self.buf[8]
90 }
91
92 #[inline]
93 pub fn protocol(&self) -> IpProtocol {
94 IpProtocol::from(self.buf[9])
95 }
96
97 #[inline]
98 pub fn checksum(&self) -> u16 {
99 read_u16be(&self.buf[10..12])
100 }
101
102 #[inline]
103 pub fn source(&self) -> Ipv4Addr {
104 Ipv4Addr::from_bits(read_u32be(&self.buf[12..16]))
105 }
106
107 #[inline]
108 pub fn destination(&self) -> Ipv4Addr {
109 Ipv4Addr::from_bits(read_u32be(&self.buf[16..20]))
110 }
111
112 #[inline]
114 pub fn options(&self) -> &'a [u8] {
115 let hdr_len = self.ihl() as usize * 4;
116 if hdr_len > IPV4_MIN_HEADER_LEN {
117 &self.buf[IPV4_MIN_HEADER_LEN..hdr_len]
118 } else {
119 &[]
120 }
121 }
122
123 #[inline]
125 pub fn payload(&self) -> &'a [u8] {
126 &self.buf[self.ihl() as usize * 4..]
127 }
128
129 #[inline]
131 pub fn header_length(&self) -> usize {
132 self.ihl() as usize * 4
133 }
134
135 pub fn verify_checksum(&self) -> bool {
137 internet_checksum(&self.buf[..self.header_length()]) == 0
138 }
139}
140
141#[derive(Debug, Clone)]
147pub struct Ipv4PacketBuilder {
148 buf: Vec<u8>,
149 payload: Option<Vec<u8>>,
150}
151
152impl Default for Ipv4PacketBuilder {
153 fn default() -> Self {
154 Self::new()
155 }
156}
157
158impl Ipv4PacketBuilder {
159 pub fn new() -> Self {
161 let mut buf = vec![0u8; IPV4_MIN_HEADER_LEN];
162 buf[0] = 0x45; buf[8] = 64; Self { buf, payload: None }
165 }
166
167 pub fn dscp(mut self, dscp: u8) -> Self {
168 self.buf[1] = (self.buf[1] & 0x03) | (dscp << 2);
169 self
170 }
171
172 pub fn ecn(mut self, ecn: u8) -> Self {
173 self.buf[1] = (self.buf[1] & 0xFC) | (ecn & 0x03);
174 self
175 }
176
177 pub fn identification(mut self, id: u16) -> Self {
178 write_u16be(&mut self.buf[4..6], id);
179 self
180 }
181
182 pub fn dont_fragment(mut self, df: bool) -> Self {
183 if df {
184 self.buf[6] |= 0x40;
185 } else {
186 self.buf[6] &= !0x40;
187 }
188 self
189 }
190
191 pub fn more_fragments(mut self, mf: bool) -> Self {
192 if mf {
193 self.buf[6] |= 0x20;
194 } else {
195 self.buf[6] &= !0x20;
196 }
197 self
198 }
199
200 pub fn fragment_offset(mut self, offset: u16) -> Self {
201 let current = read_u16be(&self.buf[6..8]);
202 let new = (current & 0xE000) | (offset & 0x1FFF);
203 write_u16be(&mut self.buf[6..8], new);
204 self
205 }
206
207 pub fn ttl(mut self, ttl: u8) -> Self {
208 self.buf[8] = ttl;
209 self
210 }
211
212 pub fn protocol(mut self, proto: IpProtocol) -> Self {
213 self.buf[9] = proto.into();
214 self
215 }
216
217 pub fn source(mut self, addr: Ipv4Addr) -> Self {
218 write_u32be(&mut self.buf[12..16], addr.to_bits());
219 self
220 }
221
222 pub fn destination(mut self, addr: Ipv4Addr) -> Self {
223 write_u32be(&mut self.buf[16..20], addr.to_bits());
224 self
225 }
226
227 pub fn payload(mut self, data: &[u8]) -> Result<Self, BuildError> {
228 if data.len() > 65535 - self.buf.len() {
229 return Err(BuildError::PayloadTooLarge {
230 max: 65535 - self.buf.len(),
231 actual: data.len(),
232 });
233 }
234 self.payload = Some(data.to_vec());
235 Ok(self)
236 }
237
238 pub fn build(self) -> Vec<u8> {
240 <Self as PacketBuilder>::build(self)
241 }
242}
243
244impl PacketBuilder for Ipv4PacketBuilder {
245 type Output = Vec<u8>;
246
247 fn build(mut self) -> Vec<u8> {
248 let total_len = self.buf.len() + self.payload.as_ref().map_or(0, |p| p.len());
249 write_u16be(&mut self.buf[2..4], total_len as u16);
250
251 seal_internet_checksum(&mut self.buf, 10);
252
253 let mut packet = self.buf;
254 if let Some(p) = self.payload {
255 packet.extend_from_slice(&p);
256 }
257 packet
258 }
259}
260
261#[cfg(test)]
262mod tests {
263 use super::*;
264
265 const SAMPLE_IPV4: &[u8] = &[
267 0x45, 0x00, 0x00, 0x3C, 0x1C, 0x46, 0x40, 0x00,
268 0x40, 0x06, 0x00, 0x00, 0xAC, 0x10, 0x0A, 0x63,
269 0xAC, 0x10, 0x0A, 0x0C,
270 ];
271
272 #[test]
273 fn parse_ipv4() {
274 let pkt = Ipv4Packet::new(SAMPLE_IPV4).unwrap();
275 assert_eq!(pkt.version(), 4);
276 assert_eq!(pkt.ihl(), 5);
277 assert_eq!(pkt.header_length(), 20);
278 assert_eq!(pkt.protocol(), IpProtocol::Tcp);
279 assert_eq!(pkt.ttl(), 64);
280 assert_eq!(pkt.source(), Ipv4Addr::new(172, 16, 10, 99));
281 assert_eq!(pkt.destination(), Ipv4Addr::new(172, 16, 10, 12));
282 assert!(pkt.dont_fragment());
283 assert!(!pkt.more_fragments());
284 }
285
286 #[test]
287 fn parse_too_short() {
288 assert!(Ipv4Packet::new(&[0u8; 19]).is_none());
289 assert!(Ipv4Packet::new(&[]).is_none());
290 }
291
292 #[test]
293 fn parse_ihl_exceeds_buffer() {
294 let mut data = [0u8; 22];
296 data[0] = 0x46; assert!(Ipv4Packet::new(&data).is_none());
298 }
299
300 #[test]
301 fn verify_checksum() {
302 let pkt_bytes = Ipv4PacketBuilder::new()
304 .source(Ipv4Addr::new(10, 0, 0, 1))
305 .destination(Ipv4Addr::new(10, 0, 0, 2))
306 .protocol(IpProtocol::Tcp)
307 .ttl(64)
308 .build();
309
310 let pkt = Ipv4Packet::new(&pkt_bytes).unwrap();
311 assert!(pkt.verify_checksum());
312 }
313
314 #[test]
315 fn build_and_parse_roundtrip() {
316 let pkt_bytes = Ipv4PacketBuilder::new()
317 .source(Ipv4Addr::new(192, 168, 1, 1))
318 .destination(Ipv4Addr::new(192, 168, 1, 2))
319 .protocol(IpProtocol::Udp)
320 .ttl(128)
321 .dont_fragment(true)
322 .identification(0x1234)
323 .build();
324
325 let pkt = Ipv4Packet::new(&pkt_bytes).unwrap();
326 assert_eq!(pkt.source(), Ipv4Addr::new(192, 168, 1, 1));
327 assert_eq!(pkt.destination(), Ipv4Addr::new(192, 168, 1, 2));
328 assert_eq!(pkt.protocol(), IpProtocol::Udp);
329 assert_eq!(pkt.ttl(), 128);
330 assert!(pkt.dont_fragment());
331 assert_eq!(pkt.identification(), 0x1234);
332 assert!(pkt.verify_checksum());
333 }
334
335 #[test]
336 fn parse_with_options() {
337 let mut data = vec![0x46u8, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00,
339 0x40, 0x06, 0x00, 0x00,
340 0xC0, 0xA8, 0x01, 0x01, 0xC0, 0xA8, 0x01, 0x02,
341 0x01, 0x02, 0x03, 0x04, 0xAA, 0xBB, ];
344 let csum = internet_checksum(&data[..24]);
346 data[10..12].copy_from_slice(&csum.to_be_bytes());
347
348 let pkt = Ipv4Packet::new(&data).unwrap();
349 assert_eq!(pkt.ihl(), 6);
350 assert_eq!(pkt.header_length(), 24);
351 assert_eq!(pkt.options(), &[0x01, 0x02, 0x03, 0x04]);
352 assert_eq!(pkt.payload(), &[0xAA, 0xBB]);
353 assert!(pkt.verify_checksum());
354 }
355}