Clean-up and minor fixes for flow-filter

[qmonnet]

Split off from #1341.

Please refer to individual commit logs for details.

[Copilot]

Pull request overview

This PR is a refactor/cleanup of the flow-filter stage’s flow-info handling, extracting common “active flow-info” logic into net::packet::Packet and reorganizing flow-filter helpers for bypassing and NAT requirement derivation.

Changes:

• Add Packet::active_flow_info() helper in net to centralize “flow-info exists and is Active” checks.
• Refactor flow-filter to use the new helper, reorganizing the bypass and multiple-match handling paths.
• Move NAT-requirements-setting helpers into FlowFilter impl and streamline call sites.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
net/src/packet/mod.rs	Adds active_flow_info() helper for retrieving only Active flow-info entries.
flow-filter/src/lib.rs	Refactors flow-info usage to use active_flow_info(), adjusts bypass logic, and reorganizes NAT requirement helpers.

Comments suppressed due to low confidence (1)

flow-filter/src/lib.rs:29

confidence: 9
tags: [logic]

`FlowInfo`/`Packet::active_flow_info()` uses `concurrency::sync::Arc` (which becomes `shuttle::sync::Arc` under the workspace’s `shuttle` feature), but this file imports `std::sync::Arc` and now uses it in `Option<&Arc<FlowInfo>>` signatures. That will fail to typecheck (and likely to compile) when building with `--features shuttle` (e.g., via `nat`’s shuttle tests), because `PacketMeta.flow_info` is `Option<concurrency::sync::Arc<FlowInfo>>`.

Use `concurrency::sync::Arc` for `FlowInfo`-related arcs here (you may need to add a direct dependency on `concurrency` in `flow-filter/Cargo.toml`), and if you still need `std::sync::Arc` for `PipelineData`, alias one of them (e.g., `StdArc`).

use std::collections::HashSet;
use std::fmt::Display;
use std::net::IpAddr;
use std::num::NonZero;
use std::sync::Arc;
use tracing::{debug, error};

</details>



---

