refactor(net): Drop the Uni<T> wrapper around packets

Uni<T> existed to tag a packet for unidirectional FlowKey extraction
back when bidirectional keys were also supported. Bidirectional keys are
gone, so Uni is not longer necessary, it simply adds dead weight that
every call site has to thread through, and makes the code less clear.

Convert the TryFrom impl to accept &Packet<Buf> directly, delete the
Uni struct, and drop the wrapper at the call sites.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Quentin Monnet <qmo@qmon.net>

Author: qmonnet

flow-entry/src/flow_table/nf_lookup.rs

@@ -12,7 +12,6 @@ use pipeline::NetworkFunction;
 
 use crate::flow_table::FlowTable;
 use net::FlowKey;
-use net::flow_key;
 
 use tracectl::trace_target;
 trace_target!("flow-lookup", LevelFilter::INFO, &["pipeline"]);
@@ -40,7 +39,7 @@ impl<Buf: PacketBufferMut> NetworkFunction<Buf> for FlowLookup {
         input.filter_map(move |mut packet| {
             let nfi = &self.name;
             if !packet.is_done() && packet.meta().is_overlay() && packet.meta().dst_vpcd.is_none() {
-                if let Ok(flow_key) = FlowKey::try_from(flow_key::Uni(&packet)) {
+                if let Ok(flow_key) = FlowKey::try_from(&packet) {
                     if let Some(flow_info) = self.flow_table.lookup(&flow_key) {
                         debug!("{nfi}: Tagging packet with flow info for flow key {flow_key}",);
                         packet.meta_mut().flow_info = Some(flow_info);
@@ -101,7 +100,7 @@ mod test {
         packet.meta_mut().set_overlay(true);
 
         // Insert matching flow entry
-        let flow_key = FlowKey::try_from(net::flow_key::Uni(&packet)).unwrap();
+        let flow_key = FlowKey::try_from(&packet).unwrap();
         let flow_info = FlowInfo::new(flow_key, Instant::now() + Duration::from_secs(10));
         flow_table.insert(flow_info).unwrap();
 
@@ -130,7 +129,7 @@ mod test {
             input: Input,
         ) -> impl Iterator<Item = Packet<Buf>> + 'a {
             input.filter_map(move |packet| {
-                let flow_key = FlowKey::try_from(net::flow_key::Uni(&packet)).unwrap();
+                let flow_key = FlowKey::try_from(&packet).unwrap();
                 let flow_info = FlowInfo::new(flow_key, Instant::now() + self.timeout);
                 self.flow_table
                     .insert(flow_info)
@@ -193,8 +192,8 @@ mod test {
             packet_2.meta_mut().set_overlay(true);
 
             // build keys for the packets
-            let key_1 = FlowKey::try_from(net::flow_key::Uni(&packet_1)).unwrap();
-            let key_2 = FlowKey::try_from(net::flow_key::Uni(&packet_2)).unwrap();
+            let key_1 = FlowKey::try_from(&packet_1).unwrap();
+            let key_2 = FlowKey::try_from(&packet_2).unwrap();
 
             // create a pair of related flow entries; flow_2 will get a longer timeout
             let expires_at = tokio::time::Instant::now().into_std() + Duration::from_secs(2);
@@ -251,8 +250,8 @@ mod test {
         let mut packet_2 = build_test_udp_ipv4_packet("192.168.1.1", "20.0.0.1", 500, 80);
         packet_1.meta_mut().set_overlay(true);
         packet_2.meta_mut().set_overlay(true);
-        let key_1 = FlowKey::try_from(net::flow_key::Uni(&packet_1)).unwrap();
-        let key_2 = FlowKey::try_from(net::flow_key::Uni(&packet_2)).unwrap();
+        let key_1 = FlowKey::try_from(&packet_1).unwrap();
+        let key_2 = FlowKey::try_from(&packet_2).unwrap();
         let input = vec![packet_1, packet_2];
         let out: Vec<_> = pipeline.process(input.into_iter()).collect();
         let packet_1 = &out[0];

flow-filter/src/lib.rs

@@ -369,7 +369,7 @@ impl FlowFilter {
             return;
         }
 
-        let Ok(flow_key) = FlowKey::try_from(net::flow_key::Uni(&*packet)) else {
+        let Ok(flow_key) = FlowKey::try_from(&*packet) else {
             return;
         };
 

flow-filter/src/tests.rs

@@ -12,7 +12,6 @@ use config::external::overlay::vpcpeering::{VpcExpose, VpcManifest, VpcPeering,
 use lpm::prefix::{L4Protocol, PortRange, Prefix, PrefixWithOptionalPorts};
 use net::FlowKey;
 use net::buffer::{PacketBufferMut, TestBuffer};
-use net::flow_key::Uni;
 use net::flows::{FlowInfo, FlowStatus};
 use net::headers::{Net, TryHeadersMut, TryIpMut};
 use net::ip::NextHeader;
@@ -191,7 +190,7 @@ fn fake_flow_session<Buf: PacketBufferMut>(
     set_port_fw_state: bool,
 ) {
     // build flow key
-    let flow_key = FlowKey::try_from(Uni(&*packet)).unwrap();
+    let flow_key = FlowKey::try_from(&*packet).unwrap();
 
     // Create flow_info with dst_vpcd and NAT info and attach it to the packet
     let flow_info = FlowInfo::new(flow_key, Instant::now() + Duration::from_secs(60));

nat/src/portfw/flow_state.rs

@@ -6,7 +6,6 @@
 #![allow(clippy::single_match_else)]
 
 use net::buffer::PacketBufferMut;
-use net::flow_key::Uni;
 use net::flows::{ExtractRef, FlowStatus};
 use net::ip::UnicastIpAddr;
 use net::packet::{Packet, VpcDiscriminant};
@@ -107,7 +106,7 @@ pub(crate) fn build_portfw_flow_keys<Buf: PacketBufferMut>(
     dst_vpcd: VpcDiscriminant, // destination VPC to forward to
 ) -> Result<(FlowKey, FlowKey), ()> {
     // Extract flow key for the current packet
-    let current_flow_key = FlowKey::try_from(Uni(&*packet)).map_err(|_| ())?;
+    let current_flow_key = FlowKey::try_from(&*packet).map_err(|_| ())?;
 
     // Retrieve initial flow key for the current packet (before any other NAT translation); if
     // we don't have the information, we didn't populate it because we don't need it and fall

nat/src/stateful/nf.rs

@@ -16,7 +16,7 @@ use crate::stateful::state::MasqueradeState;
 use concurrency::sync::{Arc, Weak};
 use flow_entry::flow_table::table::{FlowTable, FlowTableError};
 use net::buffer::PacketBufferMut;
-use net::flow_key::{IcmpProtoKey, Uni};
+use net::flow_key::IcmpProtoKey;
 use net::flows::{ExtractRef, FlowInfo};
 use net::headers::{TryIp, TryTcp};
 use net::ip::UnicastIpAddr;
@@ -375,7 +375,7 @@ impl StatefulNat {
 
         // Extract flow key for the current packet
         let current_flow_key =
-            FlowKey::try_from(Uni(&*packet)).map_err(|_| StatefulNatError::FlowKeyError)?;
+            FlowKey::try_from(&*packet).map_err(|_| StatefulNatError::FlowKeyError)?;
 
         // Retrieve initial flow key for the current packet (before any other NAT translation); if
         // we don't have the information, we didn't populate it because we don't need it and fall

nat/src/stateful/test.rs

@@ -18,7 +18,6 @@ use flow_entry::flow_table::{FlowLookup, FlowTable};
 use flow_filter::{FlowFilter, FlowFilterTable, FlowFilterTableWriter};
 use net::buffer::{PacketBufferMut, TestBuffer};
 use net::eth::mac::Mac;
-use net::flow_key::Uni;
 use net::flows::FlowStatus;
 use net::flows::flow_info_item::ExtractRef;
 use net::headers::TryTcpMut;
@@ -366,7 +365,7 @@ fn check_packet(
 }
 
 fn flow_lookup<Buf: PacketBufferMut>(flow_table: &FlowTable, packet: &mut Packet<Buf>) {
-    let flow_key = FlowKey::try_from(Uni(&*packet)).unwrap();
+    let flow_key = FlowKey::try_from(&*packet).unwrap();
     if let Some(flow_info) = flow_table.lookup(&flow_key) {
         packet.meta_mut().flow_info = Some(flow_info);
     }

net/src/flows/flow_key.rs

@@ -585,40 +585,37 @@ impl Hash for FlowKey {
     }
 }
 
-/// Wrapper to specify unidirectional `FlowKey` creation
-#[repr(transparent)]
-#[derive(Debug)]
-pub struct Uni<T>(pub T);
-
-fn flow_key_from_packet<Buf: PacketBufferMut>(packet: &Packet<Buf>) -> Option<FlowKey> {
-    let ip = packet.headers().try_ip()?;
-    let src_ip = ip.src_addr();
-    let dst_ip = ip.dst_addr();
-
-    let transport = packet.headers().try_transport()?;
-    let ip_proto_key = match transport {
-        Transport::Tcp(tcp) => IpProtoKey::Tcp(TcpProtoKey {
-            src_port: tcp.source(),
-            dst_port: tcp.destination(),
-        }),
-        Transport::Udp(udp) => IpProtoKey::Udp(UdpProtoKey {
-            src_port: udp.source(),
-            dst_port: udp.destination(),
-        }),
-        Transport::Icmp4(icmp) => IpProtoKey::Icmp(IcmpProtoKey::new_icmp_v4(packet, icmp)),
-        Transport::Icmp6(icmp) => IpProtoKey::Icmp(IcmpProtoKey::new_icmp_v6(packet, icmp)),
-        #[allow(unreachable_patterns)]
-        _ => return None,
-    };
+impl<Buf: PacketBufferMut> TryFrom<&Packet<Buf>> for FlowKey {
+    type Error = FlowKeyError;
+    fn try_from(packet: &Packet<Buf>) -> Result<Self, Self::Error> {
+        let ip = packet
+            .headers()
+            .try_ip()
+            .ok_or(FlowKeyError::NoFlowKeyData)?;
+        let src_ip = ip.src_addr();
+        let dst_ip = ip.dst_addr();
 
-    let src_vpcd = packet.meta().src_vpcd;
-    Some(FlowKey::new(src_vpcd, src_ip, dst_ip, ip_proto_key))
-}
+        let transport = packet
+            .headers()
+            .try_transport()
+            .ok_or(FlowKeyError::NoFlowKeyData)?;
+        let ip_proto_key = match transport {
+            Transport::Tcp(tcp) => IpProtoKey::Tcp(TcpProtoKey {
+                src_port: tcp.source(),
+                dst_port: tcp.destination(),
+            }),
+            Transport::Udp(udp) => IpProtoKey::Udp(UdpProtoKey {
+                src_port: udp.source(),
+                dst_port: udp.destination(),
+            }),
+            Transport::Icmp4(icmp) => IpProtoKey::Icmp(IcmpProtoKey::new_icmp_v4(packet, icmp)),
+            Transport::Icmp6(icmp) => IpProtoKey::Icmp(IcmpProtoKey::new_icmp_v6(packet, icmp)),
+            #[allow(unreachable_patterns)]
+            _ => return Err(FlowKeyError::NoFlowKeyData),
+        };
 
-impl<Buf: PacketBufferMut> TryFrom<Uni<&Packet<Buf>>> for FlowKey {
-    type Error = FlowKeyError;
-    fn try_from(packet: Uni<&Packet<Buf>>) -> Result<Self, Self::Error> {
-        flow_key_from_packet(packet.0).ok_or(FlowKeyError::NoFlowKeyData)
+        let src_vpcd = packet.meta().src_vpcd;
+        Ok(FlowKey::new(src_vpcd, src_ip, dst_ip, ip_proto_key))
     }
 }
 
@@ -1078,7 +1075,7 @@ mod tests {
             .with_generator(FlowKeyAndPacket)
             .for_each(|(flow_key, packet)| match flow_key {
                 Some(_) => {
-                    let gen_flow_key = FlowKey::try_from(Uni(packet)).unwrap();
+                    let gen_flow_key = FlowKey::try_from(packet).unwrap();
                     assert_eq!(
                         gen_flow_key,
                         flow_key.unwrap(),
@@ -1087,7 +1084,7 @@ mod tests {
                     );
                 }
                 None => {
-                    assert!(FlowKey::try_from(Uni(packet)).is_err());
+                    assert!(FlowKey::try_from(packet).is_err());
                 }
             });
     }