From d2e6ea1ba9668b85b43421dd6a2bfad361ef5fe1 Mon Sep 17 00:00:00 2001
From: Red Bear OS <build@redbearos.org>
Date: Wed, 8 Jul 2026 21:16:50 +0300
Subject: [PATCH] nat: add 6 unit tests for IP rewrite + table management

Tests cover:
- rewrite_src_ipv4: source IP changes, dst unchanged
- rewrite_dst_ipv4: destination IP changes, src unchanged
- rewrite_short_packet_returns_false: <20 bytes returns false
- rewrite_recomputes_checksum: checksum updates after rewrite
- nat_table_add_assigns_unique_ids: each rule gets unique ID
- nat_ephemeral_port_allocates_unique: ports are unique across calls

Total tests across filter module: 15 (5 table + 4 conntrack + 6 nat)
All passing.
---
 netstack/src/filter/nat.rs | 95 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 94 insertions(+), 1 deletion(-)

diff --git a/netstack/src/filter/nat.rs b/netstack/src/filter/nat.rs
index 2e9e9612d1..b0e5e0a8ed 100644
--- a/netstack/src/filter/nat.rs
+++ b/netstack/src/filter/nat.rs
@@ -358,4 +358,97 @@ pub enum NatParseError {
     BadHook(String),
     BadAddress(String),
     BadPort(String),
-}
\ No newline at end of file
+}
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn make_ipv4_udp_packet(src: [u8; 4], dst: [u8; 4]) -> Vec<u8> {
+        let mut p = vec![0u8; 28];
+        p[0] = 0x45; p[1] = 0x00;
+        p[2] = 0x00; p[3] = 0x1c;
+        p[6] = 0x00; p[7] = 0x00;
+        p[8] = 64;
+        p[9] = 17;
+        p[10] = 0x00; p[11] = 0x00;
+        p[12..16].copy_from_slice(&src);
+        p[16..20].copy_from_slice(&dst);
+        p[20] = 0x12; p[21] = 0x34;
+        p[22] = 0x56; p[23] = 0x78;
+        p[24] = 0x00; p[25] = 0x08;
+        p[26] = 0x00; p[27] = 0x00;
+        p
+    }
+
+    fn read_ipv4(p: &[u8], offset: usize) -> Ipv4Address {
+        Ipv4Address::new(p[offset], p[offset+1], p[offset+2], p[offset+3])
+    }
+
+    #[test]
+    fn rewrite_src_changes_source_address() {
+        let mut p = make_ipv4_udp_packet([10, 0, 0, 1], [192, 168, 1, 1]);
+        let new_src = Ipv4Address::new(192, 168, 99, 99);
+        assert!(rewrite_src_ipv4(&mut p, new_src));
+        assert_eq!(read_ipv4(&p, 12), new_src, "Source IP must be rewritten");
+        assert_eq!(read_ipv4(&p, 16), Ipv4Address::new(192, 168, 1, 1),
+            "Destination IP must be unchanged");
+    }
+
+    #[test]
+    fn rewrite_dst_changes_destination_address() {
+        let mut p = make_ipv4_udp_packet([10, 0, 0, 1], [192, 168, 1, 1]);
+        let new_dst = Ipv4Address::new(8, 8, 8, 8);
+        assert!(rewrite_dst_ipv4(&mut p, new_dst));
+        assert_eq!(read_ipv4(&p, 16), new_dst);
+        assert_eq!(read_ipv4(&p, 12), Ipv4Address::new(10, 0, 0, 1),
+            "Source IP must be unchanged");
+    }
+
+    #[test]
+    fn rewrite_short_packet_returns_false() {
+        let mut p = vec![0u8; 10];
+        assert!(!rewrite_src_ipv4(&mut p, Ipv4Address::new(1, 2, 3, 4)));
+        assert!(!rewrite_dst_ipv4(&mut p, Ipv4Address::new(1, 2, 3, 4)));
+    }
+
+    #[test]
+    fn rewrite_recomputes_checksum_when_previously_nonzero() {
+        let mut p = make_ipv4_udp_packet([10, 0, 0, 1], [192, 168, 1, 1]);
+        // Compute a real checksum first
+        if let Ok(mut ipv4) = smoltcp::wire::Ipv4Packet::new_checked(&mut p) {
+            ipv4.fill_checksum();
+        }
+        let original_csum = u16::from_be_bytes([p[10], p[11]]);
+        assert_ne!(original_csum, 0, "Pre-condition: checksum should be non-zero");
+        let _ = rewrite_src_ipv4(&mut p, Ipv4Address::new(10, 0, 0, 2));
+        let new_csum = u16::from_be_bytes([p[10], p[11]]);
+        assert_ne!(new_csum, original_csum,
+            "Checksum must be updated after source rewrite");
+    }
+
+    #[test]
+    fn nat_table_add_assigns_unique_ids() {
+        let mut t = NatTable::new();
+        let make_rule = || NatRule {
+            id: 0,
+            nat_type: NatType::Snat,
+            hook: Hook::OutputLocal,
+            src_match: None,
+            dst_match: None,
+            trans_addr: IpAddress::v4(192, 168, 1, 100),
+            trans_port: None,
+            match_count: 0,
+        };
+        let id1 = t.add(make_rule());
+        let id2 = t.add(make_rule());
+        assert_ne!(id1, id2, "Each rule gets a unique ID");
+    }
+
+    #[test]
+    fn nat_ephemeral_port_allocates_unique() {
+        let mut t = NatTable::new();
+        let p1 = t.alloc_ephemeral_port();
+        let p2 = t.alloc_ephemeral_port();
+        assert_ne!(p1, p2, "Ephemeral ports must be unique across calls");
+    }
+}