feat(s2n-quic-dc): implement dcQUIC client Tokio Builder (#2741)

boquan-fang · web-flow · commit 3a7aba323722 · 2025-08-06T09:46:47.000-07:00
diff --git a/dc/s2n-quic-dc/src/stream/client/tokio.rs b/dc/s2n-quic-dc/src/stream/client/tokio.rs
@@ -8,19 +8,331 @@ use crate::{
     path::secret,
     stream::{
         application::Stream,
+        client::{rpc as rpc_internal, tokio as client},
         endpoint,
         environment::{
             tokio::{self as env, Environment},
-            Environment as _,
+            udp as udp_pool, Environment as _,
         },
-        recv,
-        socket::Protocol,
+        recv, socket,
     },
 };
 use s2n_quic_core::time::Clock;
 use std::{io, net::SocketAddr, time::Duration};
 use tokio::net::TcpStream;
 
+pub mod rpc {
+    pub use crate::stream::client::rpc::{InMemoryResponse, Request, Response};
+}
+
+// This trait is a temporary solution to abstract handshake_with_entry
+// and local_addr until we implement the handshake provider
+#[allow(async_fn_in_trait)]
+pub trait Handshake: AsRef<secret::Map> + Clone {
+    /// Handshake with the remote peer
+    async fn handshake_with_entry(
+        &self,
+        remote_handshake_addr: SocketAddr,
+    ) -> std::io::Result<(secret::map::Peer, secret::HandshakeKind)>;
+
+    fn local_addr(&self) -> std::io::Result<SocketAddr>;
+}
+
+#[derive(Clone)]
+pub struct Client<H: Handshake + Clone, S: event::Subscriber + Clone> {
+    env: Environment<S>,
+    handshake: H,
+    default_protocol: socket::Protocol,
+    linger: Option<Duration>,
+}
+
+impl<H: Handshake + Clone, S: event::Subscriber + Clone> Client<H, S> {
+    #[inline]
+    pub fn new(handshake: H, subscriber: S) -> io::Result<Self> {
+        Self::builder().build(handshake, subscriber)
+    }
+
+    #[inline]
+    pub fn builder() -> Builder {
+        Builder::default()
+    }
+
+    pub fn drop_state(&self) {
+        self.handshake.as_ref().drop_state()
+    }
+
+    pub fn handshake_state(&self) -> &H {
+        &self.handshake
+    }
+
+    #[inline]
+    pub async fn handshake_with(
+        &self,
+        remote_handshake_addr: SocketAddr,
+    ) -> io::Result<secret::HandshakeKind> {
+        let (_peer, kind) = self
+            .handshake
+            .handshake_with_entry(remote_handshake_addr)
+            .await?;
+        Ok(kind)
+    }
+
+    #[inline]
+    async fn handshake_for_connect(
+        &self,
+        remote_handshake_addr: SocketAddr,
+    ) -> io::Result<secret::map::Peer> {
+        let (peer, _kind) = self
+            .handshake
+            .handshake_with_entry(remote_handshake_addr)
+            .await?;
+        Ok(peer)
+    }
+
+    /// Connects using the preferred protocol
+    #[inline]
+    pub async fn connect(
+        &self,
+        handshake_addr: SocketAddr,
+        acceptor_addr: SocketAddr,
+    ) -> io::Result<Stream<S>> {
+        match self.default_protocol {
+            socket::Protocol::Udp => self.connect_udp(handshake_addr, acceptor_addr).await,
+            socket::Protocol::Tcp => self.connect_tcp(handshake_addr, acceptor_addr).await,
+            protocol => Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!("invalid default protocol {protocol:?}"),
+            )),
+        }
+    }
+
+    /// Makes an RPC request using the preferred protocol
+    pub async fn rpc<Req, Res>(
+        &self,
+        handshake_addr: SocketAddr,
+        acceptor_addr: SocketAddr,
+        request: Req,
+        response: Res,
+    ) -> io::Result<Res::Output>
+    where
+        Req: rpc::Request,
+        Res: rpc::Response,
+    {
+        match self.default_protocol {
+            socket::Protocol::Udp => {
+                self.rpc_udp(handshake_addr, acceptor_addr, request, response)
+                    .await
+            }
+            socket::Protocol::Tcp => {
+                self.rpc_tcp(handshake_addr, acceptor_addr, request, response)
+                    .await
+            }
+            protocol => Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!("invalid default protocol {protocol:?}"),
+            )),
+        }
+    }
+
+    /// Connects using the UDP transport layer
+    #[inline]
+    pub async fn connect_udp(
+        &self,
+        handshake_addr: SocketAddr,
+        acceptor_addr: SocketAddr,
+    ) -> io::Result<Stream<S>> {
+        // ensure we have a secret for the peer
+        let handshake = self.handshake_for_connect(handshake_addr);
+
+        let mut stream = client::connect_udp(handshake, acceptor_addr, &self.env).await?;
+        Self::write_prelude(&mut stream).await?;
+        Ok(stream)
+    }
+
+    /// Makes an RPC request using the UDP transport layer
+    #[inline]
+    pub async fn rpc_udp<Req, Res>(
+        &self,
+        handshake_addr: SocketAddr,
+        acceptor_addr: SocketAddr,
+        request: Req,
+        response: Res,
+    ) -> io::Result<Res::Output>
+    where
+        Req: rpc::Request,
+        Res: rpc::Response,
+    {
+        // ensure we have a secret for the peer
+        let handshake = self.handshake_for_connect(handshake_addr);
+
+        let stream = client::connect_udp(handshake, acceptor_addr, &self.env).await?;
+        rpc_internal::from_stream(stream, request, response).await
+    }
+
+    /// Connects using the TCP transport layer
+    #[inline]
+    pub async fn connect_tcp(
+        &self,
+        handshake_addr: SocketAddr,
+        acceptor_addr: SocketAddr,
+    ) -> io::Result<Stream<S>> {
+        // ensure we have a secret for the peer
+        let handshake = self.handshake_for_connect(handshake_addr);
+
+        let mut stream =
+            client::connect_tcp(handshake, acceptor_addr, &self.env, self.linger).await?;
+        Self::write_prelude(&mut stream).await?;
+        Ok(stream)
+    }
+
+    /// Makes an RPC request using the TCP transport layer
+    #[inline]
+    pub async fn rpc_tcp<Req, Res>(
+        &self,
+        handshake_addr: SocketAddr,
+        acceptor_addr: SocketAddr,
+        request: Req,
+        response: Res,
+    ) -> io::Result<Res::Output>
+    where
+        Req: rpc::Request,
+        Res: rpc::Response,
+    {
+        // ensure we have a secret for the peer
+        let handshake = self.handshake_for_connect(handshake_addr);
+
+        let stream = client::connect_tcp(handshake, acceptor_addr, &self.env, self.linger).await?;
+        rpc_internal::from_stream(stream, request, response).await
+    }
+
+    /// Connects with a pre-existing TCP stream
+    #[inline]
+    pub async fn connect_tcp_with(
+        &self,
+        handshake_addr: SocketAddr,
+        stream: TcpStream,
+    ) -> io::Result<Stream<S>> {
+        // ensure we have a secret for the peer
+        let handshake = self.handshake_for_connect(handshake_addr).await?;
+
+        let mut stream = client::connect_tcp_with(handshake, stream, &self.env).await?;
+        Self::write_prelude(&mut stream).await?;
+        Ok(stream)
+    }
+
+    #[inline]
+    async fn write_prelude(stream: &mut Stream<S>) -> io::Result<()> {
+        // TODO should we actually write the prelude here or should we do late sealer binding on
+        // the first packet to reduce secret reordering on the peer
+
+        stream
+            .write_from(&mut s2n_quic_core::buffer::reader::storage::Empty)
+            .await
+            .map(|_| ())
+    }
+}
+
+#[derive(Default)]
+pub struct Builder {
+    default_protocol: Option<socket::Protocol>,
+    background_threads: Option<usize>,
+    linger: Option<Duration>,
+    send_buffer: Option<usize>,
+    recv_buffer: Option<usize>,
+}
+
+impl Builder {
+    pub fn with_tcp(self, enabled: bool) -> Self {
+        self.with_default_protocol(if enabled {
+            socket::Protocol::Tcp
+        } else {
+            socket::Protocol::Udp
+        })
+    }
+
+    pub fn with_udp(self, enabled: bool) -> Self {
+        self.with_default_protocol(if enabled {
+            socket::Protocol::Udp
+        } else {
+            socket::Protocol::Tcp
+        })
+    }
+
+    pub fn with_default_protocol(mut self, protocol: socket::Protocol) -> Self {
+        self.default_protocol = Some(protocol);
+        self
+    }
+
+    pub fn with_background_threads(mut self, threads: usize) -> Self {
+        self.background_threads = Some(threads);
+        self
+    }
+
+    pub fn with_linger(mut self, linger: Duration) -> Self {
+        self.linger = Some(linger);
+        self
+    }
+
+    /// Sets the send buffer for the OS socket handle.
+    ///
+    /// See `SO_SNDBUF` for more information.
+    ///
+    /// Note that this only applies to sockets that are created by SaltyLib. Any sockets
+    /// provided by the application will not inherit this value.
+    pub fn with_send_buffer(mut self, bytes: usize) -> Self {
+        self.send_buffer = Some(bytes);
+        self
+    }
+
+    /// Sets the recv buffer for the OS socket handle.
+    ///
+    /// See `SO_RCVBUF` for more information.
+    ///
+    /// Note that this only applies to sockets that are created by SaltyLib. Any sockets
+    /// provided by the application will not inherit this value.
+    pub fn with_recv_buffer(mut self, bytes: usize) -> Self {
+        self.recv_buffer = Some(bytes);
+        self
+    }
+
+    #[inline]
+    pub fn build<H: Handshake + Clone, S: event::Subscriber + Clone>(
+        self,
+        handshake: H,
+        subscriber: S,
+    ) -> io::Result<Client<H, S>> {
+        // bind the sockets to the same address family as the handshake
+        let mut local_addr = handshake.local_addr()?;
+        local_addr.set_port(0);
+        let mut options = socket::Options::new(local_addr);
+
+        options.send_buffer = self.send_buffer;
+        options.recv_buffer = self.recv_buffer;
+
+        let mut env = env::Builder::new(subscriber).with_socket_options(options);
+
+        let pool = udp_pool::Config::new((handshake.as_ref()).clone());
+        env = env.with_pool(pool);
+
+        if let Some(threads) = self.background_threads {
+            env = env.with_threads(threads);
+        }
+        let env = env.build()?;
+
+        // default to UDP
+        let default_protocol = self.default_protocol.unwrap_or(socket::Protocol::Udp);
+
+        let linger = self.linger;
+
+        Ok(Client {
+            env,
+            handshake,
+            default_protocol,
+            linger,
+        })
+    }
+}
+
 /// Connects using the UDP transport layer
 ///
 /// Callers should send data immediately after calling this to ensure minimal
@@ -52,7 +364,7 @@ where
     // build the stream inside the application context
     let stream = stream.connect()?;
 
-    debug_assert_eq!(stream.protocol(), Protocol::Udp);
+    debug_assert_eq!(stream.protocol(), socket::Protocol::Udp);
 
     Ok(stream)
 }
@@ -190,7 +502,7 @@ where
     // build the stream inside the application context
     let stream = stream.connect()?;
 
-    debug_assert_eq!(stream.protocol(), Protocol::Tcp);
+    debug_assert_eq!(stream.protocol(), socket::Protocol::Tcp);
 
     Ok(stream)
 }
@@ -229,7 +541,7 @@ where
     // build the stream inside the application context
     let stream = stream.connect()?;
 
-    debug_assert_eq!(stream.protocol(), Protocol::Tcp);
+    debug_assert_eq!(stream.protocol(), socket::Protocol::Tcp);
 
     Ok(stream)
 }
