feat(bench): add generic shutdown functionality (#5426)

jmayclin · web-flow · commit 42a5d6de73e1 · 2025-07-22T12:35:07.000-07:00
diff --git a/bindings/rust/standard/bench/benches/handshake.rs b/bindings/rust/standard/bench/benches/handshake.rs
@@ -2,8 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use bench::{
-    harness::TlsBenchConfig, CipherSuite, CryptoConfig, HandshakeType, KXGroup, Mode,
-    OpenSslConnection, RustlsConnection, S2NConnection, SigType, TlsConnPair, TlsConnection,
+    harness::{TlsBenchConfig, TlsInfo},
+    CipherSuite, CryptoConfig, HandshakeType, KXGroup, Mode, OpenSslConnection, RustlsConnection,
+    S2NConnection, SigType, TlsConnPair, TlsConnection,
 };
 use criterion::{
     criterion_group, criterion_main, measurement::WallTime, BatchSize, BenchmarkGroup, Criterion,
@@ -16,7 +17,7 @@ fn bench_handshake_for_library<T>(
     kx_group: KXGroup,
     sig_type: SigType,
 ) where
-    T: TlsConnection,
+    T: TlsConnection + TlsInfo,
     T::Config: TlsBenchConfig,
 {
     let crypto_config = CryptoConfig::new(CipherSuite::default(), kx_group, sig_type);
diff --git a/bindings/rust/standard/bench/benches/resumption.rs b/bindings/rust/standard/bench/benches/resumption.rs
@@ -2,16 +2,17 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use bench::{
-    harness::TlsBenchConfig, CipherSuite, CryptoConfig, HandshakeType, KXGroup, S2NConnection,
-    SigType, TlsConnPair, TlsConnection,
+    harness::{TlsBenchConfig, TlsInfo},
+    CipherSuite, CryptoConfig, HandshakeType, KXGroup, S2NConnection, SigType, TlsConnPair,
+    TlsConnection,
 };
 use criterion::{
     criterion_group, criterion_main, measurement::WallTime, BatchSize, BenchmarkGroup, Criterion,
 };
 
 fn bench_handshake_pair<T>(bench_group: &mut BenchmarkGroup<WallTime>, sig_type: SigType)
 where
-    T: TlsConnection,
+    T: TlsConnection + TlsInfo,
     T::Config: TlsBenchConfig,
 {
     // generate all harnesses (TlsConnPair structs) beforehand so that benchmarks
@@ -38,7 +39,7 @@ where
 
 fn bench_handshake_server_1rtt<T>(bench_group: &mut BenchmarkGroup<WallTime>, sig_type: SigType)
 where
-    T: TlsConnection,
+    T: TlsConnection + TlsInfo,
     T::Config: TlsBenchConfig,
 {
     for handshake in [HandshakeType::Resumption, HandshakeType::ServerAuth] {
diff --git a/bindings/rust/standard/bench/benches/throughput.rs b/bindings/rust/standard/bench/benches/throughput.rs
@@ -2,8 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use bench::{
-    harness::TlsBenchConfig, CipherSuite, CryptoConfig, HandshakeType, KXGroup, Mode,
-    OpenSslConnection, RustlsConnection, S2NConnection, SigType, TlsConnPair, TlsConnection,
+    harness::{TlsBenchConfig, TlsInfo},
+    CipherSuite, CryptoConfig, HandshakeType, KXGroup, Mode, OpenSslConnection, RustlsConnection,
+    S2NConnection, SigType, TlsConnPair, TlsConnection,
 };
 use criterion::{
     criterion_group, criterion_main, measurement::WallTime, BatchSize, BenchmarkGroup, Criterion,
@@ -16,7 +17,7 @@ fn bench_throughput_for_library<T>(
     shared_buf: &mut [u8],
     cipher_suite: CipherSuite,
 ) where
-    T: TlsConnection,
+    T: TlsConnection + TlsInfo,
     T::Config: TlsBenchConfig,
 {
     let crypto_config = CryptoConfig::new(cipher_suite, KXGroup::default(), SigType::default());
diff --git a/bindings/rust/standard/bench/src/harness/mod.rs b/bindings/rust/standard/bench/src/harness/mod.rs
@@ -151,9 +151,6 @@ pub trait TlsConnection: Sized {
     /// Library-specific config struct
     type Config;
 
-    /// Name of the connection type
-    fn name() -> String;
-
     /// Make connection from existing config and buffer
     fn new_from_config(
         mode: Mode,
@@ -166,23 +163,42 @@ pub trait TlsConnection: Sized {
 
     fn handshake_completed(&self) -> bool;
 
+    /// Send `data` to the peer.
+    fn send(&mut self, data: &[u8]) -> Result<(), Box<dyn Error>>;
+
+    /// Read application data from the peer into `data`.
+    fn recv(&mut self, data: &mut [u8]) -> Result<(), Box<dyn Error>>;
+
+    /// Send a `CloseNotify` to the peer.
+    ///
+    /// This does not read the `CloseNotify` from the peer.
+    ///
+    /// Must be followed by a call to [`TlsConnection::shutdown_finish`] to ensure
+    /// that any `CloseNotify` alerts from the peer are read.
+    fn shutdown_send(&mut self);
+
+    /// Attempt to read the `CloseNotify` from the peer.
+    ///
+    /// Returns `true` if the connection was successfully shutdown, `false` otherwise.
+    ///
+    /// The `CloseNotify` might already have been read by `shutdown_send`, depending
+    /// on the order of client/server [`TlsConnection::shutdown_send`] calls.
+    fn shutdown_finish(&mut self) -> bool;
+}
+
+pub trait TlsInfo: Sized {
+    fn name() -> String;
     fn get_negotiated_cipher_suite(&self) -> CipherSuite;
 
     fn negotiated_tls13(&self) -> bool;
 
     /// Describes whether a connection was resumed. This method is only valid on
     /// server connections because of rustls API limitations.
     fn resumed_connection(&self) -> bool;
-
-    /// Send application data to ConnectedBuffer
-    fn send(&mut self, data: &[u8]) -> Result<(), Box<dyn Error>>;
-
-    /// Read application data from ConnectedBuffer
-    fn recv(&mut self, data: &mut [u8]) -> Result<(), Box<dyn Error>>;
 }
 
 /// A TlsConnPair owns the client and server tls connections along with the IO buffers.
-pub struct TlsConnPair<C: TlsConnection, S: TlsConnection> {
+pub struct TlsConnPair<C, S> {
     pub client: C,
     pub server: S,
     pub io: TestPairIO,
@@ -295,18 +311,6 @@ where
         self.client.handshake_completed() && self.server.handshake_completed()
     }
 
-    pub fn get_negotiated_cipher_suite(&self) -> CipherSuite {
-        assert!(self.handshake_completed());
-        assert!(
-            self.client.get_negotiated_cipher_suite() == self.server.get_negotiated_cipher_suite()
-        );
-        self.client.get_negotiated_cipher_suite()
-    }
-
-    pub fn negotiated_tls13(&self) -> bool {
-        self.client.negotiated_tls13() && self.server.negotiated_tls13()
-    }
-
     /// Send data from client to server, and then from server to client
     pub fn round_trip_transfer(&mut self, data: &mut [u8]) -> Result<(), Box<dyn Error>> {
         // send data from client to server
@@ -319,6 +323,45 @@ where
 
         Ok(())
     }
+
+    pub fn shutdown(&mut self) -> Result<(), Box<dyn Error>> {
+        // These assertions do not _have_ to be true, but you are likely making
+        // a mistake if you are hitting it. Generally all data should have been
+        // read before attempting to shutdown
+        assert_eq!(self.io.client_tx_stream.borrow().len(), 0);
+        assert_eq!(self.io.server_tx_stream.borrow().len(), 0);
+
+        self.client.shutdown_send();
+        self.server.shutdown_send();
+
+        let client_shutdown = self.client.shutdown_finish();
+        let server_shutdown = self.server.shutdown_finish();
+        if client_shutdown && server_shutdown {
+            Ok(())
+        } else {
+            Err(
+                format!("Shutdown Failed: client - {client_shutdown} server - {server_shutdown}")
+                    .into(),
+            )
+        }
+    }
+}
+
+impl<C, S> TlsConnPair<C, S>
+where
+    C: TlsInfo,
+    S: TlsInfo,
+{
+    pub fn get_negotiated_cipher_suite(&self) -> CipherSuite {
+        assert!(
+            self.client.get_negotiated_cipher_suite() == self.server.get_negotiated_cipher_suite()
+        );
+        self.client.get_negotiated_cipher_suite()
+    }
+
+    pub fn negotiated_tls13(&self) -> bool {
+        self.client.negotiated_tls13() && self.server.negotiated_tls13()
+    }
 }
 
 #[cfg(test)]
@@ -349,8 +392,8 @@ mod tests {
 
     fn test_type<C, S>()
     where
-        S: TlsConnection,
-        C: TlsConnection,
+        S: TlsConnection + TlsInfo,
+        C: TlsConnection + TlsInfo,
         C::Config: TlsBenchConfig,
         S::Config: TlsBenchConfig,
     {
@@ -361,8 +404,8 @@ mod tests {
 
     fn handshake_configs<C, S>()
     where
-        S: TlsConnection,
-        C: TlsConnection,
+        S: TlsConnection + TlsInfo,
+        C: TlsConnection + TlsInfo,
         C::Config: TlsBenchConfig,
         S::Config: TlsBenchConfig,
     {
@@ -381,6 +424,13 @@ mod tests {
 
                         assert!(conn_pair.negotiated_tls13());
                         assert_eq!(cipher_suite, conn_pair.get_negotiated_cipher_suite());
+
+                        // read in "application data" handshake messages.
+                        // "NewSessionTicket" in the case of resumption
+                        let err = conn_pair.client_mut().recv(&mut [0]).unwrap_err();
+                        assert_eq!(&err.to_string(), "blocking");
+
+                        conn_pair.shutdown().unwrap();
                     }
                 }
             }
@@ -389,8 +439,8 @@ mod tests {
 
     fn session_resumption<C, S>()
     where
-        S: TlsConnection,
-        C: TlsConnection,
+        S: TlsConnection + TlsInfo,
+        C: TlsConnection + TlsInfo,
         C::Config: TlsBenchConfig,
         S::Config: TlsBenchConfig,
     {
@@ -399,7 +449,12 @@ mod tests {
             TlsConnPair::<C, S>::new_bench_pair(CryptoConfig::default(), HandshakeType::Resumption)
                 .unwrap();
         conn_pair.handshake().unwrap();
+        // read the session tickets which were sent
+        let err = conn_pair.client_mut().recv(&mut [0]).unwrap_err();
+        assert_eq!(&err.to_string(), "blocking");
+
         assert!(conn_pair.server().resumed_connection());
+        conn_pair.shutdown().unwrap();
     }
 
     #[test]
@@ -439,6 +494,7 @@ mod tests {
                     .unwrap();
             conn_pair.handshake().unwrap();
             conn_pair.round_trip_transfer(&mut buf).unwrap();
+            conn_pair.shutdown().unwrap();
         }
     }
 }
diff --git a/bindings/rust/standard/bench/src/openssl.rs b/bindings/rust/standard/bench/src/openssl.rs
@@ -5,13 +5,13 @@ use crate::{
     get_cert_path,
     harness::{
         self, CipherSuite, CryptoConfig, HandshakeType, KXGroup, Mode, TlsBenchConfig,
-        TlsConnection, ViewIO,
+        TlsConnection, TlsInfo, ViewIO,
     },
     PemType::*,
 };
 use openssl::ssl::{
-    ErrorCode, Ssl, SslContext, SslFiletype, SslMethod, SslSession, SslSessionCacheMode, SslStream,
-    SslVerifyMode, SslVersion,
+    ErrorCode, ShutdownResult, Ssl, SslContext, SslFiletype, SslMethod, SslOptions, SslSession,
+    SslSessionCacheMode, SslStream, SslVerifyMode, SslVersion,
 };
 use std::{
     error::Error,
@@ -128,6 +128,13 @@ impl TlsBenchConfig for OpenSslConfig {
                 }
                 if handshake_type == HandshakeType::Resumption {
                     builder.set_session_cache_mode(SslSessionCacheMode::CLIENT);
+                } else {
+                    builder.set_options(SslOptions::NO_TICKET);
+                    builder.set_session_cache_mode(SslSessionCacheMode::OFF);
+                    // OpenSSL Bug: https://github.com/openssl/openssl/issues/8077
+                    // even with the above configuration, we must explicitly specify
+                    // 0 tickets
+                    builder.set_num_tickets(0)?;
                 }
             }
         }
@@ -141,21 +148,6 @@ impl TlsBenchConfig for OpenSslConfig {
 impl TlsConnection for OpenSslConnection {
     type Config = OpenSslConfig;
 
-    fn name() -> String {
-        let version_num = openssl::version::number() as u64;
-        let patch: u8 = (version_num >> 4) as u8;
-        let fix = (version_num >> 12) as u8;
-        let minor = (version_num >> 20) as u8;
-        let major = (version_num >> 28) as u8;
-        format!(
-            "openssl{}.{}.{}{}",
-            major,
-            minor,
-            fix,
-            (b'a' + patch - 1) as char
-        )
-    }
-
     fn new_from_config(
         mode: harness::Mode,
         config: &Self::Config,
@@ -212,6 +204,50 @@ impl TlsConnection for OpenSslConnection {
         self.connection.ssl().is_init_finished()
     }
 
+    fn send(&mut self, data: &[u8]) -> Result<(), Box<dyn Error>> {
+        let mut write_offset = 0;
+        while write_offset < data.len() {
+            write_offset += self.connection.write(&data[write_offset..data.len()])?;
+            self.connection.flush()?; // make sure internal buffers don't fill up
+        }
+        Ok(())
+    }
+
+    fn recv(&mut self, data: &mut [u8]) -> Result<(), Box<dyn Error>> {
+        let data_len = data.len();
+        let mut read_offset = 0;
+        while read_offset < data.len() {
+            read_offset += self.connection.read(&mut data[read_offset..data_len])?
+        }
+        Ok(())
+    }
+
+    fn shutdown_send(&mut self) {
+        // this method will not read in a CloseNotify
+        assert_eq!(self.connection.shutdown().unwrap(), ShutdownResult::Sent);
+    }
+
+    fn shutdown_finish(&mut self) -> bool {
+        self.connection.shutdown().unwrap() == ShutdownResult::Received
+    }
+}
+
+impl TlsInfo for OpenSslConnection {
+    fn name() -> String {
+        let version_num = openssl::version::number() as u64;
+        let patch: u8 = (version_num >> 4) as u8;
+        let fix = (version_num >> 12) as u8;
+        let minor = (version_num >> 20) as u8;
+        let major = (version_num >> 28) as u8;
+        format!(
+            "openssl{}.{}.{}{}",
+            major,
+            minor,
+            fix,
+            (b'a' + patch - 1) as char
+        )
+    }
+
     fn get_negotiated_cipher_suite(&self) -> CipherSuite {
         let cipher_suite = self
             .connection
@@ -234,24 +270,6 @@ impl TlsConnection for OpenSslConnection {
             == SslVersion::TLS1_3
     }
 
-    fn send(&mut self, data: &[u8]) -> Result<(), Box<dyn Error>> {
-        let mut write_offset = 0;
-        while write_offset < data.len() {
-            write_offset += self.connection.write(&data[write_offset..data.len()])?;
-            self.connection.flush()?; // make sure internal buffers don't fill up
-        }
-        Ok(())
-    }
-
-    fn recv(&mut self, data: &mut [u8]) -> Result<(), Box<dyn Error>> {
-        let data_len = data.len();
-        let mut read_offset = 0;
-        while read_offset < data.len() {
-            read_offset += self.connection.read(&mut data[read_offset..data_len])?
-        }
-        Ok(())
-    }
-
     fn resumed_connection(&self) -> bool {
         self.connection.ssl().session_reused()
     }
diff --git a/bindings/rust/standard/bench/src/rustls.rs b/bindings/rust/standard/bench/src/rustls.rs
diff --git a/bindings/rust/standard/bench/src/s2n_tls.rs b/bindings/rust/standard/bench/src/s2n_tls.rs
