feat(s2n-quic-dc): enable cbpf for packet filtering

dc/s2n-quic-dc/src/psk/io.rs

@@ -13,6 +13,7 @@ use s2n_quic::{
     server::Name,
 };
 use s2n_quic_core::inet::SocketAddress;
+use s2n_quic_platform::syscall;
 use std::{
     hash::BuildHasher,
     io,
@@ -58,8 +59,12 @@ impl Server {
         subscriber: Subscriber,
         builder: server::Builder<Event>,
     ) -> Result<Self, Error> {
+        let socket_for_client_hello_packets = syscall::bind_udp(addr, false, true, false).unwrap();
+        let socket_for_other_packets = syscall::bind_udp(addr, false, true, false).unwrap();
+
         let io = s2n_quic::provider::io::default::Builder::default()
-            .with_receive_address(addr)?
+            .with_rx_socket(socket_for_client_hello_packets.into())?
+            .with_rx_socket(socket_for_other_packets.into())?
             .with_base_mtu(DEFAULT_BASE_MTU.min(builder.mtu))?
             .with_initial_mtu(builder.mtu)?
             .with_max_mtu(builder.mtu)?

quic/s2n-quic-platform/Cargo.toml

@@ -46,6 +46,11 @@ tempfile = "3"
 tokio = { version = "1", features = ["full"] }
 tracing = { version = "0.1" }
 
+# Tests use cbpf are only available on Linux X86_64 platform
+[target.'cfg(all(target_os = "linux", not(target_arch = "aarch64")))'.dev-dependencies]
+bytes = "1"
+s2n-quic = { path = "../s2n-quic" }
+
 [package.metadata.kani]
 flags = { tests = true }
 unstable = { stubbing = true }

quic/s2n-quic-platform/src/io/tokio.rs

@@ -14,6 +14,9 @@ use s2n_quic_core::{
 use std::{convert::TryInto, io, io::ErrorKind};
 use tokio::runtime::Handle;
 
+#[cfg(all(target_os = "linux", not(target_arch = "aarch64")))]
+use crate::bpf::cbpf::{abs, and, jeq, ldb, ret, Program};
+
 mod builder;
 mod clock;
 pub(crate) mod task;
@@ -29,6 +32,25 @@ pub struct Io {
     builder: Builder,
 }
 
+/// cBPF program to route QUIC packets across multiple sockets.
+/// Routes Initial packets with DCID length = 8 to socket 0, all other packets to socket 1.
+#[cfg(all(target_os = "linux", not(target_arch = "aarch64")))]
+pub static ROUTER: Program = Program::new(&[
+    // Load byte 0 and check if it's an Initial packet (first 4 bits = 1100)
+    ldb(abs(0)),
+    and(0b1111_0000), // Mask the last four bits of the first byte. The first four bits can confirm if the packet is a INITIAL packet.
+    // If Initial packet, continue; else jump to ret(1)
+    jeq(0b1100_0000, 0, 3), // First four bits of INITIAL packet should be 1100.
+    // Load byte 5 (DCID length) and check if it equals 8
+    ldb(abs(5)),
+    // If DCID len = 8, continue to ret(0); else jump to ret(1)
+    jeq(0x08, 0, 1),
+    // Return 0: socket 0 handles Initial packets with DCID length = 8
+    ret(0),
+    // Return 1: socket 1 handles all other packets
+    ret(1),
+]);
+
 impl Io {
     pub fn builder() -> Builder {
         Builder::default()
@@ -46,7 +68,7 @@ impl Io {
     ) -> io::Result<(tokio::task::JoinHandle<()>, SocketAddress)> {
         let Builder {
             handle,
-            rx_socket,
+            rx_sockets,
             tx_socket,
             recv_addr,
             send_addr,
@@ -89,35 +111,63 @@ impl Io {
 
         let guard = handle.enter();
 
-        let rx_socket = if let Some(rx_socket) = rx_socket {
-            rx_socket
+        // Track whether sockets were explicitly provided by the user (for ROUTER attachment)
+        let user_provided_sockets = !rx_sockets.is_empty();
+
+        // Build the list of rx sockets - either from provided sockets or create from recv_addr
+        let rx_socket_list = if user_provided_sockets {
+            rx_sockets
         } else if let Some(recv_addr) = recv_addr {
-            syscall::bind_udp(recv_addr, reuse_address, reuse_port, only_v6)?
+            // Check env var for number of sockets to create (unstable feature)
+            let rx_socket_count: usize =
+                parse_env("S2N_QUIC_UNSTABLE_RX_SOCKET_COUNT").unwrap_or(1);
+            let mut sockets = Vec::with_capacity(rx_socket_count);
+            for _ in 0..rx_socket_count {
+                sockets.push(syscall::bind_udp(
+                    recv_addr,
+                    reuse_address,
+                    reuse_port,
+                    only_v6,
+                )?);
+            }
+            sockets
         } else {
             return Err(io::Error::new(
                 io::ErrorKind::InvalidInput,
                 "missing bind address",
             ));
         };
 
-        let rx_addr = convert_addr_to_std(rx_socket.local_addr()?)?;
+        // Get the address from the first socket
+        let rx_addr = convert_addr_to_std(rx_socket_list[0].local_addr()?)?;
+
+        // Only attach ROUTER if user explicitly provided multiple sockets (Linux x86_64 only)
+        // ROUTER is not used for sockets created via S2N_QUIC_UNSTABLE_RX_SOCKET_COUNT or there is only one socket in place
+        #[cfg(all(target_os = "linux", not(target_arch = "aarch64")))]
+        if user_provided_sockets && rx_socket_list.len() > 1 {
+            for socket in &rx_socket_list {
+                ROUTER.attach(socket)?;
+            }
+        }
 
         let tx_socket = if let Some(tx_socket) = tx_socket {
             tx_socket
         } else if let Some(send_addr) = send_addr {
             syscall::bind_udp(send_addr, reuse_address, reuse_port, only_v6)?
         } else {
             // No tx_socket or send address was specified, so the tx socket
-            // will be a handle to the rx socket.
-            rx_socket.try_clone()?
+            // will be a handle to the first rx socket.
+            rx_socket_list[0].try_clone()?
         };
 
         if let Some(size) = socket_send_buffer_size {
             tx_socket.set_send_buffer_size(size)?;
         }
 
         if let Some(size) = socket_recv_buffer_size {
-            rx_socket.set_recv_buffer_size(size)?;
+            for socket in &rx_socket_list {
+                socket.set_recv_buffer_size(size)?;
+            }
         }
 
         let mut mtu_config = mtu_config_builder
@@ -149,20 +199,21 @@ impl Io {
             },
         });
 
-        // Configure the socket with GRO
-        let gro_enabled = gro_enabled.unwrap_or(true) && syscall::configure_gro(&rx_socket);
+        // Configure the socket with GRO (check first socket, apply to all)
+        let gro_enabled = gro_enabled.unwrap_or(true) && syscall::configure_gro(&rx_socket_list[0]);
 
         publisher.on_platform_feature_configured(event::builder::PlatformFeatureConfigured {
             configuration: event::builder::PlatformFeatureConfiguration::Gro {
                 enabled: gro_enabled,
             },
         });
 
-        // Configure packet info CMSG
-        syscall::configure_pktinfo(&rx_socket);
-
-        // Configure TOS/ECN
-        let tos_enabled = syscall::configure_tos(&rx_socket);
+        // Configure packet info CMSG and TOS/ECN for all rx sockets
+        let mut tos_enabled = false;
+        for socket in &rx_socket_list {
+            syscall::configure_pktinfo(socket);
+            tos_enabled = syscall::configure_tos(socket);
+        }
 
         publisher.on_platform_feature_configured(event::builder::PlatformFeatureConfigured {
             configuration: event::builder::PlatformFeatureConfiguration::Ecn {
@@ -193,34 +244,22 @@ impl Io {
 
             let mut consumers = vec![];
 
-            let rx_socket_count = parse_env("S2N_QUIC_UNSTABLE_RX_SOCKET_COUNT").unwrap_or(1);
-
             // configure the number of self-wakes before "cooling down" and waiting for epoll to
             // complete
             let rx_cooldown = cooldown("RX");
 
-            for idx in 0usize..rx_socket_count {
+            for socket in rx_socket_list {
                 let (producer, consumer) = socket::ring::pair(entries, payload_len);
                 consumers.push(consumer);
 
                 // spawn a task that actually reads from the socket into the ring buffer
-                if idx + 1 == rx_socket_count {
-                    handle.spawn(task::rx(
-                        rx_socket,
-                        producer,
-                        rx_cooldown,
-                        stats_sender.clone(),
-                    ));
-                    break;
-                } else {
-                    let rx_socket = rx_socket.try_clone()?;
-                    handle.spawn(task::rx(
-                        rx_socket,
-                        producer,
-                        rx_cooldown.clone(),
-                        stats_sender.clone(),
-                    ));
-                }
+                let rx_socket = socket.try_clone()?;
+                handle.spawn(task::rx(
+                    rx_socket,
+                    producer,
+                    rx_cooldown.clone(),
+                    stats_sender.clone(),
+                ));
             }
 
             // construct the RX side for the endpoint event loop

quic/s2n-quic-platform/src/io/tokio/builder.rs

@@ -6,7 +6,7 @@ use super::*;
 #[derive(Debug, Default)]
 pub struct Builder {
     pub(super) handle: Option<Handle>,
-    pub(super) rx_socket: Option<socket2::Socket>,
+    pub(super) rx_sockets: Vec<socket2::Socket>,
     pub(super) tx_socket: Option<socket2::Socket>,
     pub(super) recv_addr: Option<std::net::SocketAddr>,
     pub(super) send_addr: Option<std::net::SocketAddr>,
@@ -34,7 +34,7 @@ impl Builder {
     ///
     /// NOTE: this method is mutually exclusive with `with_rx_socket`
     pub fn with_receive_address(mut self, addr: std::net::SocketAddr) -> io::Result<Self> {
-        debug_assert!(self.rx_socket.is_none(), "rx socket has already been set");
+        debug_assert!(self.rx_sockets.is_empty(), "rx socket has already been set");
         self.recv_addr = Some(addr);
         Ok(self)
     }
@@ -49,16 +49,17 @@ impl Builder {
         Ok(self)
     }
 
-    /// Sets the socket used for receiving for the runtime. If no tx_socket or send address is
-    /// specified, this socket will be used for transmitting.
+    /// Adds a socket to be used for receiving packets. This method can be called multiple times
+    /// to add multiple sockets for packet filtering/routing scenarios.
+    /// If no tx_socket or send address is specified, the first socket will be used for transmitting.
     ///
     /// NOTE: this method is mutually exclusive with `with_receive_address`
     pub fn with_rx_socket(mut self, socket: std::net::UdpSocket) -> io::Result<Self> {
         debug_assert!(
             self.recv_addr.is_none(),
             "recv address has already been set"
         );
-        self.rx_socket = Some(socket.into());
+        self.rx_sockets.push(socket.into());
         Ok(self)
     }