Misc changes related to flows

acl-filter/src/tests.rs

@@ -674,10 +674,11 @@ fn attach_related_flow(reply: &mut Packet<TestBuffer>, 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

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<FlowKey> {
     }
 }
 
+/// 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<AtomicU8>) {
+    // 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<Arc<FlowInfo>> = 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<AtomicU8>, 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<AtomicU8>, 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<AtomicU8>, 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::<StubItem>() {
                 let s = stub.status.load(Ordering::Relaxed);

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);

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);

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,
         }
     }
 }

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());

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<FlowStatus> 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,23 +282,28 @@ 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,
         key1: FlowKey,
         flags1: FlowInfoFlags,
         key2: FlowKey,
         flags2: FlowInfoFlags,
-    ) -> (Arc<FlowInfo>, Arc<FlowInfo>) {
-        // keys MUST differ
-        debug_assert!(
-            key1 != key2,
-            "Attempted to build two flows with identical key {key1}"
-        );
+    ) -> Result<(Arc<FlowInfo>, Arc<FlowInfo>), 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<MaybeUninit<Self>> = Arc::new_uninit();
         let mut two: Arc<MaybeUninit<Self>> = 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()))
         }
     }
 

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;
         }