fix(s2n-quic-dc): send connection close frames with pruned streams (#2831)

Author: camshaft

dc/s2n-quic-dc/src/clock.rs

@@ -2,7 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use core::{fmt, pin::Pin, task::Poll, time::Duration};
-use s2n_quic_core::{ensure, time};
+use s2n_quic_core::{
+    ensure, time,
+    time::{timer, timer::Provider},
+};
 use tracing::trace;
 
 #[macro_use]
@@ -62,7 +65,7 @@ pub trait Sleep: Clock + core::future::Future<Output = ()> {
 
 pub struct Timer {
     /// The `Instant` at which the timer should expire
-    target: Option<Timestamp>,
+    target: timer::Timer,
     /// The handle to the timer entry in the tokio runtime
     sleep: Pin<Box<dyn Sleep>>,
 }
@@ -88,16 +91,18 @@ impl Timer {
     #[inline]
     pub fn new_with_timeout(clock: &dyn Clock, timeout: Duration) -> Self {
         let (sleep, target) = clock.sleep(timeout);
+        let mut timer = timer::Timer::default();
+        timer.set(target);
         Self {
-            target: Some(target),
+            target: timer,
             sleep,
         }
     }
 
     #[inline]
     pub fn cancel(&mut self) {
         trace!(cancel = ?self.target);
-        self.target = None;
+        self.target.cancel();
     }
 
     pub async fn sleep(&mut self, target: Timestamp) {
@@ -110,13 +115,13 @@ impl Timer {
 impl time::clock::Timer for Timer {
     #[inline]
     fn poll_ready(&mut self, cx: &mut core::task::Context) -> Poll<()> {
-        ensure!(self.target.is_some(), Poll::Ready(()));
+        ensure!(self.target.is_armed(), Poll::Ready(()));
 
         let res = self.sleep.as_mut().poll(cx);
 
         if res.is_ready() {
             // clear the target after it fires, otherwise we'll endlessly wake up the task
-            self.target = None;
+            self.target.cancel();
         }
 
         res
@@ -125,9 +130,15 @@ impl time::clock::Timer for Timer {
     #[inline]
     fn update(&mut self, target: Timestamp) {
         // no need to update if it hasn't changed
-        ensure!(self.target != Some(target));
+        ensure!(self.target.next_expiration() != Some(target));
 
         self.sleep.as_mut().update(target);
-        self.target = Some(target);
+        self.target.set(target);
+    }
+}
+
+impl timer::Provider for Timer {
+    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
+        self.target.timers(query)
     }
 }

dc/s2n-quic-dc/src/packet/control/decoder.rs

@@ -279,6 +279,13 @@ pub struct ControlFramesMut<'a> {
     buffer: &'a mut [u8],
 }
 
+impl<'a> ControlFramesMut<'a> {
+    #[inline]
+    pub(crate) fn new(buffer: &'a mut [u8]) -> Self {
+        Self { buffer }
+    }
+}
+
 impl<'a> Iterator for ControlFramesMut<'a> {
     type Item = Result<FrameMut<'a>, s2n_codec::DecoderError>;
 

dc/s2n-quic-dc/src/packet/stream/decoder.rs

@@ -5,6 +5,7 @@ use crate::{
     credentials::Credentials,
     crypto,
     packet::{
+        control::decoder::ControlFramesMut,
         stream::{self, RelativeRetransmissionOffset, Tag},
         WireVersion,
     },
@@ -202,6 +203,11 @@ impl Packet<'_> {
         self.control_data.get(self.header)
     }
 
+    #[inline]
+    pub fn control_frames_mut(&mut self) -> ControlFramesMut<'_> {
+        ControlFramesMut::new(self.control_data.get_mut(self.header))
+    }
+
     #[inline]
     pub fn header(&self) -> &[u8] {
         self.header

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

@@ -65,12 +65,19 @@ impl Server {
 
         let server = s2n_quic::Server::builder().with_io(io)?;
 
+        let initial_max_data = builder.initial_data_window.unwrap_or_else(|| {
+            // default to only receive 10 packet worth before the application accepts the connection
+            builder.mtu as u64 * 10
+        });
+
         let connection_limits = s2n_quic::provider::limits::Limits::new()
             .with_max_idle_timeout(builder.max_idle_timeout)?
-            .with_data_window(builder.data_window)?
+            .with_data_window(initial_max_data)?
+            // After the connection is established we increase the data window to the configured value
             .with_bidirectional_local_data_window(builder.data_window)?
-            .with_bidirectional_remote_data_window(builder.data_window)?
+            .with_bidirectional_remote_data_window(initial_max_data)?
             .with_initial_round_trip_time(DEFAULT_INITIAL_RTT)?;
+
         let event = (ConfirmComplete, subscriber);
 
         let server = server
@@ -191,6 +198,7 @@ impl Client {
             .with_bidirectional_local_data_window(builder.data_window)?
             .with_bidirectional_remote_data_window(builder.data_window)?
             .with_initial_round_trip_time(DEFAULT_INITIAL_RTT)?;
+
         let event = (ConfirmComplete, subscriber);
 
         let client = client

dc/s2n-quic-dc/src/psk/server/builder.rs

@@ -18,6 +18,7 @@ pub struct Builder<
     #[allow(dead_code)]
     pub(crate) event_subscriber: Event,
     pub(crate) data_window: u64,
+    pub(crate) initial_data_window: Option<u64>,
     pub(crate) mtu: u16,
     pub(crate) max_idle_timeout: Duration,
     pub(crate) pto_jitter_percentage: u8,
@@ -28,6 +29,7 @@ impl Default for Builder<s2n_quic::provider::event::default::Subscriber> {
         Self {
             event_subscriber: Default::default(),
             data_window: DEFAULT_MAX_DATA,
+            initial_data_window: None,
             mtu: DEFAULT_MTU,
             max_idle_timeout: DEFAULT_IDLE_TIMEOUT,
             pto_jitter_percentage: DEFAULT_PTO_JITTER_PERCENTAGE,
@@ -44,6 +46,7 @@ impl<Event: s2n_quic::provider::event::Subscriber> Builder<Event> {
         Builder {
             event_subscriber,
             data_window: self.data_window,
+            initial_data_window: self.initial_data_window,
             mtu: self.mtu,
             max_idle_timeout: self.max_idle_timeout,
             pto_jitter_percentage: self.pto_jitter_percentage,
@@ -56,6 +59,15 @@ impl<Event: s2n_quic::provider::event::Subscriber> Builder<Event> {
         self
     }
 
+    /// Sets the initial amount of data that the peer is allowed to send before the application
+    /// accepts the stream
+    ///
+    /// This defaults to 10x the MTU if not set.
+    pub fn with_initial_data_window(mut self, initial_data_window: u64) -> Self {
+        self.initial_data_window = Some(initial_data_window);
+        self
+    }
+
     /// Sets the largest maximum transmission unit (MTU) that will be used for transmission
     pub fn with_mtu(mut self, mtu: u16) -> Self {
         self.mtu = mtu;

dc/s2n-quic-dc/src/socket/recv/router.rs

@@ -13,6 +13,12 @@ use s2n_quic_core::{
     varint::VarInt,
 };
 
+// Use `debug` logging for unhandled packets in non-test builds to reduce noise
+#[cfg(not(test))]
+use tracing::debug as warn;
+#[cfg(test)]
+use tracing::warn;
+
 mod with_map;
 mod zero_router;
 
@@ -139,7 +145,7 @@ pub trait Router {
         credentials: Credentials,
         segment: descriptor::Filled,
     ) {
-        tracing::warn!(
+        warn!(
             unhandled_packet = "control",
             ?tag,
             ?id,
@@ -168,7 +174,7 @@ pub trait Router {
         credentials: Credentials,
         segment: descriptor::Filled,
     ) {
-        tracing::warn!(
+        warn!(
             unhandled_packet = "stream",
             ?tag,
             ?id,
@@ -196,7 +202,7 @@ pub trait Router {
         credentials: Credentials,
         segment: descriptor::Filled,
     ) {
-        tracing::warn!(
+        warn!(
             unhandled_packet = "datagram",
             ?tag,
             ?credentials,
@@ -221,7 +227,7 @@ pub trait Router {
         credentials: credentials::Id,
         segment: descriptor::Filled,
     ) {
-        tracing::warn!(
+        warn!(
             unhandled_packet = "stale_key",
             ?queue_id,
             ?credentials,
@@ -246,7 +252,7 @@ pub trait Router {
         credentials: credentials::Id,
         segment: descriptor::Filled,
     ) {
-        tracing::warn!(
+        warn!(
             unhandled_packet = "replay_detected",
             ?queue_id,
             ?credentials,
@@ -270,7 +276,7 @@ pub trait Router {
         credentials: credentials::Id,
         segment: descriptor::Filled,
     ) {
-        tracing::warn!(
+        warn!(
             unhandled_packet = "unknown_path_secret",
             ?queue_id,
             ?credentials,
@@ -281,7 +287,7 @@ pub trait Router {
 
     #[inline]
     fn on_unhandled_packet(&mut self, remote_address: SocketAddress, packet: packet::Packet) {
-        tracing::warn!(unhandled_packet = ?packet, ?remote_address)
+        warn!(unhandled_packet = ?packet, ?remote_address)
     }
 
     #[inline]
@@ -291,7 +297,7 @@ pub trait Router {
         remote_address: SocketAddress,
         segment: descriptor::Filled,
     ) {
-        tracing::warn!(
+        warn!(
             ?error,
             ?remote_address,
             packet_len = segment.len(),

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

@@ -99,7 +99,8 @@ where
                 reason,
             });
 
-        // TODO emit event
+        self.shared.receiver.on_prune();
+        self.shared.sender.on_prune();
     }
 }
 
@@ -170,6 +171,14 @@ where
         self.write.write_from(buf).await
     }
 
+    #[inline]
+    pub async fn write_all_from(
+        &mut self,
+        buf: &mut impl buffer::reader::storage::Infallible,
+    ) -> io::Result<usize> {
+        self.write.write_all_from(buf).await
+    }
+
     #[inline]
     pub async fn write_from_fin(
         &mut self,
@@ -178,6 +187,14 @@ where
         self.write.write_from_fin(buf).await
     }
 
+    #[inline]
+    pub async fn write_all_from_fin(
+        &mut self,
+        buf: &mut impl buffer::reader::storage::Infallible,
+    ) -> io::Result<usize> {
+        self.write.write_all_from_fin(buf).await
+    }
+
     #[inline]
     pub async fn read_into(
         &mut self,

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

@@ -289,7 +289,7 @@ where
         let shared = shared.clone();
 
         let task = async move {
-            let mut reader = recv::worker::Worker::new(socket, shared, endpoint_type);
+            let mut reader = recv::worker::Worker::new(socket, shared, endpoint_type, &parameters);
 
             let mut prev_waker: Option<core::task::Waker> = None;
             core::future::poll_fn(|cx| {

dc/s2n-quic-dc/src/stream/environment/bach.rs

@@ -14,6 +14,7 @@ use crate::{
 use bach::ext::*;
 use s2n_quic_platform::features;
 use std::{io, net::SocketAddr, sync::Arc};
+use tracing::{info_span, Instrument};
 
 mod pool;
 pub mod udp;
@@ -188,8 +189,9 @@ where
     }
 
     #[inline]
+    #[track_caller]
     fn spawn_reader<F: 'static + Send + std::future::Future<Output = ()>>(&self, f: F) {
-        self.rt.spawn(f.primary());
+        self.rt.spawn(f.instrument(info_span!("reader")).primary());
     }
 
     #[inline]
@@ -198,7 +200,8 @@ where
     }
 
     #[inline]
+    #[track_caller]
     fn spawn_writer<F: 'static + Send + std::future::Future<Output = ()>>(&self, f: F) {
-        self.rt.spawn(f.primary());
+        self.rt.spawn(f.instrument(info_span!("writer")).primary());
     }
 }

dc/s2n-quic-dc/src/stream/recv/application.rs

@@ -5,7 +5,11 @@ use crate::{
     clock::Timer,
     event::{self, ConnectionPublisher as _},
     msg,
-    stream::{recv, runtime, shared::ArcShared, socket, Actor},
+    stream::{
+        recv, runtime,
+        shared::{AcceptState, ArcShared, ShutdownKind},
+        socket, Actor,
+    },
 };
 use core::{
     fmt,
@@ -227,7 +231,8 @@ where
 
         loop {
             // try to process any bytes we have in the recv buffer
-            reader.process_recv_buffer(out_buf, shared, transport_features);
+            // - use the `Accepted` state since this is the application interface
+            reader.process_recv_buffer(out_buf, shared, transport_features, AcceptState::Accepted);
 
             // if we still have remaining capacity in the `out_buf` make sure the reassembler is
             // fully drained
@@ -239,6 +244,12 @@ where
             if let Err(err) = reader.receiver.check_error() {
                 self.local_state
                     .transition(LocalState::Errored(err), &self.shared);
+
+                if out_buf.written_len() > 0 {
+                    // if we've written something to the buffer then return that first
+                    break;
+                }
+
                 return Err(err.into()).into();
             }
 
@@ -370,9 +381,13 @@ where
             .transition(LocalState::Drained, &self.shared);
 
         // let the peer know if we shut down cleanly
-        let is_panicking = std::thread::panicking();
+        let kind = if std::thread::panicking() {
+            ShutdownKind::Panicking
+        } else {
+            ShutdownKind::Normal
+        };
 
-        self.shared.receiver.shutdown(is_panicking);
+        self.shared.receiver.shutdown(kind);
     }
 
     #[inline(always)]