Update NAT example to illustrate "shared IPs" issue on session collisions

nat/Cargo.toml

@@ -31,6 +31,7 @@ tracing = { workspace = true }
 # internal
 lpm = { workspace = true, features = ["testing"] }
 net = { workspace = true, features = ["bolero"] }
+tracectl = { workspace = true }
 
 # external
 bolero = { workspace = true, default-features = false, features = ["alloc"] }

nat/src/stateful/allocator_writer.rs

@@ -65,13 +65,22 @@ impl NatAllocatorWriter {
         self.get_reader().factory()
     }
 
-    pub fn update_allocator(&mut self, vpc_table: &VpcTable) -> Result<(), ConfigError> {
+    fn update_allocator_and_set_randomness(
+        &mut self,
+        vpc_table: &VpcTable,
+        #[allow(unused_variables)] disable_randomness: bool,
+    ) -> Result<(), ConfigError> {
         let new_config = StatefulNatConfig::new(vpc_table);
 
         let old_allocator_guard = self.allocator.load();
         let Some(old_allocator) = old_allocator_guard.as_deref() else {
             // No existing allocator, build a new one
+            #[cfg(test)]
+            let new_allocator =
+                Self::build_new_allocator(&new_config)?.set_disable_randomness(disable_randomness);
+            #[cfg(not(test))]
             let new_allocator = Self::build_new_allocator(&new_config)?;
+
             self.allocator.store(Some(Arc::new(new_allocator)));
             self.config = new_config;
             return Ok(());
@@ -91,6 +100,18 @@ impl NatAllocatorWriter {
         Ok(())
     }
 
+    pub fn update_allocator(&mut self, vpc_table: &VpcTable) -> Result<(), ConfigError> {
+        self.update_allocator_and_set_randomness(vpc_table, false)
+    }
+
+    #[cfg(test)]
+    pub fn update_allocator_and_turn_off_randomness(
+        &mut self,
+        vpc_table: &VpcTable,
+    ) -> Result<(), ConfigError> {
+        self.update_allocator_and_set_randomness(vpc_table, true)
+    }
+
     fn build_new_allocator(config: &StatefulNatConfig) -> Result<NatDefaultAllocator, ConfigError> {
         NatDefaultAllocator::build_nat_allocator(config)
     }

nat/src/stateful/apalloc/alloc.rs

@@ -78,9 +78,12 @@ impl<I: NatIpWithBitmap> IpAllocator<I> {
         Err(AllocatorError::NoFreeIp)
     }
 
-    fn allocate_new_ip_from_pool(&self) -> Result<Arc<AllocatedIp<I>>, AllocatorError> {
+    fn allocate_new_ip_from_pool(
+        &self,
+        disable_randomness: bool,
+    ) -> Result<Arc<AllocatedIp<I>>, AllocatorError> {
         let mut allocated_ips = self.pool.write().unwrap();
-        let new_ip = allocated_ips.use_new_ip(self.clone())?;
+        let new_ip = allocated_ips.use_new_ip(self.clone(), disable_randomness)?;
         let arc_ip = Arc::new(new_ip);
         allocated_ips.add_in_use(&arc_ip);
         Ok(arc_ip)
@@ -89,8 +92,9 @@ impl<I: NatIpWithBitmap> IpAllocator<I> {
     fn allocate_from_new_ip(
         &self,
         allow_null: bool,
+        disable_randomness: bool,
     ) -> Result<port_alloc::AllocatedPort<I>, AllocatorError> {
-        self.allocate_new_ip_from_pool()
+        self.allocate_new_ip_from_pool(disable_randomness)
             .and_then(|ip| ip.allocate_port_for_ip(allow_null))
     }
 
@@ -102,6 +106,7 @@ impl<I: NatIpWithBitmap> IpAllocator<I> {
     pub(crate) fn allocate(
         &self,
         allow_null: bool,
+        disable_randomness: bool,
     ) -> Result<port_alloc::AllocatedPort<I>, AllocatorError> {
         // FIXME: Should we clean up every time??
         self.cleanup_used_ips();
@@ -110,22 +115,27 @@ impl<I: NatIpWithBitmap> IpAllocator<I> {
             return Ok(port);
         }
 
-        self.allocate_from_new_ip(allow_null)
+        self.allocate_from_new_ip(allow_null, disable_randomness)
     }
 
-    fn get_allocated_ip(&self, ip: I) -> Result<Arc<AllocatedIp<I>>, AllocatorError> {
+    fn get_allocated_ip(
+        &self,
+        ip: I,
+        disable_randomness: bool,
+    ) -> Result<Arc<AllocatedIp<I>>, AllocatorError> {
         self.pool
             .write()
             .unwrap()
-            .reserve_from_pool(ip, self.clone())
+            .reserve_from_pool(ip, self.clone(), disable_randomness)
     }
 
     pub(crate) fn reserve(
         &self,
         ip: I,
         port: NatPort,
+        disable_randomness: bool,
     ) -> Result<port_alloc::AllocatedPort<I>, AllocatorError> {
-        self.get_allocated_ip(ip)
+        self.get_allocated_ip(ip, disable_randomness)
             .and_then(|allocated_ip| allocated_ip.reserve_port_for_ip(port))
     }
 
@@ -153,14 +163,30 @@ pub(crate) struct AllocatedIp<I: NatIpWithBitmap> {
 }
 
 impl<I: NatIpWithBitmap> AllocatedIp<I> {
-    fn new(ip: I, ip_allocator: IpAllocator<I>) -> Self {
+    fn new(ip: I, ip_allocator: IpAllocator<I>, disable_randomness: bool) -> Self {
+        // Allow opt-in deterministic port allocation for tests
+        let port_allocator = Self::get_new_portallocator(disable_randomness);
+
         Self {
             ip,
-            port_allocator: port_alloc::PortAllocator::new(),
+            port_allocator,
             ip_allocator,
         }
     }
 
+    #[cfg(test)]
+    fn get_new_portallocator(disable_randomness: bool) -> port_alloc::PortAllocator<I> {
+        if disable_randomness {
+            port_alloc::PortAllocator::new_no_randomness()
+        } else {
+            port_alloc::PortAllocator::new()
+        }
+    }
+    #[cfg(not(test))]
+    fn get_new_portallocator(_disable_randomness: bool) -> port_alloc::PortAllocator<I> {
+        port_alloc::PortAllocator::new()
+    }
+
     pub(crate) fn ip(&self) -> I {
         self.ip
     }
@@ -245,12 +271,13 @@ impl<I: NatIpWithBitmap> NatPool<I> {
     fn use_new_ip(
         &mut self,
         ip_allocator: IpAllocator<I>,
+        disable_randomness: bool,
     ) -> Result<AllocatedIp<I>, AllocatorError> {
         // Retrieve the first available offset
         let offset = self.bitmap.pop_ip()?;
 
         let ip = I::try_from_offset(offset, &self.bitmap_mapping)?;
-        Ok(AllocatedIp::new(ip, ip_allocator))
+        Ok(AllocatedIp::new(ip, ip_allocator, disable_randomness))
     }
 
     fn deallocate_from_pool(&mut self, ip: I) {
@@ -262,6 +289,7 @@ impl<I: NatIpWithBitmap> NatPool<I> {
         &mut self,
         ip: I,
         ip_allocator: IpAllocator<I>,
+        disable_randomness: bool,
     ) -> Result<Arc<AllocatedIp<I>>, AllocatorError> {
         let offset = I::try_to_offset(ip, &self.reverse_bitmap_mapping)?;
 
@@ -282,7 +310,7 @@ impl<I: NatIpWithBitmap> NatPool<I> {
         // drops an AllocatedIp and its reference count goes to 0, but it hasn't called the drop()
         // function to remove the IP from the bitmap in that other thread yet).
         let _ = self.bitmap.set_ip_allocated(offset);
-        let arc_ip = Arc::new(AllocatedIp::new(ip, ip_allocator));
+        let arc_ip = Arc::new(AllocatedIp::new(ip, ip_allocator, disable_randomness));
         self.add_in_use(&arc_ip);
         Ok(arc_ip)
     }

nat/src/stateful/apalloc/mod.rs

@@ -242,6 +242,8 @@ pub struct NatDefaultAllocator {
     pools_dst66: PoolTable<Ipv6Addr, Ipv6Addr>,
     exempt4: ExemptTable<Ipv4Addr>,
     exempt6: ExemptTable<Ipv6Addr>,
+    #[cfg(test)]
+    disable_randomness: bool,
 }
 
 impl NatAllocator<AllocatedIpPort<Ipv4Addr>, AllocatedIpPort<Ipv6Addr>> for NatDefaultAllocator {
@@ -253,21 +255,33 @@ impl NatAllocator<AllocatedIpPort<Ipv4Addr>, AllocatedIpPort<Ipv6Addr>> for NatD
             pools_dst66: PoolTable::new(),
             exempt4: ExemptTable::new(),
             exempt6: ExemptTable::new(),
+            #[cfg(test)]
+            disable_randomness: false,
         }
     }
 
     fn allocate_v4(
         &self,
         flow_key: &FlowKey,
     ) -> Result<AllocationResult<AllocatedIpPort<Ipv4Addr>>, AllocatorError> {
-        Self::allocate_from_tables(flow_key, &self.pools_src44, &self.pools_dst44)
+        Self::allocate_from_tables(
+            flow_key,
+            &self.pools_src44,
+            &self.pools_dst44,
+            self.must_disable_randomness(),
+        )
     }
 
     fn allocate_v6(
         &self,
         flow_key: &FlowKey,
     ) -> Result<AllocationResult<AllocatedIpPort<Ipv6Addr>>, AllocatorError> {
-        Self::allocate_from_tables(flow_key, &self.pools_src66, &self.pools_dst66)
+        Self::allocate_from_tables(
+            flow_key,
+            &self.pools_src66,
+            &self.pools_dst66,
+            self.must_disable_randomness(),
+        )
     }
 
     fn is_exempt_v4(&self, flow_key: &FlowKey) -> Result<bool, AllocatorError> {
@@ -302,6 +316,7 @@ impl NatDefaultAllocator {
         flow_key: &FlowKey,
         pools_src: &PoolTable<I, I>,
         pools_dst: &PoolTable<I, I>,
+        disable_randomness: bool,
     ) -> Result<AllocationResult<AllocatedIpPort<I>>, AllocatorError> {
         let next_header = Self::get_next_header(flow_key);
         Self::check_proto(next_header)?;
@@ -341,7 +356,8 @@ impl NatDefaultAllocator {
 
         // Allocate IP and ports from pools, for source and destination NAT
         let allow_null = matches!(flow_key.data().proto_key_info(), IpProtoKey::Icmp(_));
-        let (src_mapping, dst_mapping) = Self::get_mapping(pool_src_opt, pool_dst_opt, allow_null)?;
+        let (src_mapping, dst_mapping) =
+            Self::get_mapping(pool_src_opt, pool_dst_opt, allow_null, disable_randomness)?;
 
         // Now based on the previous allocation, we need to "reserve" IP and ports for the reverse
         // path for the flow. First retrieve the relevant address pools.
@@ -359,8 +375,12 @@ impl NatDefaultAllocator {
         };
 
         // Reserve IP and ports for the reverse path for the flow.
-        let (reverse_src_mapping, reverse_dst_mapping) =
-            Self::get_reverse_mapping(flow_key, reverse_pool_src_opt, reverse_pool_dst_opt)?;
+        let (reverse_src_mapping, reverse_dst_mapping) = Self::get_reverse_mapping(
+            flow_key,
+            reverse_pool_src_opt,
+            reverse_pool_dst_opt,
+            disable_randomness,
+        )?;
 
         Ok(AllocationResult {
             src: src_mapping,
@@ -372,6 +392,16 @@ impl NatDefaultAllocator {
         })
     }
 
+    #[cfg(test)]
+    const fn must_disable_randomness(&self) -> bool {
+        self.disable_randomness
+    }
+    #[cfg(not(test))]
+    #[allow(clippy::unused_self)]
+    const fn must_disable_randomness(&self) -> bool {
+        false
+    }
+
     fn check_proto(next_header: NextHeader) -> Result<(), AllocatorError> {
         match next_header {
             NextHeader::TCP | NextHeader::UDP | NextHeader::ICMP | NextHeader::ICMP6 => Ok(()),
@@ -411,6 +441,7 @@ impl NatDefaultAllocator {
         pool_src_opt: Option<&alloc::IpAllocator<I>>,
         pool_dst_opt: Option<&alloc::IpAllocator<I>>,
         allow_null: bool,
+        disable_randomness: bool,
     ) -> Result<AllocationMapping<I>, AllocatorError> {
         // Allocate IP and ports for source and destination NAT.
         //
@@ -424,12 +455,12 @@ impl NatDefaultAllocator {
         // port/identifier value for the src_mapping, even though we'll never use it. (This does not
         // apply to TCP or UDP, for which we need and use both ports).
         let src_mapping = match pool_src_opt {
-            Some(pool_src) => Some(pool_src.allocate(allow_null)?),
+            Some(pool_src) => Some(pool_src.allocate(allow_null, disable_randomness)?),
             None => None,
         };
 
         let dst_mapping = match pool_dst_opt {
-            Some(pool_dst) => Some(pool_dst.allocate(allow_null)?),
+            Some(pool_dst) => Some(pool_dst.allocate(allow_null, disable_randomness)?),
             None => None,
         };
 
@@ -440,6 +471,7 @@ impl NatDefaultAllocator {
         flow_key: &FlowKey,
         reverse_pool_src_opt: Option<&alloc::IpAllocator<I>>,
         reverse_pool_dst_opt: Option<&alloc::IpAllocator<I>>,
+        disable_randomness: bool,
     ) -> Result<AllocationMapping<I>, AllocatorError> {
         let reverse_src_mapping = match reverse_pool_src_opt {
             Some(pool_src) => {
@@ -462,6 +494,7 @@ impl NatDefaultAllocator {
                         )
                     })?,
                     reservation_src_port_number,
+                    disable_randomness,
                 )?)
             }
             None => None,
@@ -482,6 +515,7 @@ impl NatDefaultAllocator {
                         )
                     })?,
                     reservation_dst_port_number,
+                    disable_randomness,
                 )?)
             }
             None => None,
@@ -500,6 +534,13 @@ impl NatDefaultAllocator {
             IcmpProtoKey::Unsupported => Err(AllocatorError::UnsupportedIcmpCategory),
         }
     }
+
+    #[cfg(test)]
+    #[must_use]
+    pub fn set_disable_randomness(mut self, disable_randomness: bool) -> Self {
+        self.disable_randomness = disable_randomness;
+        self
+    }
 }
 
 // This method is for setting a range end field that is not usually relevant for table lookups, for

nat/src/stateful/apalloc/port_alloc.rs

@@ -95,6 +95,20 @@ impl<I: NatIpWithBitmap> PortAllocator<I> {
         }
     }
 
+    #[cfg(test)]
+    pub(crate) fn new_no_randomness() -> Self {
+        let base_ports = (0..=255).collect::<Vec<_>>();
+        // Do not shuffle
+        let blocks = std::array::from_fn(|i| AllocatorPortBlock::new(base_ports[i]));
+        Self {
+            blocks,
+            usable_blocks: AtomicU16::new(256),
+            current_alloc_index: AtomicUsize::new(0),
+            thread_blocks: ThreadPortMap::new(),
+            allocated_blocks: AllocatedPortBlockMap::new(),
+        }
+    }
+
     #[concurrency_mode(std)]
     fn shuffle_slice<T>(slice: &mut [T]) {
         let mut rng = rand::rng();