You can also share your feedback on Copilot code review. [Take the survey](https://www.surveymonkey.com/r/XP6L3XJ).

[qmonnet]

As discussed offline, I dropped the last commit and reverted the change in set_nat_requirements_from_flow_info() (I also added a note in the relevant commit description as to why we don't do the flow status check in that function).

Diff (click to expand)

diff --git c/flow-filter/src/lib.rs w/flow-filter/src/lib.rs
index b7ffade26aec..5a8b712eddcd 100644
--- c/flow-filter/src/lib.rs
+++ w/flow-filter/src/lib.rs
@@ -67,11 +67,9 @@ impl FlowFilter {
     ) {
         let nfi = &self.name;
         let genid = self.pipeline_data.genid();
-        let flow_info = packet.active_flow_info().cloned();
-        let dst_vpcd_from_flow_info = Self::dst_vpcd_from_flow_info(flow_info.as_ref());
 
-        // Bypass flow-filter if packet has flow-info and it is not outdated
-        if self.bypass_with_flow_info(packet, flow_info.as_ref(), genid, dst_vpcd_from_flow_info) {
+        // bypass flow-filter if packet has flow-info and it is not outdated
+        if self.bypass_with_flow_info(packet, genid) {
             return;
         }
 
@@ -112,12 +110,9 @@ impl FlowFilter {
                     "{nfi}: Found multiple matches for destination VPC for flow {tuple}. Checking for a flow table entry..."
                 );
 
-                match self.check_packet_flow_info(flow_info.as_ref(), dst_vpcd_from_flow_info) {
+                match self.check_packet_flow_info(packet) {
                     Ok(Some(dst_vpcd)) => {
-                        if self
-                            .set_nat_requirements_from_flow_info(packet, flow_info.as_ref())
-                            .is_ok()
-                        {
+                        if Self::set_nat_requirements_from_flow_info(packet).is_ok() {
                             Some(dst_vpcd)
                         } else {
                             debug!("{nfi}: Failed to set NAT requirements from flow info");
@@ -172,17 +167,14 @@ impl FlowFilter {
     }
 
     /// Check if flow-info is up-to-date and allows bypassing the main filtering logic.
-    /// Assumes `dst_vpcd_from_flow_info` was retrieved from `flow_info_opt`.
     fn bypass_with_flow_info<Buf: PacketBufferMut>(
         &self,
         packet: &mut Packet<Buf>,
-        flow_info_opt: Option<&Arc<FlowInfo>>,
         genid: i64,
-        dst_vpcd_from_flow_info: Option<VpcDiscriminant>,
     ) -> bool {
         let nfi = &self.name;
 
-        let Some(flow_info) = flow_info_opt else {
+        let Some(flow_info) = packet.active_flow_info() else {
             return false;
         };
         let flow_genid = flow_info.genid();
@@ -191,33 +183,32 @@ impl FlowFilter {
             return false;
         }
 
+        let vpcd = Self::dst_vpcd_from_flow_info(flow_info);
+
         debug!("{nfi}: Packet can bypass filter due to flow {flow_info}");
 
-        if self
-            .set_nat_requirements_from_flow_info(packet, flow_info_opt)
-            .is_err()
-        {
+        if Self::set_nat_requirements_from_flow_info(packet).is_err() {
             debug!("{nfi}: Failed to set nat requirements");
             return false;
         }
-        packet.meta_mut().dst_vpcd = dst_vpcd_from_flow_info;
+        packet.meta_mut().dst_vpcd = vpcd;
         true
     }
 
     /// Attempt to determine destination VPC from packet's flow-info.
-    /// Assumes `dst_vpcd_from_flow_info` was retrieved from `flow_info`.
-    fn check_packet_flow_info(
+    fn check_packet_flow_info<Buf: PacketBufferMut>(
         &self,
-        flow_info: Option<&Arc<FlowInfo>>,
-        dst_vpcd_from_flow_info: Option<VpcDiscriminant>,
+        packet: &mut Packet<Buf>,
     ) -> Result<Option<VpcDiscriminant>, DoneReason> {
         let nfi = &self.name;
 
-        if flow_info.is_none() {
+        let Some(flow_info) = packet.active_flow_info() else {
             return Ok(None);
-        }
+        };
 
-        let Some(dst_vpcd) = dst_vpcd_from_flow_info else {
+        let vpcd = Self::dst_vpcd_from_flow_info(flow_info);
+
+        let Some(dst_vpcd) = vpcd else {
             debug!("{nfi}: No VPC discriminant found, dropping packet");
             return Err(DoneReason::Unroutable);
         };
@@ -337,11 +328,16 @@ impl FlowFilter {
 
     /// Set NAT requirements on the packet based on packet's flow-info, if any.
     fn set_nat_requirements_from_flow_info<Buf: PacketBufferMut>(
-        &self,
         packet: &mut Packet<Buf>,
-        flow_info: Option<&Arc<FlowInfo>>,
     ) -> Result<(), ()> {
-        let locked_info = flow_info.ok_or(())?.locked.read().map_err(|_| ())?;
+        let locked_info = packet
+            .meta()
+            .flow_info
+            .as_ref()
+            .ok_or(())?
+            .locked
+            .read()
+            .map_err(|_| ())?;
         let needs_stateful_nat = locked_info.nat_state.is_some();
         let needs_port_forwarding = locked_info.port_fw_state.is_some();
         drop(locked_info);
@@ -361,8 +357,8 @@ impl FlowFilter {
 
     /// Extract the destination VPC discriminant from a flow-info entry.
     /// Panic if the lock has been poisoned.
-    fn dst_vpcd_from_flow_info(flow_info: Option<&Arc<FlowInfo>>) -> Option<VpcDiscriminant> {
-        flow_info?
+    fn dst_vpcd_from_flow_info(flow_info: &Arc<FlowInfo>) -> Option<VpcDiscriminant> {
+        flow_info
             .locked
             .read()
             .unwrap()

Thanks a lot for the review!