fix(s2n-quic-transport): close all streams after Endpoint drop (#3170)

quic/s2n-quic-platform/src/io/testing.rs

@@ -222,6 +222,10 @@ impl Handle {
             queue_send_buffer_size: None,
         }
     }
+
+    pub fn close_buffers(&self) {
+        self.buffers.close();
+    }
 }
 
 pub struct Builder {

quic/s2n-quic-tests/src/tests.rs

@@ -26,6 +26,7 @@ use std::{
 
 mod blackhole;
 mod buffer_limit;
+mod close_all_stream_after_endpoint_drop;
 mod connection_limits;
 mod connection_migration;
 mod deduplicate;

quic/s2n-quic-tests/src/tests/close_all_stream_after_endpoint_drop.rs

@@ -0,0 +1,199 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use super::*;
+use bach::time::timeout;
+
+const CLOSE_PROPAGATION_DURATION: Duration = Duration::from_millis(5);
+
+/// Verifies client endpoint-drop tx stream behavior.
+///
+/// After endpoint drop, tx stream send should complete with an error.
+#[test]
+fn client_tx_fails_after_endpoint_drop() {
+    let model = Model::default();
+    test(model.clone(), |handle| {
+        let server = build_server(handle, model.clone())?;
+        let server_addr = start_server(server)?;
+
+        let h = handle.clone();
+
+        let client = build_client(handle, model, true).unwrap();
+
+        primary::spawn(async move {
+            let connect = Connect::new(server_addr).with_server_name("localhost");
+            let mut client_connection = client.connect(connect).await.unwrap();
+
+            let mut stream = client_connection.open_bidirectional_stream().await.unwrap();
+
+            let sent = Bytes::from("hello");
+            stream.send(sent.clone()).await.unwrap();
+            let received = stream.receive().await.unwrap().unwrap();
+            assert_eq!(sent, received);
+
+            // Drop endpoint
+            h.close_buffers();
+
+            delay(CLOSE_PROPAGATION_DURATION).await;
+
+            let send_res = stream.send(Bytes::from("world")).await;
+            assert!(send_res.is_err());
+        });
+
+        Ok(())
+    })
+    .unwrap();
+}
+
+/// Verifies client endpoint-drop rx stream behavior.
+///
+/// After endpoint drop, rx stream accept should complete with an error instead of hanging.
+#[test]
+fn client_rx_fails_after_endpoint_drop() {
+    let model = Model::default();
+    test(model.clone(), |handle| {
+        let server = build_server(handle, model.clone())?;
+        let server_addr = start_server(server)?;
+
+        let h = handle.clone();
+
+        let client = build_client(handle, model, true).unwrap();
+
+        primary::spawn(async move {
+            let connect = Connect::new(server_addr).with_server_name("localhost");
+            let mut client_connection = client.connect(connect).await.unwrap();
+
+            let mut stream = client_connection.open_bidirectional_stream().await.unwrap();
+
+            let sent = Bytes::from("hello");
+            stream.send(sent.clone()).await.unwrap();
+            let received = stream.receive().await.unwrap().unwrap();
+            assert_eq!(sent, received);
+
+            // Drop endpoint
+            h.close_buffers();
+
+            // The async call should return an error
+            let res = timeout(
+                CLOSE_PROPAGATION_DURATION,
+                async move { stream.receive().await },
+            )
+            .await
+            .unwrap();
+            assert!(res.is_err());
+        });
+
+        Ok(())
+    })
+    .unwrap();
+}
+
+/// Verifies server endpoint-drop tx stream behavior.
+///
+/// After server endpoint drop, tx stream send should complete with an error.
+#[test]
+fn server_tx_fails_after_endpoint_drop() {
+    let model = Model::default();
+    test(model.clone(), |handle| {
+        let mut server = build_server(handle, model.clone())?;
+        let server_addr = server.local_addr()?;
+
+        let h = handle.clone();
+
+        let client = build_client(handle, model, true).unwrap();
+
+        run_echo_client(client, server_addr);
+
+        primary::spawn(async move {
+            let mut connection = server.accept().await.unwrap();
+
+            let mut stream = connection
+                .accept_bidirectional_stream()
+                .await
+                .unwrap()
+                .unwrap();
+
+            let data = stream.receive().await.unwrap().unwrap();
+            stream.send(Bytes::from("hello")).await.unwrap();
+
+            // Drop endpoint
+            h.close_buffers();
+
+            delay(CLOSE_PROPAGATION_DURATION).await;
+
+            let send_res = stream.send(data).await;
+            assert!(send_res.is_err());
+        });
+
+        Ok(())
+    })
+    .unwrap();
+}
+
+/// Verifies server endpoint-drop rx stream behavior.
+///
+/// After server endpoint drop, rx stream receive should complete with an error instead of hanging.
+#[test]
+fn server_rx_fails_after_endpoint_drop() {
+    let model = Model::default();
+    test(model.clone(), |handle| {
+        let mut server = build_server(handle, model.clone())?;
+        let server_addr = server.local_addr()?;
+
+        let h = handle.clone();
+
+        let client = build_client(handle, model, true).unwrap();
+
+        run_echo_client(client, server_addr);
+
+        primary::spawn(async move {
+            let mut connection = server.accept().await.unwrap();
+
+            let mut stream = connection
+                .accept_bidirectional_stream()
+                .await
+                .unwrap()
+                .unwrap();
+
+            let _data = stream.receive().await.unwrap().unwrap();
+            stream.send(Bytes::from("hello")).await.unwrap();
+
+            // Drop endpoint
+            h.close_buffers();
+
+            // The async call should return an error
+            let res = timeout(
+                CLOSE_PROPAGATION_DURATION,
+                async move { stream.receive().await },
+            )
+            .await
+            .unwrap();
+            assert!(res.is_err());
+        });
+
+        Ok(())
+    })
+    .unwrap();
+}
+
+fn run_echo_client(client: Client, server_addr: SocketAddr) {
+    spawn(async move {
+        let connect = Connect::new(server_addr).with_server_name("localhost");
+        let Ok(mut client_connection) = client.connect(connect).await else {
+            return;
+        };
+
+        let Ok(mut stream) = client_connection.open_bidirectional_stream().await else {
+            return;
+        };
+
+        let sent = Bytes::from("hello");
+        let Ok(()) = stream.send(sent.clone()).await else {
+            return;
+        };
+        _ = stream.receive().await;
+
+        // Prevent dropping the connection and the stream
+        delay(Duration::from_secs(60)).await;
+    });
+}

quic/s2n-quic-transport/src/connection/api.rs

@@ -163,7 +163,8 @@ impl Connection {
     /// This will immediately terminate all outstanding streams.
     #[inline]
     pub fn close(&self, error_code: application::Error) {
-        self.api.close_connection(Some(error_code));
+        self.api
+            .close_connection(Some(connection::Error::application(error_code)));
     }
 
     #[inline]

quic/s2n-quic-transport/src/connection/api_provider.rs

@@ -16,7 +16,6 @@ use core::{
     task::{Context, Poll},
 };
 use s2n_quic_core::{
-    application,
     application::ServerName,
     inet::SocketAddress,
     query::{Query, QueryMut},
@@ -53,7 +52,7 @@ pub(crate) trait ConnectionApiProvider: Sync + Send {
         context: &Context,
     ) -> Poll<Result<Stream, connection::Error>>;
 
-    fn close_connection(&self, code: Option<application::Error>);
+    fn close_connection(&self, code: Option<connection::Error>);
 
     fn server_name(&self) -> Result<Option<ServerName>, connection::Error>;
 

quic/s2n-quic-transport/src/connection/connection_container.rs

@@ -32,7 +32,6 @@ use intrusive_collections::{
     intrusive_adapter, KeyAdapter, LinkedList, LinkedListLink, RBTree, RBTreeLink,
 };
 use s2n_quic_core::{
-    application,
     application::ServerName,
     event::supervisor,
     inet::SocketAddress,
@@ -296,7 +295,7 @@ impl<C: connection::Trait, L: connection::Lock<C>> ConnectionApiProvider for Con
         }
     }
 
-    fn close_connection(&self, error: Option<application::Error>) {
+    fn close_connection(&self, error: Option<connection::Error>) {
         let _: Result<(), connection::Error> = self.api_write_call(|conn| {
             conn.application_close(error);
             Ok(())
@@ -1146,6 +1145,17 @@ impl<C: connection::Trait, L: connection::Lock<C>> ConnectionContainer<C, L> {
     }
 }
 
+impl<C: connection::Trait, L: connection::Lock<C>> Drop for ConnectionContainer<C, L> {
+    // ConnectionContainer dropped means Endpoint is dropped. Close all connections.
+    fn drop(&mut self) {
+        let mut cursor = self.connection_map.front();
+        while let Some(node) = cursor.get() {
+            node.close_connection(Some(connection::Error::endpoint_closing()));
+            cursor.move_next();
+        }
+    }
+}
+
 /// Return values for iterations over a `Connection` list.
 /// The value instructs the iterator whether iteration will be continued.
 #[derive(Clone, Copy, Debug)]

quic/s2n-quic-transport/src/connection/connection_container/tests.rs

@@ -14,11 +14,12 @@ use bolero::{check, generator::*};
 use bytes::Bytes;
 use core::{
     any::Any,
+    cell::Cell,
     task::{Context, Poll},
     time::Duration,
 };
 use s2n_quic_core::{
-    application, event,
+    event,
     event::builder::DatagramDropReason,
     inet::{DatagramInfo, SocketAddress},
     io::tx,
@@ -36,6 +37,12 @@ use s2n_quic_core::{
 };
 use std::sync::Mutex;
 
+thread_local! {
+    static CLOSE_CALL_COUNT: Cell<usize> = const { Cell::new(0) };
+    static APPLICATION_CLOSE_CALL_COUNT: Cell<usize> = const { Cell::new(0) };
+    static APPLICATION_CLOSE_WITH_ERROR_COUNT: Cell<usize> = const { Cell::new(0) };
+}
+
 struct TestConnection {
     accept_state: AcceptState,
     is_closed: bool,
@@ -90,6 +97,7 @@ impl connection::Trait for TestConnection {
     ) {
         assert!(!self.is_closed);
         assert!(!self.close_timer.is_armed());
+        CLOSE_CALL_COUNT.with(|count| count.set(count.get() + 1));
         self.close_timer.set(timestamp + Duration::from_secs(1));
     }
 
@@ -293,7 +301,11 @@ impl connection::Trait for TestConnection {
         _stream_type: Option<stream::StreamType>,
         _context: &Context,
     ) -> Poll<Result<Option<stream::StreamId>, connection::Error>> {
-        todo!()
+        if self.is_closed {
+            return Poll::Ready(Err(connection::Error::unspecified()));
+        }
+
+        Poll::Pending
     }
 
     fn poll_open_stream(
@@ -302,11 +314,20 @@ impl connection::Trait for TestConnection {
         _token: &mut connection::OpenToken,
         _context: &Context,
     ) -> Poll<Result<stream::StreamId, connection::Error>> {
-        todo!()
+        if self.is_closed {
+            return Poll::Ready(Err(connection::Error::unspecified()));
+        }
+
+        Poll::Ready(Ok(stream::StreamId::from_varint(
+            s2n_quic_core::varint::VarInt::from_u8(0),
+        )))
     }
 
-    fn application_close(&mut self, _error: Option<application::Error>) {
-        // no-op
+    fn application_close(&mut self, _error: Option<connection::Error>) {
+        APPLICATION_CLOSE_CALL_COUNT.with(|count| count.set(count.get() + 1));
+        if _error.is_some() {
+            APPLICATION_CLOSE_WITH_ERROR_COUNT.with(|count| count.set(count.get() + 1));
+        }
     }
 
     fn server_name(&self) -> Option<ServerName> {
@@ -627,3 +648,25 @@ fn container_test() {
         assert!(connections.next().is_none());
     });
 }
+
+#[test]
+fn drop_closes_all_connections() {
+    APPLICATION_CLOSE_CALL_COUNT.with(|count| count.set(0));
+    APPLICATION_CLOSE_WITH_ERROR_COUNT.with(|count| count.set(0));
+
+    let mut id_gen = InternalConnectionIdGenerator::new();
+    let (_handle, acceptor, connector, _close_handle) = endpoint::handle::Handle::new(100);
+
+    {
+        let mut container: ConnectionContainer<TestConnection, TestLock> =
+            ConnectionContainer::new(acceptor, connector);
+
+        for _ in 0..3 {
+            let id = id_gen.generate_id();
+            container.insert_connection(TestConnection::default(), id);
+        }
+    }
+
+    APPLICATION_CLOSE_CALL_COUNT.with(|count| assert_eq!(count.get(), 3));
+    APPLICATION_CLOSE_WITH_ERROR_COUNT.with(|count| assert_eq!(count.get(), 3));
+}