Set correct packet length after defragmenting a packet (#2)

Co-authored-by: Adam Michalik <[email protected]>

Author: aaronddavies

src/iface/interface/ipv4.rs

@@ -101,7 +101,7 @@ impl InterfaceInner {
         ipv4_packet: &Ipv4Packet<&'a [u8]>,
         frag: &'a mut FragmentsBuffer,
     ) -> Option<Packet<'a>> {
-        let ipv4_repr = check!(Ipv4Repr::parse(ipv4_packet, &self.caps.checksum));
+        let mut ipv4_repr = check!(Ipv4Repr::parse(ipv4_packet, &self.caps.checksum));
         if !self.is_unicast_v4(ipv4_repr.src_addr) && !ipv4_repr.src_addr.is_unspecified() {
             // Discard packets with non-unicast source addresses but allow unspecified
             net_debug!("non-unicast or unspecified source address");
@@ -134,10 +134,10 @@ impl InterfaceInner {
                     return None;
                 }
 
-                // NOTE: according to the standard, the total length needs to be
-                // recomputed, as well as the checksum. However, we don't really use
-                // the IPv4 header after the packet is reassembled.
-                f.assemble()?
+                let payload = f.assemble()?;
+                // Update the payload length, so that the raw sockets get the correct value.
+                ipv4_repr.payload_len = payload.len();
+                payload
             } else {
                 ipv4_packet.payload()
             }

src/iface/interface/tests/ipv4.rs

@@ -1409,6 +1409,129 @@ fn test_raw_socket_tx_fragmentation(#[case] medium: Medium) {
     }
 }
 
+#[rstest]
+#[case(Medium::Ip)]
+#[cfg(all(
+    feature = "socket-raw",
+    feature = "proto-ipv4-fragmentation",
+    feature = "medium-ip"
+))]
+#[case(Medium::Ethernet)]
+#[cfg(all(
+    feature = "socket-raw",
+    feature = "proto-ipv4-fragmentation",
+    feature = "medium-ethernet"
+))]
+fn test_raw_socket_rx_fragmentation(#[case] medium: Medium) {
+    use crate::wire::{IpProtocol, IpVersion, Ipv4Address, Ipv4Packet, Ipv4Repr};
+
+    let (mut iface, mut sockets, _device) = setup(medium);
+
+    // Raw socket bound to IPv4 and a custom protocol.
+    let packets = 1;
+    let rx_buffer = raw::PacketBuffer::new(vec![raw::PacketMetadata::EMPTY; packets], vec![0; 64]);
+    let tx_buffer = raw::PacketBuffer::new(vec![raw::PacketMetadata::EMPTY; packets], vec![0; 64]);
+    let raw_socket = raw::Socket::new(
+        Some(IpVersion::Ipv4),
+        Some(IpProtocol::Unknown(99)),
+        rx_buffer,
+        tx_buffer,
+    );
+    let handle = sockets.add(raw_socket);
+
+    // Build two IPv4 fragments that together form one packet.
+    let src_addr = Ipv4Address::new(127, 0, 0, 2);
+    let dst_addr = Ipv4Address::new(127, 0, 0, 1);
+    let proto = IpProtocol::Unknown(99);
+    let ident: u16 = 0x1234;
+
+    let total_payload_len = 30usize;
+    let first_payload_len = 24usize; // must be a multiple of 8
+    let last_payload_len = total_payload_len - first_payload_len;
+
+    // Helper to build one fragment as on-the-wire bytes
+    let build_fragment = |payload_len: usize,
+                          more_frags: bool,
+                          frag_offset_octets: u16,
+                          payload_byte: u8|
+     -> Vec<u8> {
+        let repr = Ipv4Repr {
+            src_addr,
+            dst_addr,
+            next_header: proto,
+            hop_limit: 64,
+            payload_len,
+        };
+        let header_len = repr.buffer_len();
+        let mut bytes = vec![0u8; header_len + payload_len];
+        {
+            let mut pkt = Ipv4Packet::new_unchecked(&mut bytes[..]);
+            repr.emit(&mut pkt, &ChecksumCapabilities::default());
+            pkt.set_ident(ident);
+            pkt.set_dont_frag(false);
+            pkt.set_more_frags(more_frags);
+            pkt.set_frag_offset(frag_offset_octets);
+            // Recompute checksum after changing fragmentation fields.
+            pkt.fill_checksum();
+        }
+        // Fill payload with a simple pattern for validation
+        for b in &mut bytes[header_len..] {
+            *b = payload_byte;
+        }
+        bytes
+    };
+
+    let frag1_bytes = build_fragment(first_payload_len, true, 0, 0xAA);
+    let frag2_bytes = build_fragment(last_payload_len, false, first_payload_len as u16, 0xBB);
+
+    let frag1 = Ipv4Packet::new_unchecked(&frag1_bytes[..]);
+    let frag2 = Ipv4Packet::new_unchecked(&frag2_bytes[..]);
+
+    // First fragment alone should not be delivered to the raw socket.
+    assert_eq!(
+        iface.inner.process_ipv4(
+            &mut sockets,
+            PacketMeta::default(),
+            HardwareAddress::default(),
+            &frag1,
+            &mut iface.fragments
+        ),
+        None
+    );
+    {
+        let socket = sockets.get_mut::<raw::Socket>(handle);
+        assert!(!socket.can_recv());
+    }
+
+    // After the last fragment, the reassembled packet should be delivered.
+    assert_eq!(
+        iface.inner.process_ipv4(
+            &mut sockets,
+            PacketMeta::default(),
+            HardwareAddress::default(),
+            &frag2,
+            &mut iface.fragments
+        ),
+        None
+    );
+
+    // Validate the raw socket received one defragmented packet with correct payload.
+    let socket = sockets.get_mut::<raw::Socket>(handle);
+    assert!(socket.can_recv());
+    let data = socket.recv().expect("raw socket should have a packet");
+    let packet = Ipv4Packet::new_unchecked(data);
+    let repr = Ipv4Repr::parse(&packet, &ChecksumCapabilities::default()).unwrap();
+    assert_eq!(repr.src_addr, src_addr);
+    assert_eq!(repr.dst_addr, dst_addr);
+    assert_eq!(repr.next_header, proto);
+    assert_eq!(repr.payload_len, total_payload_len);
+
+    let payload = packet.payload();
+    assert_eq!(payload.len(), total_payload_len);
+    assert!(payload[..first_payload_len].iter().all(|&b| b == 0xAA));
+    assert!(payload[first_payload_len..].iter().all(|&b| b == 0xBB));
+}
+
 #[rstest]
 #[case(Medium::Ip)]
 #[cfg(all(feature = "socket-udp", feature = "medium-ip"))]