diff --git a/acl-filter/src/tests.rs b/acl-filter/src/tests.rs index 5998aaecf2..eab77405b8 100644 --- a/acl-filter/src/tests.rs +++ b/acl-filter/src/tests.rs @@ -674,10 +674,11 @@ fn attach_related_flow(reply: &mut Packet, fwd_key: FlowKey) -> Arc< let (fwd_flow, reply_flow) = FlowInfo::related_pair( expiry, fwd_key, - FlowInfoFlags::default(), + FlowInfoFlags::default() | FlowInfoFlags::INITIATOR, reply_key, FlowInfoFlags::default(), - ); + ) + .unwrap(); reply_flow.update_status(FlowStatus::Active); reply.meta_mut().flow_info = Some(reply_flow); fwd_flow diff --git a/flow-entry/src/flow_table/concurrent_fuzz.rs b/flow-entry/src/flow_table/concurrent_fuzz.rs index 8403c5ca8d..9a757cd6b7 100644 --- a/flow-entry/src/flow_table/concurrent_fuzz.rs +++ b/flow-entry/src/flow_table/concurrent_fuzz.rs @@ -45,14 +45,14 @@ #![cfg(not(feature = "loom"))] use crate::flow_table::FlowTable; -use concurrency::sync::Arc; use concurrency::sync::atomic::{AtomicU8, Ordering}; +use concurrency::sync::{Arc, Weak}; use concurrency::thread; // `spawn_scoped` is inherent on std's `Builder`, but supplied by `BuilderExt` under shuttle #[cfg_attr(not(feature = "shuttle"), allow(unused_imports))] use concurrency::thread::BuilderExt; use net::FlowKey; -use net::flows::{ExtractRef, FlowInfo}; +use net::flows::{ExtractRef, FlowInfo, FlowInfoFlags}; use std::fmt; use std::time::{Duration, Instant}; @@ -115,7 +115,7 @@ impl bolero::TypeGenerator for Scenario { /// runnable at some scheduling point. Rather than skip degenerate /// shapes, we guarantee at least two of the three op streams contain /// an `Insert` — an `Insert` does model-visible work (writes - /// `FlowInfoLocked` + the atomic status) *and* creates a flow for any + /// `FlowInfoLocked` + the atomic status) *and* creates flows for any /// `Read`/`Flip` ops to land on. Missing inserts are spliced in at a /// driver-chosen offset, so the normalization stays a deterministic /// function of the input and a failure still reproduces from its seed. @@ -154,27 +154,49 @@ fn key_set(base: FlowKey) -> Vec { } } +/// Insert one flow per key. Insert two flows as a pair if possible, with only the "base" key marked +/// as the initiator. Insert the single key otherwise. +fn insert_flows(table: &FlowTable, keys: &[FlowKey], stub_status: &Arc) { + // Far-future expiry so the per-flow timer never fires inside the test + // window — we race the insert path, not the expiry path. (The timer task + // is also cfg'd out entirely under shuttle.) + let expires_at = Instant::now() + Duration::from_hours(1); + let flows: Vec> = match keys { + [fwd_key, rev_key] => { + let (fwd, rev) = FlowInfo::related_pair( + expires_at, + *fwd_key, + FlowInfoFlags::INITIATOR, + *rev_key, + FlowInfoFlags::default(), + ) + .expect("related_pair should succeed for distinct keys"); + vec![fwd, rev] + } + [key] => vec![Arc::new(FlowInfo::new(*key, expires_at))], + _ => panic!("key set should be 1 or 2 keys"), + }; + for fi in flows { + // Stuff a stub item into the locked state so readers and + // flippers have something to race on. + { + let mut guard = fi.locked.write(); + guard.nat_state = Some(Box::new(StubItem { + status: stub_status.clone(), + })); + } + let _ = table.insert_from_arc(&fi); + } +} + /// Apply one [`Op`] across the whole key set. fn apply_op(table: &FlowTable, keys: &[FlowKey], stub_status: &Arc, op: Op) { - for k in keys { - match op { - Op::Insert => { - // Far-future expiry so the per-flow timer never fires - // inside the test window — we race the insert path, not - // the expiry path. (The timer task is also cfg'd out - // entirely under shuttle.) - let fi = Arc::new(FlowInfo::new(*k, Instant::now() + Duration::from_hours(1))); - // Stuff a stub item into the locked state so readers and - // flippers have something to race on. - { - let mut guard = fi.locked.write(); - guard.nat_state = Some(Box::new(StubItem { - status: stub_status.clone(), - })); - } - let _ = table.insert_from_arc(&fi); - } - Op::Lookup => { + match op { + Op::Insert => { + insert_flows(table, keys, stub_status); + } + Op::Lookup => { + for k in keys { // A returned entry must always carry a legal status; // AtomicFlowStatus::load panics on a corrupt u8, so the // load itself is the assertion against torn writes / @@ -183,17 +205,23 @@ fn apply_op(table: &FlowTable, keys: &[FlowKey], stub_status: &Arc, op let _ = fi.status(); } } - Op::Invalidate => { + } + Op::Invalidate => { + for k in keys { if let Some(fi) = table.lookup(k) { fi.invalidate(); } } - Op::ExtendExpiry => { + } + Op::ExtendExpiry => { + for k in keys { if let Some(fi) = table.lookup(k) { let _ = fi.extend_expiry(Duration::from_mins(1)); } } - Op::ReadStubStatus => { + } + Op::ReadStubStatus => { + for k in keys { if let Some(fi) = table.lookup(k) && let Some(stub) = fi .locked @@ -209,7 +237,9 @@ fn apply_op(table: &FlowTable, keys: &[FlowKey], stub_status: &Arc, op assert!(v < STATE_COUNT, "stub status out of range: {v}"); } } - Op::AdvanceStatus => { + } + Op::AdvanceStatus => { + for k in keys { if let Some(fi) = table.lookup(k) && let Some(stub) = fi .locked @@ -284,6 +314,15 @@ impl Scenario { // to make sure the locked state isn't corrupted. table.for_each_flow(|_k, v| { let _ = v.status(); + // We don't always have a related flow, it may have been dropped already, or we might + // have a key that is identical in both directions. + if let Some(related_flow) = v.related.as_ref().and_then(Weak::upgrade) { + assert_ne!( + v.get_flags().is_initiator(), + related_flow.get_flags().is_initiator(), + "exactly one flow of a pair must be the initiator" + ); + } let guard = v.locked.read(); if let Some(stub) = guard.nat_state.as_ref().extract_ref::() { let s = stub.status.load(Ordering::Relaxed); diff --git a/flow-entry/src/flow_table/nf_lookup.rs b/flow-entry/src/flow_table/nf_lookup.rs index d57827b9c6..e741a05a42 100644 --- a/flow-entry/src/flow_table/nf_lookup.rs +++ b/flow-entry/src/flow_table/nf_lookup.rs @@ -200,10 +200,12 @@ mod test { let (flow_1, flow_2) = FlowInfo::related_pair( expires_at, key_1, - FlowInfoFlags::default(), + FlowInfoFlags::default() | FlowInfoFlags::INITIATOR, key_2, FlowInfoFlags::default(), - ); + ) + .unwrap(); + assert_eq!(Arc::weak_count(&flow_1), 1); assert_eq!(Arc::weak_count(&flow_2), 1); assert_eq!(Arc::strong_count(&flow_1), 1); diff --git a/flow-filter/src/tests.rs b/flow-filter/src/tests.rs index 091f7c1627..b89e2032ce 100644 --- a/flow-filter/src/tests.rs +++ b/flow-filter/src/tests.rs @@ -60,7 +60,8 @@ fn attach_flow( packet.meta().compute_flow_flags_forward(), flow_key.reverse(dst_vpcd), packet.meta().compute_flow_flags_reverse(), - ); + ) + .unwrap(); if active { flow_info.update_status(FlowStatus::Active); diff --git a/nat/src/masquerade/nf.rs b/nat/src/masquerade/nf.rs index c03274f0f5..e7e6137afd 100644 --- a/nat/src/masquerade/nf.rs +++ b/nat/src/masquerade/nf.rs @@ -18,7 +18,7 @@ use config::GenId; use flow_entry::flow_table::table::{FlowTable, FlowTableError}; use net::buffer::PacketBufferMut; use net::flow_key::IcmpProtoKey; -use net::flows::{ExtractRef, FlowInfo}; +use net::flows::{ExtractRef, FlowInfo, FlowInfoError}; use net::headers::{TryIp, TryTcp}; use net::ip::UnicastIpAddr; use net::packet::{DoneReason, Packet, VpcDiscriminant}; @@ -59,6 +59,8 @@ pub(crate) enum MasqueradeError { IntendedDrop(&'static str), #[error("Failed to NAT packet: {0}")] NatError(#[from] NatPacketError), + #[error("Failed to create flow state: {0}")] + FlowError(#[from] FlowInfoError), } /// A stateful NAT processor, implementing the [`NetworkFunction`] trait. [`Masquerade`] processes @@ -292,7 +294,7 @@ impl Masquerade { packet.meta().compute_flow_flags_forward(), reverse_key, packet.meta().compute_flow_flags_reverse(), - ); + )?; // set up their NAT state Self::setup_flow_masquerade_state(&forward, forward_state, dst_vpc_id); @@ -537,6 +539,7 @@ impl From<&MasqueradeError> for DoneReason { | MasqueradeError::NatError(_) => DoneReason::NatFailure, MasqueradeError::Bug(_) | MasqueradeError::IntendedDrop(_) => DoneReason::Filtered, MasqueradeError::AllocationFailure(inner) => inner.into(), + MasqueradeError::FlowError(_) => DoneReason::InternalFailure, } } } diff --git a/nat/src/portfw/nf.rs b/nat/src/portfw/nf.rs index 88781bf444..22a46b03ae 100644 --- a/nat/src/portfw/nf.rs +++ b/nat/src/portfw/nf.rs @@ -118,9 +118,7 @@ impl PortForwarder { let Ok((fw_key, rev_key)) = build_portfw_flow_keys(packet, new_dst_ip, new_dst_port, entry.dst_vpcd) else { - warn!( - "Failed to build flow keys for port forwarding: {dst_ip}:{dst_port} -> {new_dst_ip}:{new_dst_port}" - ); + warn!("Failed to build flow keys: {dst_ip}:{dst_port} -> {new_dst_ip}:{new_dst_port}"); packet.done(DoneReason::InternalFailure); return; }; @@ -144,13 +142,17 @@ impl PortForwarder { // create a pair of related flow entries (outside the flow table). Timeout is set according to the rule matched let timeout = Instant::now() + entry.init_timeout(); - let (fw_flow, rev_flow) = FlowInfo::related_pair( + let Ok((fw_flow, rev_flow)) = FlowInfo::related_pair( timeout, fw_key, packet.meta().compute_flow_flags_forward(), rev_key, packet.meta().compute_flow_flags_reverse(), - ); + ) else { + debug!("Failed to build flow pair for port forwarded flow"); + packet.done(DoneReason::InternalFailure); + return; + }; // set the generation id for the flow fw_flow.set_genid_pair(self.pipeline_data.genid()); diff --git a/net/src/flows/flow_info.rs b/net/src/flows/flow_info.rs index 23e86981df..b23e52b838 100644 --- a/net/src/flows/flow_info.rs +++ b/net/src/flows/flow_info.rs @@ -30,6 +30,8 @@ pub enum FlowInfoError { NoSuchStatus(u8), #[error("Timeout unchanged: would go backwards")] TimeoutUnchanged, + #[error("Invalid flow pair: {0}")] + InvalidPair(String), } #[repr(u8)] @@ -140,12 +142,18 @@ impl From for AtomicFlowStatus { bitflags! { #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] pub struct FlowInfoFlags: u8 { - const REQ_STATIC_NAT_SRC = 0b0000_0001; /* Packet requires static NAT (source) */ - const REQ_STATIC_NAT_DST = 0b0000_0010; /* Packet requires static NAT (destination) */ + const INITIATOR = 0b0000_0001; /* the flow is the initiator within a pair */ + const REQ_STATIC_NAT_SRC = 0b0000_0010; /* Packet requires static NAT (source) */ + const REQ_STATIC_NAT_DST = 0b0000_0100; /* Packet requires static NAT (destination) */ } } impl FlowInfoFlags { + #[must_use] + pub const fn is_initiator(&self) -> bool { + self.contains(FlowInfoFlags::INITIATOR) + } + #[must_use] pub const fn requires_static_nat_src(self) -> bool { self.contains(FlowInfoFlags::REQ_STATIC_NAT_SRC) @@ -274,10 +282,10 @@ impl FlowInfo { /// to call this function when a couple of related flow entries are needed and later insert them in the /// flow-table. /// - /// # Panics - /// This function panics if two equal keys are provided + /// # Errors + /// This function fails if two identical keys are provided or if one (and only one) of the flows + /// is not flagged as initiator #[allow(clippy::missing_panics_doc)] - #[must_use] #[allow(clippy::unwrap_used)] pub fn related_pair( expires_at: Instant, @@ -285,12 +293,17 @@ impl FlowInfo { flags1: FlowInfoFlags, key2: FlowKey, flags2: FlowInfoFlags, - ) -> (Arc, Arc) { - // keys MUST differ - debug_assert!( - key1 != key2, - "Attempted to build two flows with identical key {key1}" - ); + ) -> Result<(Arc, Arc), FlowInfoError> { + if key1 == key2 { + return Err(FlowInfoError::InvalidPair(format!( + "Attempted to build a flow pair with identical keys {key1}" + ))); + } + if flags1.is_initiator() == flags2.is_initiator() { + return Err(FlowInfoError::InvalidPair( + "One of the flows must be the initiator".to_string(), + )); + } let mut one: Arc> = Arc::new_uninit(); let mut two: Arc> = Arc::new_uninit(); @@ -321,7 +334,7 @@ impl FlowInfo { .set_related(one_weak), ); // turn back into Arc's - (one.assume_init(), two.assume_init()) + Ok((one.assume_init(), two.assume_init())) } } diff --git a/net/src/packet/meta.rs b/net/src/packet/meta.rs index 356e72cdc9..826c8fcb21 100644 --- a/net/src/packet/meta.rs +++ b/net/src/packet/meta.rs @@ -263,6 +263,8 @@ impl PacketMeta { #[must_use] pub fn compute_flow_flags_forward(&self) -> FlowInfoFlags { let mut flags = FlowInfoFlags::default(); + flags.insert(FlowInfoFlags::INITIATOR); + if self.requires_static_nat_src() { flags |= FlowInfoFlags::REQ_STATIC_NAT_SRC; }