feat(s2n-quic-dc): accept linger parameter instead of always setting it

Author: camshaft

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

@@ -11,7 +11,7 @@ use crate::{
         socket::Protocol,
     },
 };
-use std::{io, net::SocketAddr};
+use std::{io, net::SocketAddr, time::Duration};
 use tokio::net::TcpStream;
 
 /// Connects using the UDP transport layer
@@ -54,6 +54,7 @@ pub async fn connect_tcp<H, Sub>(
     acceptor_addr: SocketAddr,
     env: &Environment<Sub>,
     subscriber: Sub,
+    linger: Option<Duration>,
 ) -> io::Result<Stream<Sub>>
 where
     H: core::future::Future<Output = io::Result<secret::map::Peer>>,
@@ -64,7 +65,10 @@ where
 
     // Make sure TCP_NODELAY is set
     let _ = socket.set_nodelay(true);
-    let _ = socket.set_linger(Some(core::time::Duration::ZERO));
+
+    if linger.is_some() {
+        let _ = socket.set_linger(linger);
+    }
 
     // if the acceptor_ip isn't known, then ask the socket to resolve it for us
     let peer_addr = if acceptor_addr.ip().is_unspecified() {

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

@@ -9,6 +9,7 @@ use crate::{
 };
 use core::{future::poll_fn, task::Poll};
 use s2n_quic_core::{inet::SocketAddress, time::Clock};
+use std::time::Duration;
 use tokio::net::TcpListener;
 use tracing::debug;
 
@@ -26,6 +27,7 @@ where
     secrets: secret::Map,
     backlog: usize,
     accept_flavor: accept::Flavor,
+    linger: Option<Duration>,
     subscriber: Sub,
 }
 
@@ -42,6 +44,7 @@ where
         secrets: &secret::Map,
         backlog: usize,
         accept_flavor: accept::Flavor,
