Add MutablePacket trait and protocol implementation

examples/mutable_chaining.rs

@@ -0,0 +1,93 @@
+//! Demonstrates chaining mutable packet views across Ethernet/IPv4/UDP layers.
+
+use nex::net::mac::MacAddr;
+use nex::packet::ethernet::{
+    EtherType, EthernetPacket, MutableEthernetPacket, ETHERNET_HEADER_LEN,
+};
+use nex::packet::ip::IpNextProtocol;
+use nex::packet::ipv4::{self, Ipv4Packet, MutableIpv4Packet, IPV4_HEADER_LEN};
+use nex::packet::packet::{MutablePacket, Packet};
+use nex::packet::udp::{self, MutableUdpPacket, UdpPacket, UDP_HEADER_LEN};
+use std::net::Ipv4Addr;
+
+fn main() {
+    // Build a simple Ethernet/IPv4/UDP frame in-place.
+    let payload = b"hello mutable packets";
+    let frame_len = ETHERNET_HEADER_LEN + IPV4_HEADER_LEN + UDP_HEADER_LEN + payload.len();
+    let mut frame = vec![0u8; frame_len];
+
+    {
+        let mut ethernet = MutableEthernetPacket::new(&mut frame).expect("ethernet");
+        ethernet.set_source(MacAddr::new(0x00, 0x11, 0x22, 0x33, 0x44, 0x55));
+        ethernet.set_destination(MacAddr::new(0x08, 0x00, 0x27, 0xaa, 0xbb, 0xcc));
+        ethernet.set_ethertype(EtherType::Ipv4);
+
+        let ipv4_len = (IPV4_HEADER_LEN + UDP_HEADER_LEN + payload.len()) as u16;
+        {
+            // Use `new_unchecked` because the buffer starts zeroed; we will
+            // populate all required header fields before freezing it back into
+            // an immutable packet for validation.
+            let mut ipv4 = MutableIpv4Packet::new_unchecked(ethernet.payload_mut());
+            ipv4.set_version(4);
+            ipv4.set_header_length(5);
+            ipv4.set_total_length(ipv4_len);
+            ipv4.set_ttl(64);
+            ipv4.set_next_level_protocol(IpNextProtocol::Udp);
+            ipv4.set_source(Ipv4Addr::new(192, 0, 2, 1));
+            ipv4.set_destination(Ipv4Addr::new(198, 51, 100, 1));
+            ipv4.set_identification(0x1337);
+            ipv4.set_checksum(0);
+
+            {
+                let mut udp = MutableUdpPacket::new(ipv4.payload_mut()).expect("udp");
+                udp.set_source(5353);
+                udp.set_destination(8080);
+                udp.set_length((UDP_HEADER_LEN + payload.len()) as u16);
+                udp.set_checksum(0);
+
+                let udp_payload = udp.payload_mut();
+                udp_payload[..payload.len()].copy_from_slice(payload);
+            }
+
+            let snapshot = ipv4.freeze().expect("snapshot ipv4");
+            let udp_snapshot = UdpPacket::from_buf(&snapshot.payload).expect("snapshot udp");
+            let udp_checksum = udp::ipv4_checksum(
+                &udp_snapshot,
+                &snapshot.header.source,
+                &snapshot.header.destination,
+            );
+            MutableUdpPacket::new(ipv4.payload_mut())
+                .expect("udp checksum")
+                .set_checksum(udp_checksum);
+            let ipv4_checksum = ipv4::checksum(&snapshot);
+            ipv4.set_checksum(ipv4_checksum);
+        }
+    }
+
+    // Inspect immutable packet views to confirm changes persisted across layers.
+    let ethernet_packet = EthernetPacket::from_buf(&frame).expect("immutable ethernet");
+    let ipv4_packet = Ipv4Packet::from_buf(&ethernet_packet.payload).expect("immutable ipv4");
+    let udp_packet = UdpPacket::from_buf(&ipv4_packet.payload).expect("immutable udp");
+
+    println!(
+        "Ethernet: {} -> {} ({:?})",
+        ethernet_packet.header.source,
+        ethernet_packet.header.destination,
+        ethernet_packet.header.ethertype
+    );
+    println!(
+        "IPv4: {} -> {} ttl={} checksum=0x{:04x}",
+        ipv4_packet.header.source,
+        ipv4_packet.header.destination,
+        ipv4_packet.header.ttl,
+        ipv4_packet.header.checksum
+    );
+    println!(
+        "UDP: {} -> {} len={} checksum=0x{:04x}",
+        udp_packet.header.source,
+        udp_packet.header.destination,
+        udp_packet.header.length,
+        udp_packet.header.checksum
+    );
+    println!("Payload: {}", String::from_utf8_lossy(&udp_packet.payload));
+}

nex-packet/src/arp.rs

