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

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

@@ -13,6 +13,9 @@ use s2n_quic::{
     server::Name,
 };
 use s2n_quic_core::inet::SocketAddress;
+#[cfg(target_os = "linux")]
+use s2n_quic_platform::bpf::cbpf::{abs, and, jeq, ldb, ret, Program};
+use s2n_quic_platform::syscall;
 use std::{
     hash::BuildHasher,
     io,
@@ -38,6 +41,25 @@ const DEFAULT_INITIAL_RTT: Duration = Duration::from_millis(1);
 
 const BUFFER_SIZE: usize = 16 * 1024;
 
+/// 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(target_os = "linux")]
+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),
+]);
+
 pub type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 
 pub type Result<T = (), E = Error> = core::result::Result<T, E>;
@@ -58,8 +80,30 @@ impl Server {
         subscriber: Subscriber,
         builder: server::Builder<Event>,
     ) -> Result<Self, Error> {
+        let socket_for_client_hello_packets = syscall::bind_udp(addr, false, false, false)?;
+
+        // Acquire the bound address with a port assigned
+        let bound_addr = socket_for_client_hello_packets
+            .local_addr()?
+            .as_socket()
+            .unwrap();
+
+        socket_for_client_hello_packets
+            .set_reuse_port(true)
+            .unwrap();
+
+        let socket_for_other_packets = syscall::bind_udp(bound_addr, false, true, false)?;
+
+        // Attach ROUTER to both sockets for packet filtering
+        #[cfg(target_os = "linux")]
+        {
+            ROUTER.attach(&socket_for_client_hello_packets)?;
+            ROUTER.attach(&socket_for_other_packets)?;
+        }
+
         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)?
@@ -573,6 +617,7 @@ mod tests {
     };
     use s2n_quic_core::time::StdClock;
     use std::time::Instant;
+    use tokio::net::UdpSocket;
     use tokio_util::sync::DropGuard;
 
     /// A test limiter that closes all incoming connections immediately
@@ -739,4 +784,152 @@ mod tests {
             duration
         );
     }