+        linger: Option<Duration>,
         subscriber: Sub,
     ) -> Self {
         let acceptor = Self {
@@ -51,6 +54,7 @@ where
             secrets: secrets.clone(),
             backlog,
             accept_flavor,
+            linger,
             subscriber,
         };
 
@@ -98,6 +102,7 @@ where
                 workers.insert(
                     remote_address,
                     socket,
+                    self.linger,
                     &mut context,
                     subscriber_ctx,
                     &publisher,

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

@@ -153,6 +153,7 @@ where
         &mut self,
         remote_address: SocketAddress,
         stream: W::Stream,
+        linger: Option<Duration>,
         cx: &mut W::Context,
         connection_context: W::ConnectionContext,
         publisher: &Pub,
@@ -179,6 +180,7 @@ where
         self.inner.workers[idx].worker.replace(
             remote_address,
             stream,
+            linger,
             connection_context,
             publisher,
             clock,
@@ -377,6 +379,7 @@ pub trait Worker {
         &mut self,
         remote_address: SocketAddress,
         stream: Self::Stream,
+        linger: Option<Duration>,
         connection_context: Self::ConnectionContext,
         publisher: &Pub,
         clock: &C,

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

@@ -64,6 +64,7 @@ impl super::Worker for Worker {
         &mut self,
         _remote_address: SocketAddress,
         _stream: Self::Stream,
+        _linger: Option<Duration>,
         _connection_context: Self::ConnectionContext,
         _publisher: &Pub,
         clock: &C,
@@ -160,6 +161,7 @@ impl Harness {
         self.manager.insert(
             SocketAddress::default(),
             (),
+            None,
             &mut (),
             (),
             &publisher(&self.subscriber, &self.clock),

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

@@ -98,6 +98,7 @@ where
         &mut self,
         remote_address: SocketAddress,
         stream: TcpStream,
+        linger: Option<Duration>,
         subscriber_ctx: Self::ConnectionContext,
         publisher: &Pub,
         clock: &C,
@@ -107,7 +108,10 @@ where
     {
         // Make sure TCP_NODELAY is set
         let _ = stream.set_nodelay(true);
-        let _ = stream.set_linger(Some(Duration::ZERO));
+
+        if linger.is_some() {
+            let _ = stream.set_linger(linger);
+        }
 
         let now = clock.get_time();
 
@@ -331,6 +335,10 @@ impl WorkerState {
                                 error: error.error,
                             };
                             continue;
+                        } else {
+                            // close the stream immediately and send a reset to the client
+                            let _ = socket.set_linger(Some(Duration::ZERO));
+                            drop(socket);
                         }
                     }
                     return Err(Some(error.error)).into();
@@ -381,16 +389,15 @@ impl WorkerState {
     }
 
     #[inline]
-    fn poll_initial_packet<S, Pub>(
+    fn poll_initial_packet<Pub>(
         cx: &mut task::Context,
-        stream: &mut S,
+        stream: &mut TcpStream,
         remote_address: &SocketAddress,
         recv_buffer: &mut msg::recv::Message,
         sojourn_time: Duration,
         publisher: &Pub,
     ) -> Poll<Result<server::InitialPacket, Option<io::Error>>>
     where
-        S: Socket,
         Pub: EndpointPublisher,
     {
         loop {
@@ -403,6 +410,10 @@ impl WorkerState {
                         sojourn_time,
                     },
                 );
+
+                // close the stream immediately and send a reset to the client
+                let _ = stream.set_linger(Some(Duration::ZERO));
+
                 return Err(None).into();
             }
 
@@ -437,6 +448,9 @@ impl WorkerState {
                         },
                     );
 
+                    // close the stream immediately and send a reset to the client
+                    let _ = stream.set_linger(Some(Duration::ZERO));
+
                     return Err(None).into();
                 }
             }

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

@@ -68,8 +68,14 @@ impl Client {
 
         match server.protocol {
             Protocol::Tcp => {
-                stream_client::connect_tcp(handshake, server.local_addr, &self.env, subscriber)
-                    .await
+                stream_client::connect_tcp(
+                    handshake,
+                    server.local_addr,
+                    &self.env,
+                    subscriber,
+                    None,
+                )
+                .await
             }
             Protocol::Udp => {
                 stream_client::connect_udp(handshake, server.local_addr, &self.env, subscriber)
@@ -181,6 +187,8 @@ mod drop_handle {
 }
 
 pub mod server {
+    use std::time::Duration;
+
     use super::*;
 
     #[derive(Clone)]
@@ -201,6 +209,7 @@ pub mod server {
         flavor: accept::Flavor,
         protocol: Protocol,
         map_capacity: usize,
+        linger: Option<Duration>,
         subscriber: event::testing::Subscriber,
     }
 
@@ -211,6 +220,7 @@ pub mod server {
                 flavor: accept::Flavor::default(),
                 protocol: Protocol::Tcp,
                 map_capacity: 16,
+                linger: None,
                 subscriber: event::testing::Subscriber::no_snapshot(),
             }
         }
@@ -255,6 +265,11 @@ pub mod server {
             self
         }
 
+        pub fn linger(mut self, linger: Duration) -> Self {
+            self.linger = Some(linger);
+            self
+        }
+
         pub fn subscriber(mut self, subscriber: event::testing::Subscriber) -> Self {
             self.subscriber = subscriber;
             self
@@ -266,6 +281,7 @@ pub mod server {
                 flavor,
                 protocol,
                 map_capacity,
+                linger,
                 subscriber,
             } = self;
 
@@ -291,7 +307,7 @@ pub mod server {
                     let socket = tokio::net::TcpListener::from_std(socket).unwrap();
 
                     let acceptor = stream_server::tcp::Acceptor::new(
-                        0, socket, &sender, &env, &map, backlog, flavor, subscriber,
+                        0, socket, &sender, &env, &map, backlog, flavor, linger, subscriber,
                     );
                     let acceptor = drop_handle_receiver.wrap(acceptor.run());
                     let acceptor = acceptor.instrument(tracing::info_span!("tcp"));