diff --git a/acl-filter/src/tests.rs b/acl-filter/src/tests.rs index c7a29f71dd..42af55b92f 100644 --- a/acl-filter/src/tests.rs +++ b/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(), diff --git a/mgmt/src/processor/proc.rs b/mgmt/src/processor/proc.rs index 04b8f7e039..84b76ab33e 100644 --- a/mgmt/src/processor/proc.rs +++ b/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"); } diff --git a/nat/src/masquerade/allocator_writer.rs b/nat/src/masquerade/allocator_writer.rs index e8305303a8..fe8469f477 100644 --- a/nat/src/masquerade/allocator_writer.rs +++ b/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, 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,8 +106,12 @@ 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 @@ -128,6 +119,7 @@ impl NatAllocatorWriter { && 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..."); diff --git a/nat/src/masquerade/apalloc/mod.rs b/nat/src/masquerade/apalloc/mod.rs index 1b78cdb315..51338fccfc 100644 --- a/nat/src/masquerade/apalloc/mod.rs +++ b/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,6 +203,7 @@ impl Display for Allocation { #[derive(Debug)] pub struct NatAllocator { config: MasqueradeConfig, + genid: AtomicI64, pools_src44: PoolTable, pools_src66: PoolTable, randomize: bool, @@ -224,10 +211,11 @@ pub struct NatAllocator { 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>, 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>, 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( - &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 { 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) } } diff --git a/nat/src/masquerade/apalloc/port_alloc.rs b/nat/src/masquerade/apalloc/port_alloc.rs index fd394e74a8..2fa67d0ef6 100644 --- a/nat/src/masquerade/apalloc/port_alloc.rs +++ b/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 Drop for AllocatedPortBlock { pub struct AllocatedPort { port: NatPort, // the actual allocated value block_allocator: Arc>, // block/IP the allocated value belongs to - genid: GenId, // the generation id of the allocator (late set) } impl AllocatedPort { @@ -589,7 +587,6 @@ impl AllocatedPort { Self { port, block_allocator, - genid: 0, // initially zero } } #[must_use] @@ -600,13 +597,6 @@ impl AllocatedPort { 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 Drop for AllocatedPort { @@ -668,7 +658,7 @@ struct AllocatedPortBlockMap( impl Display for AllocatedPort { 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()) } } diff --git a/nat/src/masquerade/apalloc/test_alloc.rs b/nat/src/masquerade/apalloc/test_alloc.rs index 7c3520db62..a8116404bf 100644 --- a/nat/src/masquerade/apalloc/test_alloc.rs +++ b/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) } } diff --git a/nat/src/masquerade/flows.rs b/nat/src/masquerade/flows.rs index fe95f73023..6fffeeb6f7 100644 --- a/nat/src/masquerade/flows.rs +++ b/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(), diff --git a/nat/src/masquerade/nf.rs b/nat/src/masquerade/nf.rs index a6ed1b6748..f34b1d9a6a 100644 --- a/nat/src/masquerade/nf.rs +++ b/nat/src/masquerade/nf.rs @@ -8,12 +8,13 @@ use crate::common::NatFlowStatus; use crate::masquerade::NatAllocatorWriter; use crate::masquerade::allocation::{AllocationResult, AllocatorError}; use crate::masquerade::allocator_writer::NatAllocatorReader; -use crate::masquerade::apalloc::Allocation; +use crate::masquerade::apalloc::{Allocation, NatAllocator}; use crate::masquerade::flows::check_masquerading_flow; use crate::masquerade::packet::{NatPacketError, NatTranslate, masquerade}; use crate::masquerade::protocol::next_flow_status; use crate::masquerade::state::MasqueradeState; use concurrency::sync::{Arc, Weak}; +use config::GenId; use flow_entry::flow_table::table::{FlowTable, FlowTableError}; use net::buffer::PacketBufferMut; use net::flow_key::IcmpProtoKey; @@ -31,7 +32,7 @@ use std::time::{Duration, Instant}; use tracing::{debug, error, warn}; #[derive(Debug, thiserror::Error)] -enum MasqueradeError { +pub(crate) enum MasqueradeError { #[error("Unexpected failure: {0}")] Bug(&'static str), #[error("failure to get transport header")] @@ -244,9 +245,9 @@ impl Masquerade { initial_flow_key: &FlowKey, current_flow_key: &FlowKey, alloc: AllocationResult, - ) -> Result<(), MasqueradeError> { + genid: GenId, + ) -> Result, MasqueradeError> { let idle_timeout = alloc.idle_timeout; - let genid = alloc.allocation.genid(); // src and dst vpc of this packet let src_vpc_id = packet.meta().src_vpcd.unwrap_or_else(|| unreachable!()); @@ -298,7 +299,7 @@ impl Masquerade { debug_assert!(false, "reverse flow insert failed: {e:?}"); return Err(MasqueradeError::CapacityExceeded); } - Ok(()) + Ok(forward) } fn new_reverse_session( @@ -395,6 +396,9 @@ impl Masquerade { .allocate(dst_vpcd, src_ip, initial_flow_key.proto()) .map_err(MasqueradeError::AllocationFailure)?; + // The generation the installed allocator serves + let genid = allocator.genid(); + // Forbid addresses we won't know how to translate. This is a work around of a larger change if let Err(addr) = UnicastIpAddr::try_from(alloc.allocation.ip()) { error!("Allocated address {addr} won't be usable: not unicast"); @@ -404,13 +408,8 @@ impl Masquerade { debug!("{nfi}: Allocated: {alloc}"); // create flow pair - self.create_flow_pair(packet, &initial_flow_key, ¤t_flow_key, alloc)?; - - // lookup the flow (forward) just created. We should always find it. - let installed = self - .flow_table - .lookup(&initial_flow_key) - .ok_or(MasqueradeError::Bug("Unexpected flow lookup failure"))?; + let installed = + self.create_flow_pair(packet, &initial_flow_key, ¤t_flow_key, alloc, genid)?; // check that the masquerade state is readable let translate = installed @@ -427,32 +426,42 @@ impl Masquerade { return Err(e.into()); } - // .. and check whether the allocation we made and stored in the flows is still fine - // with the current allocator. This counters for the potential race where we got a port - // allocated but before we could install the flows, a new config was applied. If that - // happened, our flow would not be checked against the new config. So we'd have a flow - // with an allocation drawn from an allocator that was replaced by a newer one, and the - // new allocator would not be aware of that allocation. So, here we repeat the logic that - // checks flows against a new config / allocation. - match self.allocator.get() { - None => { - // allocator got removed. Get rid of the flows and drop the packet. - installed.invalidate_pair(); - Err(MasqueradeError::IntendedDrop("allocator was removed")) - } - Some(allocator) => { - check_masquerading_flow( - installed.flowkey(), - installed.as_ref(), - allocator.as_ref(), - ); - if installed.is_active() { - Ok(()) - } else { - // we invalidated the flow. Signal that packet should be dropped - Err(MasqueradeError::IntendedDrop("Config changed")) - } + // It may happen that between the time we got an allocation and the moment we installed the flows + // the allocator was swapped. So, here we have to check if the allocator we used is still there: + // it may have been removed or replaced. If so, the newly installed flows may no longer be valid + // and we have to remove them. Also, the genid may have changed and we need to bump it. + self.recheck_flow(&allocator, &installed) + } + + /// Re-check a freshly installed flow against the latest allocator, that could have been installed + /// while we were installing a flow. + pub(crate) fn recheck_flow( + &self, + used_allocator: &Arc, + flow: &Arc, + ) -> Result<(), MasqueradeError> { + let Some(current) = self.allocator.get() else { + debug!("Allocator got removed!"); + flow.invalidate_pair(); + return Err(MasqueradeError::IntendedDrop("Allocator got removed")); + }; + if Arc::ptr_eq(used_allocator, ¤t) { + // Allocator did not change. So the allocation of the newly installed flow is + // still valid. However, the genid of the allocator may have been bumped. + // So, update it in the new flow. + if flow.genid() != current.genid() { + flow.set_genid_pair(current.genid()); } + return Ok(()); + } + debug!("NAT allocator got updated. Re-checking newly-installed flow..."); + check_masquerading_flow(flow.flowkey(), flow.as_ref(), current.as_ref()); + if flow.is_active() { + Ok(()) + } else { + Err(MasqueradeError::IntendedDrop( + "Flow is not valid with the new allocator", + )) } } diff --git a/nat/src/masquerade/test.rs b/nat/src/masquerade/test.rs index 962b5f5a3b..c87a2f4aad 100644 --- a/nat/src/masquerade/test.rs +++ b/nat/src/masquerade/test.rs @@ -6,7 +6,7 @@ use crate::common::{NatAction, NatFlowStatus}; use crate::masquerade::state::MasqueradeState; use crate::masquerade::{MasqueradeConfig, NatAllocatorWriter}; -use crate::{IcmpErrorHandler, Masquerade}; +use crate::{IcmpErrorHandler, Masquerade, NatPort}; use ahash::HashMap; use common::cliprovider::Frame; use concurrency::sync::Arc; @@ -21,8 +21,8 @@ use flow_entry::flow_table::{FlowLookup, FlowTable}; use flow_filter::{FlowFilter, FlowFilterContext, FlowFilterContextWriter}; use net::buffer::{PacketBufferMut, TestBuffer}; use net::eth::mac::Mac; -use net::flows::FlowStatus; use net::flows::flow_info_item::ExtractRef; +use net::flows::{FlowInfo, FlowStatus}; use net::headers::TryTcpMut; use net::headers::{ EmbeddedTransport, TryEmbeddedTransport as _, TryIcmp4, TryInnerIpv4, TryIpv4, TryUdp, @@ -33,7 +33,7 @@ use net::ip::NextHeader; use net::packet::test_utils::build_test_tcp_ipv4_packet; use net::packet::test_utils::{ IcmpEchoDirection, build_test_icmp4_destination_unreachable_packet, build_test_icmp4_echo, - build_test_udp_ipv4_frame, + build_test_udp_ipv4_frame, build_test_udp_ipv4_packet, }; use net::packet::{DoneReason, Packet, VpcDiscriminant}; use net::tcp::TruncatedTcp; @@ -53,7 +53,7 @@ const ONE_MINUTE: Duration = Duration::from_mins(1); use crate::static_nat::test::build_gwconfig_from_overlay; fn test_case(msg: &str) { - debug!("{}", Frame(msg)); + println!("{}", Frame(msg)); } #[derive(Default)] @@ -113,7 +113,7 @@ fn test_setup( let overlay = overlay.validate().unwrap(); // build the configuration for the nat allocator - let nat_config = MasqueradeConfig::new(overlay.vpc_table(), genid); + let nat_config = MasqueradeConfig::new(overlay.vpc_table()); // build the config for the test flow filter and the flow filter let peerings: Vec<_> = nat_config @@ -126,7 +126,7 @@ fn test_setup( let (flow_table, pipeline, mut alloc_writer) = setup_pipeline_masquerade(flow_filter); // setup the NAT allocator - alloc_writer.update_nat_allocator(nat_config, &flow_table); + alloc_writer.update_nat_allocator(nat_config, genid, &flow_table); (flow_table, pipeline, alloc_writer) } @@ -386,8 +386,8 @@ async fn test_full_config() { // Check that we can validate the allocator let (mut nat, mut allocator) = Masquerade::new_with_defaults(); - let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table(), 1); - allocator.update_nat_allocator(nat_config, &flow_table); + let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table()); + allocator.update_nat_allocator(nat_config, 1, &flow_table); // No NAT let (orig_src, orig_dst) = ("8.8.8.8", "9.9.9.9"); @@ -461,8 +461,8 @@ async fn test_full_config() { let new_config = build_gwconfig_from_overlay(build_overlay_2vpcs()) .validate() .unwrap(); - let nat_config = MasqueradeConfig::new(new_config.external().overlay().vpc_table(), 2); - allocator.update_nat_allocator(nat_config, &flow_table); + let nat_config = MasqueradeConfig::new(new_config.external().overlay().vpc_table()); + allocator.update_nat_allocator(nat_config, 2, &flow_table); // Check existing connection // TODO: We should drop this connection after updating the allocator in the future, as a @@ -554,8 +554,8 @@ fn test_full_config_no_nat() { // Check that we can validate the allocator let (_, mut allocator) = Masquerade::new_with_defaults(); - let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table(), 1); - allocator.update_nat_allocator(nat_config, &FlowTable::new(16)); + let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table()); + allocator.update_nat_allocator(nat_config, 1, &FlowTable::new(16)); } fn check_packet_icmp_echo( @@ -625,8 +625,8 @@ async fn test_icmp_echo_nat() { // Check that we can validate the allocator let (mut nat, mut allocator) = Masquerade::new_with_defaults(); - let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table(), 1); - allocator.update_nat_allocator(nat_config, &FlowTable::new(16)); + let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table()); + allocator.update_nat_allocator(nat_config, 1, &FlowTable::new(16)); // No NAT let (orig_src, orig_dst, orig_identifier) = (addr_v4("8.8.8.8"), addr_v4("9.9.9.9"), 1337); @@ -940,8 +940,8 @@ async fn test_default_expose() { // Check that we can validate the allocator let (mut nat, mut allocator) = Masquerade::new_with_defaults(); - let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table(), 1); - allocator.update_nat_allocator(nat_config, &FlowTable::new(16)); + let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table()); + allocator.update_nat_allocator(nat_config, 1, &FlowTable::new(16)); // Using the expose with a prefix let (orig_src, orig_dst, orig_src_port, orig_dst_port) = ("1.1.0.1", "3.3.3.3", 9999, 443); @@ -1160,10 +1160,10 @@ async fn test_full_config_unidirectional_nat_overlapping_destination() { // Build NAT stage let (mut nat, mut allocator) = Masquerade::new_with_defaults(); - let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table(), 1); + let nat_config = MasqueradeConfig::new(config.external().overlay().vpc_table()); // Check that we can validate the allocator - allocator.update_nat_allocator(nat_config, &FlowTable::new(16)); + allocator.update_nat_allocator(nat_config, 1, &FlowTable::new(16)); // NAT: expose12 <-> expose21 let (orig_src, orig_dst, orig_src_port, orig_dst_port) = ("1.0.0.18", "5.0.0.5", 9998, 443); @@ -1228,13 +1228,13 @@ async fn test_full_config_unidirectional_nat_overlapping_destination() { let mut allocator = NatAllocatorWriter::new(); let mut nat = Masquerade::new("masquerade", flow_table.clone(), allocator.get_reader()); let nat_config = - MasqueradeConfig::new(config.external().overlay().vpc_table(), 2).set_randomize(false); + MasqueradeConfig::new(config.external().overlay().vpc_table()).set_randomize(false); // Check that we can validate the allocator // // When we build the allocator, turn off randomness to check whether we may get collisions // for port allocation - allocator.update_nat_allocator(nat_config, &flow_table); + allocator.update_nat_allocator(nat_config, 2, &flow_table); // NAT: expose12 <-> expose21 let (orig_src, orig_dst, orig_src_port, orig_dst_port) = ("1.0.0.18", "5.0.0.5", 9998, 443); @@ -1620,8 +1620,8 @@ async fn test_masquerade_reconfig_keep_flow() { // update the NAT allocator with an identical config let overlay = build_overlay_2vpcs().validate().unwrap(); - let nat_config = MasqueradeConfig::new(overlay.vpc_table(), genid + 1); - allocw.update_nat_allocator(nat_config, &flow_table); + let nat_config = MasqueradeConfig::new(overlay.vpc_table()); + allocw.update_nat_allocator(nat_config, genid + 1, &flow_table); // process a packet: it should hit identical flows, except for genid let packet = tcp_packet_to_masquerade(); @@ -1656,8 +1656,8 @@ async fn test_masquerade_reconfig_drop_flow() { // update the NAT allocator with an identical config let overlay = build_overlay_2vpcs_modified().validate().unwrap(); - let nat_config = MasqueradeConfig::new(overlay.vpc_table(), genid + 1); - allocw.update_nat_allocator(nat_config, &flow_table); + let nat_config = MasqueradeConfig::new(overlay.vpc_table()); + allocw.update_nat_allocator(nat_config, genid + 1, &flow_table); // process a packet: it should hit identical flows let packet = tcp_packet_to_masquerade(); @@ -1671,3 +1671,99 @@ async fn test_masquerade_reconfig_drop_flow() { tokio::time::sleep(Duration::from_secs(1)).await; assert_eq!(flow_table.active_len(), Some(0)); } + +#[tokio::test] +#[cfg_attr(not(emulated), traced_test)] +// tests that on a new config the genid is updated, whether the allocator changed or not +async fn test_genid_updated_on_reconfig() { + // build some config with genid 1 + let genid = 1; + let (flow_table, _pipeline, mut allocw) = test_setup(genid, &build_overlay_2vpcs()); + + // check that allocator's genid matches + let observed_genid = allocw.get_reader().get().unwrap().genid(); + assert_eq!(observed_genid, genid); + + // update the masquerade config, but with NO change so that allocator is the same + let overlay = &build_overlay_2vpcs().validate().unwrap(); + let nat_config = MasqueradeConfig::new(overlay.vpc_table()); + allocw.update_nat_allocator(nat_config, genid + 1, &flow_table); + + let observed_genid = allocw.get_reader().get().unwrap().genid(); + assert_eq!(observed_genid, genid + 1); + + // update the NAT allocator with a distinct config, causing allocator to be updated + let overlay = build_overlay_2vpcs_modified().validate().unwrap(); + let nat_config = MasqueradeConfig::new(overlay.vpc_table()); + allocw.update_nat_allocator(nat_config, genid + 2, &flow_table); + + let observed_genid = allocw.get_reader().get().unwrap().genid(); + assert_eq!(observed_genid, genid + 2); +} + +fn get_flow_allocation(flow_info: &FlowInfo) -> Option<(IpAddr, NatPort)> { + let locked = flow_info.locked.read(); + let alloc = locked + .nat_state + .as_ref()? + .extract_ref::()? + .allocation()?; + Some((alloc.ip(), alloc.port())) +} + +#[tokio::test] +#[cfg_attr(not(emulated), traced_test)] +async fn test_recheck_flow_when_allocator_is_kept() { + let genid = 1; + let (mut nat, mut allocw) = Masquerade::new_with_defaults(); + let flow_table = nat.sessions().clone(); + + let overlay = build_overlay_2vpcs().validate().unwrap(); + let nat_config = MasqueradeConfig::new(overlay.vpc_table()); + allocw.update_nat_allocator(nat_config, genid, &flow_table); + + test_case("Create a masqueraded flow with the current config"); + let mut packet: Packet = build_test_udp_ipv4_packet("1.1.0.1", "3.3.3.1", 4321, 80); + packet.meta_mut().set_overlay(true); + packet.meta_mut().set_masquerade(true); + packet.meta_mut().src_vpcd = Some(vpcd(100)); + packet.meta_mut().dst_vpcd = Some(vpcd(200)); + let flow_key = FlowKey::try_from(&packet).unwrap(); + let out: Vec<_> = nat.process(std::iter::once(packet)).collect(); + assert_eq!(out[0].get_done(), None); + + test_case("Check that the flow is there"); + let installed = flow_table.lookup(&flow_key).expect("Flow must be there"); + let reverse = installed + .related + .as_ref() + .expect("Flow must have a related flow") + .upgrade() + .expect("Reverse flow must be there"); + + assert!(installed.is_active()); + assert!(reverse.is_active()); + assert_eq!(reverse.genid(), genid); + assert_eq!(installed.genid(), genid); + let allocation = get_flow_allocation(&installed).expect("Flow must have an allocation"); + + test_case("Advance the genid of the installed allocator (without checking flows)"); + // we advance the genid of the allocator, without checking the flows to simulate the + // race where a flow would be installed and race against the check made when a new + // config would be applied, to exercise `recheck_flow` + let allocator = allocw.get_reader().get().expect("Allocator must be there"); + allocator.set_genid(genid + 1); + assert_eq!(installed.genid(), genid); // flow was left behind by the sweep + + test_case("Re-check the flow: it must be upgraded, not invalidated"); + assert!(nat.recheck_flow(&allocator, &installed).is_ok()); + assert!(installed.is_active()); + assert!(reverse.is_active()); + assert_eq!(installed.genid(), genid + 1); + assert_eq!(reverse.genid(), genid + 1); + assert_eq!(get_flow_allocation(&installed), Some(allocation)); // flow keeps its allocation + + test_case("Done (logs beyond this point are irrelevant)"); + tokio::time::sleep(Duration::from_secs(1)).await; + assert_eq!(flow_table.active_len(), Some(2)); +} diff --git a/nat/src/test.rs b/nat/src/test.rs index 664daeacd4..09d9b48596 100644 --- a/nat/src/test.rs +++ b/nat/src/test.rs @@ -119,8 +119,8 @@ fn setup_masq_pipeline( // Masquerade let mut allocator = NatAllocatorWriter::new(); let masquerade = Masquerade::new("masquerade", flow_table.clone(), allocator.get_reader()); - let masquerade_config = MasqueradeConfig::new(overlay.vpc_table(), 1); - allocator.update_nat_allocator(masquerade_config, &flow_table); + let masquerade_config = MasqueradeConfig::new(overlay.vpc_table()); + allocator.update_nat_allocator(masquerade_config, 1, &flow_table); if let Some(state) = allocator.get_reader().get() { println!("{state}"); }