test(s2n-quic-dc): implement tokio stream tests for dcQUIC (#2772)

Author: boquan-fang

dc/s2n-quic-dc/Cargo.toml

@@ -68,6 +68,7 @@ bolero = "0.13"
 bolero-generator = "0.13"
 insta = "1"
 s2n-codec = { path = "../../common/s2n-codec", features = ["testing"] }
+s2n-quic = { path = "../../quic/s2n-quic" }
 s2n-quic-core = { path = "../../quic/s2n-quic-core", features = ["testing"] }
 s2n-quic-platform = { path = "../../quic/s2n-quic-platform", features = [
     "testing",
@@ -77,4 +78,4 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 
 [lints.rust.unexpected_cfgs]
 level = "warn"
-check-cfg = ['cfg(fuzzing)', 'cfg(kani)', 'cfg(todo)']
+check-cfg = ['cfg(future)', 'cfg(fuzzing)', 'cfg(kani)', 'cfg(todo)']

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

@@ -201,7 +201,7 @@ impl Provider {
         Ok((peer, HandshakeKind::Fresh))
     }
 
-    /// Handshake with a peer in the background.∂
+    /// Handshake with a peer in the background.
     #[inline]
     pub fn background_handshake_with(
         &self,
@@ -215,7 +215,6 @@ impl Provider {
 
         let client = self.state.client.clone();
         if let Some((runtime, _)) = self.state.runtime.as_ref() {
-            let server_name = server_name.clone();
             // Drop the JoinHandle -- we're not actually going to block on the join handle's
             // result. The future will keep running in the background.
             runtime.spawn(async move {

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

@@ -8,6 +8,8 @@ pub const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(30);
 /// The maximum length of a single packet written to a stream
 pub const MAX_DATAGRAM_SIZE: usize = 1 << 15; // 32k
 
+pub use crate::stream::socket::Protocol;
+
 pub mod application;
 pub mod client;
 pub mod crypto;

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

@@ -26,8 +26,8 @@ pub mod rpc {
     pub use crate::stream::client::rpc::{InMemoryResponse, Request, Response};
 }
 
-// This trait is a temporary solution to abstract handshake_with_entry,
-// local_addr, and map methods until we implement the handshake provider
+// This trait is a solution to abstract handshake_with_entry,
+// local_addr, and map methods
 #[allow(async_fn_in_trait)]
 pub trait Handshake: Clone {
     /// Handshake with the remote peer

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

@@ -25,7 +25,7 @@ use tracing::{trace, Instrument as _};
 pub mod tcp;
 pub mod udp;
 
-// This trait is a temporary solution to abstract local_addr and map methods until we implement the handshake provider
+// This trait is a solution to abstract local_addr and map methods
 pub trait Handshake: Clone {
     fn local_addr(&self) -> SocketAddr;
 

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

@@ -1,17 +1,19 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use super::{client::tokio::Client as ClientTokio, server::tokio::Server as ServerTokio};
 use crate::{
     either::Either,
     event::{self, testing},
     path::secret,
+    psk::{client::Provider as ClientProvider, server::Provider as ServerProvider},
     stream::{
         application, client as stream_client,
         environment::{bach, tokio, udp, Environment},
-        recv, send,
         server::{self as stream_server, accept, stats},
         socket::Protocol,
     },
+    testing::NoopSubscriber,
 };
 use s2n_quic_core::dc::{self, ApplicationParams};
 use s2n_quic_platform::socket;
@@ -31,8 +33,6 @@ thread_local! {
 pub type Subscriber = (Arc<event::testing::Subscriber>, event::tracing::Subscriber);
 
 pub type Stream = application::Stream<Subscriber>;
-pub type Writer = send::application::Writer<Subscriber>;
-pub type Reader = recv::application::Reader<Subscriber>;
 
 const DEFAULT_POOLED: bool = true;
 
@@ -47,6 +47,264 @@ pub(crate) const MAX_DATAGRAM_SIZE: u16 = if cfg!(target_os = "linux") {
 
 type Env = Either<tokio::Environment<Subscriber>, bach::Environment<Subscriber>>;
 
+pub fn bind_pair(
+    protocol: Protocol,
+    server_addr: SocketAddr,
+    client: ClientProvider,
+    server: ServerProvider,
+) -> (
+    ClientTokio<ClientProvider, NoopSubscriber>,
+    ServerTokio<ServerProvider, NoopSubscriber>,
+) {
+    let test_subscriber = NoopSubscriber {};
+    let client = ClientTokio::<ClientProvider, NoopSubscriber>::builder()
+        .with_default_protocol(protocol)
+        .build(client, test_subscriber.clone())
+        .unwrap();
+
+    let server = ServerTokio::<ServerProvider, NoopSubscriber>::builder()
+        .with_address(server_addr)
+        .with_protocol(protocol)
+        .with_workers(1.try_into().unwrap())
+        .build(server, test_subscriber)
+        .unwrap();
+
+    (client, server)
+}
+
+macro_rules! check_pair_addrs {
+    ($local:ident, $peer:ident) => {
+        debug_assert_eq!(
+            $local.local_addr().ok().map(|addr| addr.port()),
+            $peer.peer_addr().ok().map(|addr| addr.port())
+        );
+    };
+}
+
+macro_rules! dcquic_context {
+    ($protocol:ident) => {
+        use super::Protocol;
+        use std::net::SocketAddr;
+
+        pub type Stream = crate::stream::application::Stream<super::NoopSubscriber>;
+
+        pub struct Context(super::Context);
+
+        #[allow(dead_code)]
+        impl Context {
+            pub async fn new() -> Self {
+                Self(super::Context::new(Protocol::$protocol).await)
+            }
+
+            pub fn new_sync(addr: SocketAddr) -> Self {
+                Self(super::Context::new_sync(Protocol::$protocol, addr))
+            }
+
+            pub async fn bind(addr: SocketAddr) -> Self {
+                Self(super::Context::bind(Protocol::$protocol, addr).await)
+            }
+
+            pub fn acceptor_addr(&self) -> SocketAddr {
+                self.0.acceptor_addr()
+            }
+
+            pub fn handshake_addr(&self) -> SocketAddr {
+                self.0.handshake_addr()
+            }
+
+            pub async fn pair(&self) -> (Stream, Stream) {
+                self.0.pair().await
+            }
+
+            pub async fn pair_with(&self, acceptor_addr: SocketAddr) -> (Stream, Stream) {
+                self.0.pair_with(acceptor_addr).await
+            }
+
+            pub fn protocol(&self) -> Protocol {
+                self.0.protocol
+            }
+        }
+    };
+}
+
+pub mod dcquic {
+    use crate::{
+        psk::{client::Provider as ClientProvider, server::Provider as ServerProvider},
+        stream::{
+            client::tokio::Client as ClientTokio,
+            server::tokio::Server as ServerTokio,
+            socket::Protocol,
+            testing::{bind_pair, NoopSubscriber},
+        },
+        testing::server_name,
+    };
+    use std::net::SocketAddr;
+
+    pub type Stream = crate::stream::application::Stream<NoopSubscriber>;
+
+    pub mod tcp {
+        dcquic_context!(Tcp);
+    }
+
+    pub mod udp {
+        dcquic_context!(Udp);
+    }
+
+    pub struct Context {
+        pub(crate) server: ServerTokio<ServerProvider, NoopSubscriber>,
+        pub(crate) client: ClientTokio<ClientProvider, NoopSubscriber>,
+        protocol: Protocol,
+    }
+
+    impl Context {
+        pub async fn new(protocol: Protocol) -> Self {
+            if protocol.is_udp() {
+                Self::bind(protocol, "[::1]:0".parse().unwrap()).await
+            } else {
+                Self::bind(protocol, "127.0.0.1:0".parse().unwrap()).await
+            }
+        }
+
+        pub fn new_sync(protocol: Protocol, addr: SocketAddr) -> Self {
+            let (client, server) = crate::testing::pair_sync();
+            let (client, server) = bind_pair(protocol, addr, client, server);
+            Self {
+                client,
+                server,
+                protocol,
+            }
+        }
+
+        pub async fn bind(protocol: Protocol, addr: SocketAddr) -> Self {
+            Self::bind_with(protocol, addr, crate::testing::Pair::default()).await
+        }
+
+        pub async fn bind_with(
+            protocol: Protocol,
+            addr: SocketAddr,
+            pair: crate::testing::Pair,
+        ) -> Self {
+            let (client, server) = pair.build().await;
+            let (client, server) = bind_pair(protocol, addr, client, server);
+            Self {
+                client,
+                server,
+                protocol,
+            }
+        }
+
+        pub fn acceptor_addr(&self) -> SocketAddr {
+            self.server.acceptor_addr().expect("acceptor_addr")
+        }
+
+        pub fn handshake_addr(&self) -> SocketAddr {
+            self.server.handshake_addr().expect("handshake_addr")
+        }
+
+        pub async fn pair(&self) -> (Stream, Stream) {
+            self.pair_with(self.acceptor_addr()).await
+        }
+
+        pub async fn pair_with(&self, acceptor_addr: SocketAddr) -> (Stream, Stream) {
+            let handshake_addr = self.handshake_addr();
+            let (client, server) = tokio::join!(
+                async move {
+                    self.client
+                        .connect(handshake_addr, acceptor_addr, server_name())
+                        .await
+                        .expect("connect")
+                },
+                async move {
+                    let (conn, _) = self.server.accept().await.expect("accept");
+                    conn
+                }
+            );
+
+            check_pair_addrs!(client, server);
+
+            // the client doesn't have a response from the server so it might not know its
+            // port yet
+            if self.protocol().is_udp() {
+                let acceptor_port = acceptor_addr.port();
+                let server_local_port = server.local_addr().unwrap().port();
+                let client_peer_port = client.peer_addr().unwrap().port();
+                assert!(
+                    [acceptor_port, server_local_port].contains(&client_peer_port),
+                    "acceptor_port={acceptor_port}, server_local_port={server_local_port}, client_peer_port={client_peer_port}"
+                );
+            } else {
+                check_pair_addrs!(server, client);
+            }
+
+            (client, server)
+        }
+
+        pub fn protocol(&self) -> Protocol {
+            self.protocol
+        }
+
+        #[allow(dead_code)]
+        pub fn split(
+            self,
+        ) -> (
+            ClientTokio<ClientProvider, NoopSubscriber>,
+            ServerTokio<ServerProvider, NoopSubscriber>,
+        ) {
+            (self.client, self.server)
+        }
+    }
+}
+
+pub mod tcp {
+    use super::Protocol;
+    use std::net::SocketAddr;
+    use tokio::net::{TcpListener, TcpStream};
+
+    pub type Stream = TcpStream;
+
+    pub struct Context {
+        acceptor: TcpListener,
+    }
+
+    impl Context {
+        pub async fn new() -> Self {
+            Self::bind("127.0.0.1:0".parse().unwrap()).await
+        }
+
+        pub async fn bind(addr: SocketAddr) -> Self {
+            let acceptor = TcpListener::bind(addr).await.expect("bind");
+            Self { acceptor }
+        }
+
+        pub fn acceptor_addr(&self) -> SocketAddr {
+            self.acceptor.local_addr().expect("acceptor_addr")
+        }
+
+        pub async fn pair(&self) -> (Stream, Stream) {
+            self.pair_with(self.acceptor_addr()).await
+        }
+
+        pub async fn pair_with(&self, acceptor_addr: SocketAddr) -> (Stream, Stream) {
+            let (client, server) = tokio::join!(
+                async move { Stream::connect(acceptor_addr).await.expect("connect") },
+                async move {
+                    let (conn, _addr) = self.acceptor.accept().await.expect("accept");
+                    conn
+                }
+            );
+
+            check_pair_addrs!(client, server);
+            check_pair_addrs!(server, client);
+
+            (client, server)
+        }
+
+        pub fn protocol(&self) -> Protocol {
+            Protocol::Tcp
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct Client {
     map: secret::Map,

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

@@ -2,8 +2,13 @@
 // SPDX-License-Identifier: Apache-2.0
 
 mod accept_queue;
+mod behavior;
+/// A set of tests ensuring we support a large number of peers.
+#[cfg(future)] // TODO remove this since they're quite expensive
+mod cardinality;
 mod deterministic;
 mod idle_timeout;
 mod key_update;
 mod request_response;
+mod restart;
 mod rpc;