fix(masquerade): fix flow genid for masqueraded flows

acl-filter/src/tests.rs

@@ -948,7 +948,7 @@ mod end_to_end {
 
         // Masquerade (creates the related flow pair used by 'flow'-scoped replies)
         let mut allocator = NatAllocatorWriter::new();
-        allocator.update_nat_allocator(MasqueradeConfig::new(overlay.vpc_table(), 1), &flow_table);
+        allocator.update_nat_allocator(MasqueradeConfig::new(overlay.vpc_table()), 1, &flow_table);
         pipeline = pipeline.add_stage(Masquerade::new(
             "masquerade",
             flow_table.clone(),

mgmt/src/processor/proc.rs

@@ -538,8 +538,8 @@ fn apply_masquerade_config(
     natallocatorw: &mut NatAllocatorWriter,
     genid: GenId,
 ) {
-    let nat_config = MasqueradeConfig::new(vpc_table, genid).set_randomize(true);
-    natallocatorw.update_nat_allocator(nat_config, flow_table);
+    let nat_config = MasqueradeConfig::new(vpc_table).set_randomize(true);
+    natallocatorw.update_nat_allocator(nat_config, genid, flow_table);
     debug!("Updated masquerade NAT allocator");
 }
 

nat/src/masquerade/allocator_writer.rs

@@ -21,22 +21,15 @@ pub(crate) struct MasqueradePeering {
     pub(crate) dst_vpcd: VpcDiscriminant,
     pub(crate) peering: ValidatedPeering,
 }
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Default, Clone, PartialEq)]
 pub struct MasqueradeConfig {
-    genid: GenId,
     peerings: Vec<MasqueradePeering>,
     randomize: bool,
 }
-impl PartialEq for MasqueradeConfig {
-    fn eq(&self, other: &Self) -> bool {
-        // we exclude genid from comparison
-        self.peerings == other.peerings && self.randomize == other.randomize
-    }
-}
 
 impl MasqueradeConfig {
     #[must_use]
-    pub fn new(vpc_table: &ValidatedVpcTable, genid: GenId) -> Self {
+    pub fn new(vpc_table: &ValidatedVpcTable) -> Self {
         let mut peerings = Vec::new();
         for vpc in vpc_table.values() {
             for peering in vpc.local_stateful_nat_peerings() {
@@ -48,17 +41,11 @@ impl MasqueradeConfig {
             }
         }
         Self {
-            genid,
             peerings,
             randomize: true, // randomize by default
         }
     }
 
-    #[must_use]
-    pub fn genid(&self) -> GenId {
-        self.genid
-    }
-
     #[must_use]
     pub fn set_randomize(mut self, value: bool) -> Self {
         self.randomize = value;
@@ -119,15 +106,20 @@ impl NatAllocatorWriter {
     /// masquerading config is provided, a new allocator will be installed and the flows using the
     /// previous one be either invalidated or adapted to use the new allocator: their ports/ips
     /// will be transferred (reserved) in the new allocator.
-    pub fn update_nat_allocator(&mut self, nat_config: MasqueradeConfig, flow_table: &FlowTable) {
-        let genid = nat_config.genid();
+    pub fn update_nat_allocator(
+        &mut self,
+        nat_config: MasqueradeConfig,
+        genid: GenId,
+        flow_table: &FlowTable,
+    ) {
         let curr_allocator = self.0.load_full();
 
         // keep state as-is if config did not change, and just upgrade flows
         if let Some(current) = curr_allocator.as_ref()
             && current.config() == &nat_config
         {
             debug!("No need to update NAT allocator: NAT peerings did not change");
+            current.set_genid(genid);
             upgrade_all_masquerading_flows(flow_table, genid);
             return;
         }
@@ -142,7 +134,7 @@ impl NatAllocatorWriter {
             return;
         }
 
-        let mut allocator = NatAllocator::new(nat_config);
+        let mut allocator = NatAllocator::new(nat_config, genid);
         if curr_allocator.is_some() {
             let guard = check_masquerading_flows(flow_table, &mut allocator);
             debug!("Replacing masquerade NAT allocator...");

nat/src/masquerade/apalloc/mod.rs

@@ -68,6 +68,7 @@ use crate::NatPort;
 use crate::masquerade::MasqueradeConfig;
 pub use crate::masquerade::apalloc::natip_with_bitmap::NatIpWithBitmap;
 use crate::masquerade::natip::NatIp;
+use concurrency::sync::atomic::{AtomicI64, Ordering};
 use config::GenId;
 use net::ip::NextHeader;
 use net::packet::VpcDiscriminant;
@@ -183,21 +184,6 @@ impl Allocation {
             Self::V6(a) => a.port(),
         }
     }
-
-    #[must_use]
-    pub fn genid(&self) -> GenId {
-        match self {
-            Self::V4(a) => a.genid(),
-            Self::V6(a) => a.genid(),
-        }
-    }
-
-    pub fn set_genid(&mut self, genid: GenId) {
-        match self {
-            Self::V4(a) => a.set_genid(genid),
-            Self::V6(a) => a.set_genid(genid),
-        }
-    }
 }
 
 impl Display for Allocation {
@@ -217,17 +203,19 @@ impl Display for Allocation {
 #[derive(Debug)]
 pub struct NatAllocator {
     config: MasqueradeConfig,
+    genid: AtomicI64,
     pools_src44: PoolTable<Ipv4Addr, Ipv4Addr>,
     pools_src66: PoolTable<Ipv6Addr, Ipv6Addr>,
     randomize: bool,
 }
 
 impl NatAllocator {
     #[must_use]
-    pub(crate) fn new(config: MasqueradeConfig) -> Self {
-        debug!("Building NAT allocator for genid {}", config.genid());
+    pub(crate) fn new(config: MasqueradeConfig, genid: GenId) -> Self {
+        debug!("Building NAT allocator for genid {genid}");
         let mut allocator = Self {
             config: MasqueradeConfig::default(),
+            genid: AtomicI64::new(genid),
             pools_src44: PoolTable::new(),
             pools_src66: PoolTable::new(),
             randomize: config.randomize(),
@@ -243,13 +231,24 @@ impl NatAllocator {
         &self.config
     }
 
+    /// The configuration generation this allocator currently serves.
+    pub(crate) fn genid(&self) -> GenId {
+        self.genid.load(Ordering::Relaxed)
+    }
+
+    /// Advance the genid served, for a config update that left the NAT peerings untouched and
+    /// therefore kept this allocator
+    pub(crate) fn set_genid(&self, genid: GenId) {
+        self.genid.store(genid, Ordering::Relaxed);
+    }
+
     fn allocate_v4(
         &self,
         dst_vpcd: VpcDiscriminant,
         src_ip: Ipv4Addr,
         next_header: NextHeader,
     ) -> Result<AllocationResult<AllocatedPort<Ipv4Addr>>, AllocatorError> {
-        self.allocate_from_tables(src_ip.into(), dst_vpcd, next_header, &self.pools_src44)
+        Self::allocate_from_tables(src_ip.into(), dst_vpcd, next_header, &self.pools_src44)
     }
 
     fn allocate_v6(
@@ -258,7 +257,7 @@ impl NatAllocator {
         src_ip: Ipv6Addr,
         next_header: NextHeader,
     ) -> Result<AllocationResult<AllocatedPort<Ipv6Addr>>, AllocatorError> {
-        self.allocate_from_tables(src_ip.into(), dst_vpcd, next_header, &self.pools_src66)
+        Self::allocate_from_tables(src_ip.into(), dst_vpcd, next_header, &self.pools_src66)
     }
 
     /// Allocate an IP address and port for the given source IP, dispatching on IP version.
@@ -292,7 +291,6 @@ impl NatAllocator {
         }
     }
     fn allocate_from_tables<I: NatIpWithBitmap>(
-        &self,
         src_ip: IpAddr,
         dst_vpcd: VpcDiscriminant,
         next_header: NextHeader,
@@ -317,8 +315,7 @@ impl NatAllocator {
             })?;
 
         let allow_null = next_header == NextHeader::ICMP || next_header == NextHeader::ICMP6;
-        let mut allocation = pool.allocate(allow_null)?;
-        allocation.set_genid(self.config.genid());
+        let allocation = pool.allocate(allow_null)?;
         let idle_timeout = pool.idle_timeout();
 
         Ok(AllocationResult {
@@ -372,7 +369,7 @@ impl NatAllocator {
         port: NatPort,
     ) -> Result<Allocation, AllocatorError> {
         debug!("Re-reserving {ip} {protocol}:{port}, dst_vpcd:{dst_vpcd}");
-        let mut allocation = match (src_ip, ip) {
+        let allocation = match (src_ip, ip) {
             (IpAddr::V4(src), IpAddr::V4(allocated)) => self
                 .reserve_ipv4_port(protocol, dst_vpcd, src, allocated, port)
                 .map(Allocation::V4)?,
@@ -385,7 +382,6 @@ impl NatAllocator {
                 )));
             }
         };
-        allocation.set_genid(self.config.genid());
         Ok(allocation)
     }
 }

nat/src/masquerade/apalloc/port_alloc.rs

@@ -16,7 +16,6 @@ use concurrency::concurrency_mode;
 use concurrency::sync::atomic::{AtomicBool, AtomicU16, AtomicUsize};
 use concurrency::sync::{Arc, Mutex, RwLock, Weak};
 use concurrency::thread::ThreadId;
-use config::GenId;
 use lpm::prefix::PortRange;
 use std::collections::{BTreeSet, HashMap};
 use std::fmt::Display;
@@ -580,7 +579,6 @@ impl<I: NatIpWithBitmap> Drop for AllocatedPortBlock<I> {
 pub struct AllocatedPort<I: NatIpWithBitmap> {
     port: NatPort,                               // the actual allocated value
     block_allocator: Arc<AllocatedPortBlock<I>>, // block/IP the allocated value belongs to
-    genid: GenId,                                // the generation id of the allocator (late set)
 }
 
 impl<I: NatIpWithBitmap> AllocatedPort<I> {
@@ -589,7 +587,6 @@ impl<I: NatIpWithBitmap> AllocatedPort<I> {
         Self {
             port,
             block_allocator,
-            genid: 0, // initially zero
         }
     }
     #[must_use]
@@ -600,13 +597,6 @@ impl<I: NatIpWithBitmap> AllocatedPort<I> {
     pub fn ip(&self) -> I {
         self.block_allocator.ip()
     }
-    #[must_use]
-    pub fn genid(&self) -> GenId {
-        self.genid
-    }
-    pub fn set_genid(&mut self, genid: GenId) {
-        self.genid = genid;
-    }
 }
 
 impl<I: NatIpWithBitmap> Drop for AllocatedPort<I> {
@@ -668,7 +658,7 @@ struct AllocatedPortBlockMap<I: NatIpWithBitmap>(
 
 impl<I: NatIpWithBitmap> Display for AllocatedPort<I> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}:{} (genid: {})", self.ip(), self.port(), self.genid())
+        write!(f, "{}:{}", self.ip(), self.port())
     }
 }
 

nat/src/masquerade/apalloc/test_alloc.rs

@@ -130,8 +130,8 @@ mod context {
 
     pub fn build_allocator() -> NatAllocator {
         let vpc_table = build_context();
-        let config = MasqueradeConfig::new(&vpc_table, 1);
-        NatAllocator::new(config)
+        let config = MasqueradeConfig::new(&vpc_table);
+        NatAllocator::new(config, 1)
     }
 }
 

nat/src/masquerade/flows.rs

@@ -91,7 +91,7 @@ pub(crate) fn check_masquerading_flow(
     allocator: &NatAllocator,
 ) {
     let config = allocator.config();
-    let genid = config.genid();
+    let genid = allocator.genid();
     if flow_info.genid() == genid {
         return;
     }
@@ -172,7 +172,7 @@ pub(crate) fn check_masquerading_flows<'a>(
     flow_table: &'a FlowTable,
     new_allocator: &mut NatAllocator,
 ) -> FlowTableReadGuard<'a> {
-    let genid = new_allocator.config().genid();
+    let genid = new_allocator.genid();
     debug!("CHECKING flows against new masquerade configuration with genid {genid}...");
     let guard = flow_table.for_each_flow_filtered(
         |_, f| f.is_active(),