refactor(s2n-quic-dc): Add trait for TCP WorkerState poll behavior (#2836)

Author: mehnazyunus

dc/s2n-quic-dc/src/stream/endpoint.rs

@@ -7,7 +7,7 @@ use crate::{
     path::secret::{self, map, Map},
     random::Random,
     stream::{
-        application,
+        self, application,
         environment::{Environment, Peer},
         recv,
         send::{self, flow},
@@ -28,9 +28,8 @@ use super::environment::{ReadWorkerSocket as _, WriteWorkerSocket as _};
 
 type Result<T = (), E = io::Error> = core::result::Result<T, E>;
 
-pub struct AcceptError<Peer> {
+pub struct AcceptError {
     pub secret_control: Vec<u8>,
-    pub peer: Option<Peer>,
     pub error: io::Error,
 }
 
@@ -84,39 +83,47 @@ where
 }
 
 #[inline]
-pub fn accept_stream<Env, P>(
-    now: Timestamp,
-    env: &Env,
-    peer: P,
+pub fn derive_stream_credentials(
     packet: &server::InitialPacket,
     map: &Map,
-    subscriber_ctx: <Env::Subscriber as event::Subscriber>::ConnectionContext,
-    parameter_override: Option<&dyn Fn(dc::ApplicationParams) -> dc::ApplicationParams>,
-) -> Result<application::Builder<Env::Subscriber>, AcceptError<P>>
-where
-    Env: Environment,
-    P: Peer<Env>,
-{
+    features: &stream::TransportFeatures,
+    secret_control: &mut Vec<u8>,
+) -> Result<(secret::map::Bidirectional, dc::ApplicationParams), io::Error> {
     let credentials = &packet.credentials;
-    let mut secret_control = vec![];
-    let Some((crypto, mut parameters)) = map.pair_for_credentials(
+
+    let Some((crypto, parameters)) = map.pair_for_credentials(
         credentials,
         packet.source_queue_id,
-        &peer.features(),
-        &mut secret_control,
+        features,
+        secret_control,
     ) else {
         let error = io::Error::new(
             io::ErrorKind::NotFound,
             format!("missing credentials for client: {credentials:?}"),
         );
-        let error = AcceptError {
-            secret_control,
-            peer: Some(peer),
-            error,
-        };
         return Err(error);
     };
 
+    Ok((crypto, parameters))
+}
+
+#[inline]
+pub fn accept_stream<Env, P>(
+    now: Timestamp,
+    env: &Env,
+    peer: P,
+    packet: &server::InitialPacket,
+    map: &Map,
+    subscriber_ctx: <Env::Subscriber as event::Subscriber>::ConnectionContext,
+    parameter_override: Option<&dyn Fn(dc::ApplicationParams) -> dc::ApplicationParams>,
+    crypto: secret::map::Bidirectional,
+    mut parameters: dc::ApplicationParams,
+    secret_control: Vec<u8>,
+) -> Result<application::Builder<Env::Subscriber>, AcceptError>
+where
+    Env: Environment,
+    P: Peer<Env>,
+{
     if let Some(o) = parameter_override {
         parameters = o(parameters);
     }
@@ -148,7 +155,6 @@ where
         Err(error) => {
             let error = AcceptError {
                 secret_control,
-                peer: None,
                 error,
             };
             Err(error)

dc/s2n-quic-dc/src/stream/server/tokio.rs

@@ -523,6 +523,7 @@ impl<H: Handshake + Clone, S: event::Subscriber + Clone> Start<'_, H, S> {
             self.backlog,
             self.accept_flavor,
             self.linger,
+            tcp::worker::DefaultBehavior,
         )?
         .run();
 

dc/s2n-quic-dc/src/stream/server/tokio/tcp.rs

@@ -6,7 +6,7 @@ use crate::{
     path::secret,
     stream::{
         environment::{tokio::Environment, Environment as _},
-        server::accept,
+        server::{accept, tokio::tcp::worker::PollBehavior},
     },
 };
 use core::{future::poll_fn, task::Poll};
@@ -18,13 +18,14 @@ use tracing::debug;
 mod fresh;
 mod lazy;
 mod manager;
-mod worker;
+pub mod worker;
 
 pub(crate) use lazy::LazyBoundStream;
 
-pub struct Acceptor<Sub>
+pub struct Acceptor<Sub, B>
 where
     Sub: Subscriber + Clone,
+    B: PollBehavior<Sub> + Clone,
 {
     sender: accept::Sender<Sub>,
     socket: AsyncFd<TcpListener>,
@@ -33,11 +34,13 @@ where
     backlog: usize,
     accept_flavor: accept::Flavor,
     linger: Option<Duration>,
+    poll_behavior: B,
 }
 
-impl<Sub> Acceptor<Sub>
+impl<Sub, B> Acceptor<Sub, B>
 where
     Sub: event::Subscriber + Clone,
+    B: PollBehavior<Sub> + Clone,
 {
     #[inline]
     pub fn new(
@@ -49,6 +52,7 @@ where
         backlog: usize,
         accept_flavor: accept::Flavor,
         linger: Option<Duration>,
+        poll_behavior: B,
     ) -> std::io::Result<Self> {
         let acceptor = Self {
             sender: sender.clone(),
@@ -58,6 +62,7 @@ where
             backlog,
             accept_flavor,
             linger,
+            poll_behavior,
         };
 
         #[cfg(target_os = "linux")]
@@ -100,8 +105,12 @@ where
         let drop_guard = DropLog;
         let mut fresh = fresh::Queue::new(self.backlog);
         let mut workers = {
-            let workers =
-                (0..self.backlog).map(|_| worker::Worker::new(self.env.clock().get_time()));
+            let workers = (0..self.backlog).map(|_| {
+                worker::Worker::<Sub, B>::new(
+                    self.env.clock().get_time(),
+                    self.poll_behavior.clone(),
+                )
+            });
             manager::Manager::new(workers)
         };
         let mut context = worker::Context::new(&self);