+
+    // Tests that the ROUTER cBPF filter correctly routes packets to the appropriate socket.
+    #[cfg(target_os = "linux")]
+    #[tokio::test]
+    #[cfg_attr(miri, ignore)]
+    async fn router_cbpf_packet_filtering_test() -> io::Result<()> {
+        static IPV4_LOCALHOST: &str = "127.0.0.1:0";
+
+        // Create two rx sockets bound to same port with SO_REUSEPORT
+        let rx_socket_0 = syscall::bind_udp(IPV4_LOCALHOST, false, false, false)?;
+        rx_socket_0.set_nonblocking(true)?;
+        let port = rx_socket_0.local_addr()?.as_socket().unwrap().port();
+        rx_socket_0.set_reuse_port(true)?;
+
+        let rx_socket_1 = syscall::bind_udp(("127.0.0.1", port), false, true, false)?;
+        rx_socket_1.set_nonblocking(true)?;
+
+        // Attach ROUTER to both sockets
+        ROUTER.attach(&rx_socket_0)?;
+        ROUTER.attach(&rx_socket_1)?;
+
+        // Convert to tokio sockets for async recv
+        let rx_socket_0 = UdpSocket::from_std(rx_socket_0.into())?;
+        let rx_socket_1 = UdpSocket::from_std(rx_socket_1.into())?;
+
+        // Create sender socket
+        let sender = UdpSocket::bind("127.0.0.1:0").await?;
+        let target_addr: std::net::SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap();
+
+        // Test packet A: Initial packet with DCID length = 8
+        // Should route to socket 0
+        // Format: [header byte, version (4 bytes), dcid_len, ...]
+        let packet_a = {
+            let mut p = vec![0u8; 32];
+            p[0] = 0xC0; // Initial packet (first 4 bits = 1100)
+            p[1..5].copy_from_slice(&[0x00, 0x00, 0x00, 0x01]); // version
+            p[5] = 0x08; // DCID length = 8
+            p
+        };
+
+        // Test packet B: Handshake packet
+        // Should route to socket 1
+        let packet_b = {
+            let mut p = vec![0u8; 32];
+            p[0] = 0xE0; // Handshake packet first four bits are 1110
+            p
+        };
+
+        // Test packet C: Initial packet but DCID length != 8
+        // Should route to socket 1
+        let packet_c = {
+            let mut p = vec![0u8; 32];
+            p[0] = 0xC0; // Initial packet (first 4 bits = 1100)
+            p[1..5].copy_from_slice(&[0x00, 0x00, 0x00, 0x01]); // version
+            p[5] = 0x10; // DCID length = 16
+            p
+        };
+
+        // Send packets
+        sender.send_to(&packet_a, target_addr).await?;
+        sender.send_to(&packet_b, target_addr).await?;
+        sender.send_to(&packet_c, target_addr).await?;
+
+        // Receive and verify routing
+        let mut buf_socket0 = [0u8; 1024];
+        let mut buf_packet1_socket1 = [0u8; 1024];
+        let mut buf_packet2_socket1 = [0u8; 1024];
+
+        // Socket 0 should receive packet_a (Initial with DCID len = 8)
+        let recv_result = tokio::time::timeout(
+            Duration::from_millis(500),
+            rx_socket_0.recv_from(&mut buf_socket0),
+        )
+        .await;
+        assert!(
+            recv_result.is_ok(),
+            "Socket 0 should receive packet_a (Initial with DCID len=8)"
+        );
+        let (len, _) = recv_result.unwrap()?;
+        assert_eq!(
+            buf_socket0[0], 0xC0,
+            "Socket 0 should receive Initial packet"
+        );
+        assert_eq!(
+            buf_socket0[5], 0x08,
+            "Socket 0 should receive packet with DCID len=8"
+        );
+        assert_eq!(len, 32);
+
+        // Socket 1 should receive packet_b and packet_c
+        let recv_result = tokio::time::timeout(
+            Duration::from_millis(500),
+            rx_socket_1.recv_from(&mut buf_packet1_socket1),
+        )
+        .await;
+        assert!(
+            recv_result.is_ok(),
+            "Socket 1 should receive packet_b or packet_c"
+        );
+        let (len1, _) = recv_result.unwrap()?;
+        assert_eq!(len1, 32);
+
+        let recv_result = tokio::time::timeout(
+            Duration::from_millis(500),
+            rx_socket_1.recv_from(&mut buf_packet2_socket1),
+        )
+        .await;
+        assert!(
+            recv_result.is_ok(),
+            "Socket 1 should receive packet_b or packet_c"
+        );
+        let (len2, _) = recv_result.unwrap()?;
+        assert_eq!(len2, 32);
+
+        // Verify that socket 1 received exactly packet_b and packet_c in either order
+        let received_packets = [&buf_packet1_socket1[..32], &buf_packet2_socket1[..32]];
+
+        // Check that one packet matches packet_b (Handshake: header = 0xE0)
+        let has_packet_b = received_packets
+            .iter()
+            .any(|p| p[0] == 0xE0 && p[..] == packet_b[..]);
+        assert!(
+            has_packet_b,
+            "Socket 1 should receive packet_b (Handshake packet with header 0xE0)"
+        );
+
+        // Check that one packet matches packet_c (Initial with DCID len = 16)
+        let has_packet_c = received_packets
+            .iter()
+            .any(|p| p[0] == 0xC0 && p[5] == 0x10 && p[..] == packet_c[..]);
+        assert!(
+            has_packet_c,
+            "Socket 1 should receive packet_c (Initial packet with DCID len=16)"
+        );
+
+        // Socket 0 should not have any more packets
+        let recv_result = tokio::time::timeout(
+            Duration::from_millis(100),
+            rx_socket_0.recv_from(&mut buf_socket0),
+        )
+        .await;
+        assert!(
+            recv_result.is_err(),
+            "Socket 0 should not receive any more packets"
+        );
+
+        Ok(())
+    }
 }

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

@@ -46,7 +46,7 @@ impl Io {
     ) -> io::Result<(tokio::task::JoinHandle<()>, SocketAddress)> {
         let Builder {
             handle,
-            rx_socket,
+            rx_sockets,
             tx_socket,
             recv_addr,
             send_addr,
@@ -89,35 +89,51 @@ impl Io {
 
         let guard = handle.enter();
 
-        let rx_socket = if let Some(rx_socket) = rx_socket {
-            rx_socket
+        // Build the list of rx sockets - either from provided sockets or create from recv_addr
+        let rx_socket_list = if !rx_sockets.is_empty() {
+            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()?)?;
 
         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
@@ -150,19 +166,20 @@ impl Io {
         });
 
         // Configure the socket with GRO
-        let gro_enabled = gro_enabled.unwrap_or(true) && syscall::configure_gro(&rx_socket);
+        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 +210,21 @@ 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(),
-                    ));
-                }
+                handle.spawn(task::rx(
+                    socket,
+                    producer,
+                    rx_cooldown.clone(),
+                    stats_sender.clone(),
+                ));
             }
 
             // construct the RX side for the endpoint event loop