@@ -2,7 +2,7 @@
 
 use crate::{
     ethernet::{EtherType, ETHERNET_HEADER_LEN},
-    packet::Packet,
+    packet::{MutablePacket, Packet},
 };
 
 use bytes::{Bytes, BytesMut};
@@ -448,6 +448,146 @@ impl fmt::Display for ArpPacket {
     }
 }
 
+/// Represents a mutable ARP Packet.
+pub struct MutableArpPacket<'a> {
+    buffer: &'a mut [u8],
+}
+
+impl<'a> MutablePacket<'a> for MutableArpPacket<'a> {
+    type Packet = ArpPacket;
+
+    fn new(buffer: &'a mut [u8]) -> Option<Self> {
+        if buffer.len() < ARP_HEADER_LEN {
+            None
+        } else {
+            Some(Self { buffer })
+        }
+    }
+
+    fn packet(&self) -> &[u8] {
+        &*self.buffer
+    }
+
+    fn packet_mut(&mut self) -> &mut [u8] {
+        &mut *self.buffer
+    }
+
+    fn header(&self) -> &[u8] {
+        &self.packet()[..ARP_HEADER_LEN]
+    }
+
+    fn header_mut(&mut self) -> &mut [u8] {
+        let (header, _) = (&mut *self.buffer).split_at_mut(ARP_HEADER_LEN);
+        header
+    }
+
+    fn payload(&self) -> &[u8] {
+        &self.packet()[ARP_HEADER_LEN..]
+    }
+
+    fn payload_mut(&mut self) -> &mut [u8] {
+        let (_, payload) = (&mut *self.buffer).split_at_mut(ARP_HEADER_LEN);
+        payload
+    }
+}
+
+impl<'a> MutableArpPacket<'a> {
+    /// Create a packet without performing length checks.
+    pub fn new_unchecked(buffer: &'a mut [u8]) -> Self {
+        Self { buffer }
+    }
+
+    fn raw(&self) -> &[u8] {
+        &*self.buffer
+    }
+
+    fn raw_mut(&mut self) -> &mut [u8] {
+        &mut *self.buffer
+    }
+
+    pub fn get_hardware_type(&self) -> ArpHardwareType {
+        ArpHardwareType::new(u16::from_be_bytes([self.raw()[0], self.raw()[1]]))
+    }
+
+    pub fn set_hardware_type(&mut self, ty: ArpHardwareType) {
+        self.raw_mut()[0..2].copy_from_slice(&ty.value().to_be_bytes());
+    }
+
+    pub fn get_protocol_type(&self) -> EtherType {
+        EtherType::new(u16::from_be_bytes([self.raw()[2], self.raw()[3]]))
+    }
+
+    pub fn set_protocol_type(&mut self, ty: EtherType) {
+        self.raw_mut()[2..4].copy_from_slice(&ty.value().to_be_bytes());
+    }
+
+    pub fn get_hw_addr_len(&self) -> u8 {
+        self.raw()[4]
+    }
+
+    pub fn set_hw_addr_len(&mut self, len: u8) {
+        self.raw_mut()[4] = len;
+    }
+
+    pub fn get_proto_addr_len(&self) -> u8 {
+        self.raw()[5]
+    }
+
+    pub fn set_proto_addr_len(&mut self, len: u8) {
+        self.raw_mut()[5] = len;
+    }
+
+    pub fn get_operation(&self) -> ArpOperation {
+        ArpOperation::new(u16::from_be_bytes([self.raw()[6], self.raw()[7]]))
+    }
+
+    pub fn set_operation(&mut self, op: ArpOperation) {
+        self.raw_mut()[6..8].copy_from_slice(&op.value().to_be_bytes());
+    }
+
+    pub fn get_sender_hw_addr(&self) -> MacAddr {
+        MacAddr::from_octets(self.raw()[8..14].try_into().unwrap())
+    }
+
+    pub fn set_sender_hw_addr(&mut self, addr: MacAddr) {
+        self.raw_mut()[8..14].copy_from_slice(&addr.octets());
+    }
+
+    pub fn get_sender_proto_addr(&self) -> Ipv4Addr {
+        Ipv4Addr::new(
+            self.raw()[14],
+            self.raw()[15],
+            self.raw()[16],
+            self.raw()[17],
+        )
+    }
+
+    pub fn set_sender_proto_addr(&mut self, addr: Ipv4Addr) {
+        self.raw_mut()[14..18].copy_from_slice(&addr.octets());
+    }
+
+    pub fn get_target_hw_addr(&self) -> MacAddr {
+        MacAddr::from_octets(self.raw()[18..24].try_into().unwrap())
+    }
+
+    pub fn set_target_hw_addr(&mut self, addr: MacAddr) {
+        self.raw_mut()[18..24].copy_from_slice(&addr.octets());
+    }
+
+    pub fn get_target_proto_addr(&self) -> Ipv4Addr {
+        Ipv4Addr::new(
+            self.raw()[24],
+            self.raw()[25],
+            self.raw()[26],
+            self.raw()[27],
+        )
+    }
+
+    pub fn set_target_proto_addr(&mut self, addr: Ipv4Addr) {
+        self.raw_mut()[24..28].copy_from_slice(&addr.octets());
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -543,4 +683,31 @@ mod tests {
             _ => panic!("Expected unknown operation"),
         }
     }
+
+    #[test]
+    fn test_mutable_arp_packet_updates() {
+        let mut raw = [
+            0x00, 0x01, // Hardware Type: Ethernet
+            0x08, 0x00, // Protocol Type: IPv4
+            0x06, // HW Addr Len
+            0x04, // Proto Addr Len
+            0x00, 0x01, // Operation: Request
+            0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff, // Sender MAC
+            192, 168, 1, 1, // Sender IP
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Target MAC
+            192, 168, 1, 2, // Target IP
+            0xde, 0xad, 0xbe, 0xef, // payload
+        ];
+
+        let mut packet = MutableArpPacket::new(&mut raw).expect("mutable arp");
+        assert_eq!(packet.get_operation(), ArpOperation::Request);
+        packet.set_operation(ArpOperation::Reply);
+        packet.set_sender_proto_addr(Ipv4Addr::new(10, 0, 0, 1));
+        packet.payload_mut()[0] = 0xaa;
+
+        let frozen = packet.freeze().expect("freeze");
+        assert_eq!(frozen.header.operation, ArpOperation::Reply);
+        assert_eq!(frozen.header.sender_proto_addr, Ipv4Addr::new(10, 0, 0, 1));
+        assert_eq!(frozen.payload[0], 0xaa);
+    }
 }

