tokio-quiche: support per-connection DSCP marking

Add optional initial DSCP marking for server connections and a handle to update it as the connection progresses. Apply the marking to outgoing IPv4 and IPv6 packets, including GSO sends and stateless replies.

Cover per-connection isolation, IPv4-mapped addresses, and IPv6 in tests.

Author: inamdarminaz

tokio-quiche/src/quic/connection/mod.rs

@@ -67,6 +67,7 @@ use super::io::worker::IoWorkerParams;
 use super::io::worker::Running;
 use super::io::worker::RunningOrClosing;
 use super::io::worker::WriteState;
+use super::DscpHandle;
 use super::QuicheConnection;
 use crate::metrics::Metrics;
 use crate::quic::io::worker::IoWorker;
@@ -281,6 +282,15 @@ where
         (*self.params.quiche_conn).as_mut()
     }
 
+    /// Returns the outgoing DSCP handle when server marking is configured.
+    ///
+    /// Clone it before starting the handshake to update the worker's marking
+    /// from a certificate-selection callback. It is initialized with the
+    /// listener's pre-handshake value.
+    pub fn dscp_handle(&self) -> Option<&DscpHandle> {
+        self.params.dscp_handle.as_ref()
+    }
+
     /// A handle to the [`QuicAuditStats`] for this connection.
     ///
     /// # Note
@@ -337,6 +347,7 @@ where
             cfg: self.params.writer_cfg,
             audit_log_stats: self.audit_log_stats,
             write_state: WriteState::default(),
+            dscp_handle: self.params.dscp_handle,
             conn_map_cmd_tx: self.params.conn_map_cmd_tx,
             cid_generator: self.params.cid_generator,
             #[cfg(feature = "perf-quic-listener-metrics")]
@@ -468,6 +479,7 @@ where
 {
     pub writer_cfg: WriterConfig,
     pub initial_pkt: Option<Incoming>,
+    pub dscp_handle: Option<DscpHandle>,
     pub shutdown_tx: mpsc::Sender<()>,
     pub conn_map_cmd_tx: mpsc::UnboundedSender<ConnectionMapCommand>, /* channel that signals connection map changes */
     pub scid: ConnectionId<'static>,

tokio-quiche/src/quic/dscp.rs

@@ -0,0 +1,112 @@
+// Copyright (C) 2025, Cloudflare, Inc.
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+//     * Redistributions of source code must retain the above copyright notice,
+//       this list of conditions and the following disclaimer.
+//     * Redistributions in binary form must reproduce the above copyright
+//       notice, this list of conditions and the following disclaimer in the
+//       documentation and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+
+use std::sync::atomic::AtomicU8;
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
+
+const UNMARKED: u8 = 64;
+
+/// A six-bit Differentiated Services Code Point for outgoing UDP packets.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct Dscp(u8);
+
+impl Dscp {
+    /// Creates a DSCP value, rejecting values outside the six-bit range.
+    pub const fn new(value: u8) -> Option<Self> {
+        if value < UNMARKED {
+            Some(Self(value))
+        } else {
+            None
+        }
+    }
+
+    /// Returns the six-bit DSCP value.
+    pub const fn value(self) -> u8 {
+        self.0
+    }
+
+    #[cfg(any(all(target_os = "linux", not(feature = "fuzzing")), test))]
+    pub(crate) const fn tos(self) -> u8 {
+        self.0 << 2
+    }
+}
+
+/// Per-connection DSCP shared with the QUIC I/O worker.
+///
+/// Clone this handle before starting the handshake and call [`Self::set`] once
+/// the connection's final marking is known. `Dscp::new(0)` explicitly
+/// clears an earlier marking even when the shared socket has a nonzero default;
+/// `None` omits the per-packet control message.
+#[derive(Clone, Debug)]
+pub struct DscpHandle(Arc<AtomicU8>);
+
+impl DscpHandle {
+    pub(crate) fn new(initial: Dscp) -> Self {
+        Self(Arc::new(AtomicU8::new(initial.value())))
+    }
+
+    /// Changes the DSCP used for subsequent packets on this connection.
+    pub fn set(&self, dscp: Option<Dscp>) {
+        // This atomic only holds the marking byte; it synchronizes no other data.
+        self.0
+            .store(dscp.map_or(UNMARKED, Dscp::value), Ordering::Relaxed);
+    }
+
+    /// Returns the currently selected marking, or `None` if disabled.
+    pub fn get(&self) -> Option<Dscp> {
+        Dscp::new(self.0.load(Ordering::Relaxed))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn dscp_is_six_bits() {
+        assert_eq!(Dscp::new(0).map(Dscp::tos), Some(0));
+        assert_eq!(Dscp::new(34).map(Dscp::tos), Some(136));
+        assert_eq!(Dscp::new(63).map(Dscp::tos), Some(252));
+        assert_eq!(Dscp::new(64), None);
+        assert_eq!(Dscp::new(255), None);
+    }
+
+    #[test]
+    fn handles_share_only_their_own_connection() {
+        let first = DscpHandle::new(Dscp::new(34).unwrap());
+        let cloned = first.clone();
+        let second = DscpHandle::new(Dscp::new(1).unwrap());
+
+        cloned.set(Dscp::new(28));
+        assert_eq!(first.get(), Dscp::new(28));
+        assert_eq!(second.get(), Dscp::new(1));
+
+        cloned.set(Dscp::new(0));
+        assert_eq!(first.get(), Dscp::new(0));
+        cloned.set(None);
+        assert_eq!(first.get(), None);
+        assert_eq!(second.get(), Dscp::new(1));
+    }
+}