diff --git a/quic/s2n-quic-platform/src/io/testing.rs b/quic/s2n-quic-platform/src/io/testing.rs index b0895f3100..527903314b 100644 --- a/quic/s2n-quic-platform/src/io/testing.rs +++ b/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 { diff --git a/quic/s2n-quic-tests/src/tests.rs b/quic/s2n-quic-tests/src/tests.rs index 3b28a549fd..291990a220 100644 --- a/quic/s2n-quic-tests/src/tests.rs +++ b/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; diff --git a/quic/s2n-quic-tests/src/tests/close_all_stream_after_endpoint_drop.rs b/quic/s2n-quic-tests/src/tests/close_all_stream_after_endpoint_drop.rs new file mode 100644 index 0000000000..0ce2f6d446 --- /dev/null +++ b/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; + }); +} diff --git a/quic/s2n-quic-transport/src/connection/api.rs b/quic/s2n-quic-transport/src/connection/api.rs index 4e6c947f47..2cb1d12a54 100644 --- a/quic/s2n-quic-transport/src/connection/api.rs +++ b/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] diff --git a/quic/s2n-quic-transport/src/connection/api_provider.rs b/quic/s2n-quic-transport/src/connection/api_provider.rs index b892805c0b..d383cca29f 100644 --- a/quic/s2n-quic-transport/src/connection/api_provider.rs +++ b/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>; - fn close_connection(&self, code: Option); + fn close_connection(&self, code: Option); fn server_name(&self) -> Result, connection::Error>; diff --git a/quic/s2n-quic-transport/src/connection/connection_container.rs b/quic/s2n-quic-transport/src/connection/connection_container.rs index ed92f073b5..7fda66403e 100644 --- a/quic/s2n-quic-transport/src/connection/connection_container.rs +++ b/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> ConnectionApiProvider for Con } } - fn close_connection(&self, error: Option) { + fn close_connection(&self, error: Option) { let _: Result<(), connection::Error> = self.api_write_call(|conn| { conn.application_close(error); Ok(()) @@ -1146,6 +1145,17 @@ impl> ConnectionContainer { } } +impl> Drop for ConnectionContainer { + // 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)] diff --git a/quic/s2n-quic-transport/src/connection/connection_container/tests.rs b/quic/s2n-quic-transport/src/connection/connection_container/tests.rs index ab6b42d8aa..811bb37df2 100644 --- a/quic/s2n-quic-transport/src/connection/connection_container/tests.rs +++ b/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 = const { Cell::new(0) }; + static APPLICATION_CLOSE_CALL_COUNT: Cell = const { Cell::new(0) }; + static APPLICATION_CLOSE_WITH_ERROR_COUNT: Cell = 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, _context: &Context, ) -> Poll, 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> { - 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) { - // no-op + fn application_close(&mut self, _error: Option) { + 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 { @@ -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 = + 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)); +} diff --git a/quic/s2n-quic-transport/src/connection/connection_impl.rs b/quic/s2n-quic-transport/src/connection/connection_impl.rs index 4c6385331d..65791d767f 100644 --- a/quic/s2n-quic-transport/src/connection/connection_impl.rs +++ b/quic/s2n-quic-transport/src/connection/connection_impl.rs @@ -33,7 +33,7 @@ use core::{ }; use s2n_codec::DecoderBufferMut; use s2n_quic_core::{ - application::{self, ServerName}, + application::ServerName, connection::{ error::Error, id::{Classification, Generator as _}, @@ -2281,7 +2281,7 @@ impl connection::Trait for ConnectionImpl { ) } - fn application_close(&mut self, error: Option) { + fn application_close(&mut self, error: Option) { if self.error.is_err() { return; } @@ -2290,7 +2290,11 @@ impl connection::Trait for ConnectionImpl { self.open_registry = None; if let Some(error) = error { - self.error = Err(connection::Error::application(error)); + self.error = Err(error); + // This will put all streams into Reset state and wake all tasks + if let Some((space, _)) = self.space_manager.application_mut() { + space.stream_manager.close(error); + } } else { // give the connection some time to flush all outstanding streams self.state = ConnectionState::Flushing; diff --git a/quic/s2n-quic-transport/src/connection/connection_trait.rs b/quic/s2n-quic-transport/src/connection/connection_trait.rs index 481cf48315..54f3de2eeb 100644 --- a/quic/s2n-quic-transport/src/connection/connection_trait.rs +++ b/quic/s2n-quic-transport/src/connection/connection_trait.rs @@ -21,7 +21,6 @@ use core::{ }; use s2n_codec::DecoderBufferMut; use s2n_quic_core::{ - application, application::ServerName, event::{self, builder::DatagramDropReason, supervisor, ConnectionPublisher, IntoEvent}, inet::{DatagramInfo, SocketAddress}, @@ -518,7 +517,7 @@ pub trait ConnectionTrait: 'static + Send + Sized { context: &Context, ) -> Poll>; - fn application_close(&mut self, error: Option); + fn application_close(&mut self, error: Option); fn server_name(&self) -> Option;