nex-packet/src/dhcp.rs

@@ -2,7 +2,7 @@ use bytes::{Buf, BufMut, Bytes, BytesMut};
 use nex_core::mac::MacAddr;
 use std::net::Ipv4Addr;
 
-use crate::packet::Packet;
+use crate::packet::{GenericMutablePacket, Packet};
 
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
@@ -406,9 +406,13 @@ impl Packet for DhcpPacket {
     }
 }
 
+/// Represents a mutable DHCP packet.
+pub type MutableDhcpPacket<'a> = GenericMutablePacket<'a, DhcpPacket>;
+
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::packet::MutablePacket;
     use nex_core::mac::MacAddr;
 
     #[test]
@@ -447,4 +451,20 @@ mod tests {
         let rebuilt = packet.to_bytes();
         assert_eq!(rebuilt, raw);
     }
+
+    #[test]
+    fn test_mutable_dhcp_packet_alias() {
+        let mut raw = [0u8; DHCP_MIN_PACKET_SIZE + 4];
+        raw[0] = DhcpOperation::Request.value();
+        raw[1] = DhcpHardwareType::Ethernet.value();
+        raw[2] = 6; // hardware length
+
+        let mut packet = <MutableDhcpPacket as MutablePacket>::new(&mut raw).expect("mutable dhcp");
+        packet.header_mut()[0] = DhcpOperation::Reply.value();
+        packet.payload_mut()[0] = 0xaa;
+
+        let frozen = packet.freeze().expect("freeze");
+        assert_eq!(frozen.header.op, DhcpOperation::Reply);
+        assert_eq!(frozen.payload[0], 0xaa);
+    }
 }

nex-packet/src/dns.rs

@@ -1,4 +1,4 @@
-use crate::packet::Packet;
+use crate::packet::{GenericMutablePacket, Packet};
 use bytes::{BufMut, Bytes, BytesMut};
 use core::str;
 use nex_core::bitfield::{u1, u16be, u32be};
@@ -1208,9 +1208,13 @@ impl std::fmt::Display for DnsName {
     }
 }
 
+/// Represents a mutable DNS packet.
+pub type MutableDnsPacket<'a> = GenericMutablePacket<'a, DnsPacket>;
+
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::packet::MutablePacket;
 
     #[test]
     fn test_dns_query() {
@@ -1294,4 +1298,19 @@ mod tests {
         assert_eq!(packet.responses[0].data_len, 4);
         assert_eq!(packet.responses[0].data, vec![192, 168, 122, 189]);
     }
+
+    #[test]
+    fn test_mutable_dns_packet_header_edit() {
+        let mut raw = [0u8; 16];
+        raw[1] = 0x01; // id
+
+        let mut packet = <MutableDnsPacket as MutablePacket>::new(&mut raw).expect("mutable dns");
+        packet.header_mut()[0] = 0x12;
+        packet.header_mut()[1] = 0x34;
+        packet.payload_mut()[0] = 0xaa;
+
+        let frozen = packet.freeze().expect("freeze");
+        assert_eq!(frozen.header.id, 0x1234);
+        assert_eq!(frozen.payload[0], 0xaa);
+    }
 }