test(integration): add async cert verify and offload 'stress' test

bindings/rust/standard/integration/Cargo.toml

@@ -22,7 +22,7 @@ pq = [ "s2n-tls/pq" ]
 s2n-tls = { path = "../../extended/s2n-tls", features = ["unstable-testing", "unstable-crl"]}
 s2n-tls-hyper = { path = "../s2n-tls-hyper" }
 s2n-tls-tokio = { path = "../../extended/s2n-tls-tokio" }
-s2n-tls-sys = { path = "../../extended/s2n-tls-sys" }
+s2n-tls-sys = { path = "../../extended/s2n-tls-sys", features = ["unstable-async_offload"] }
 aws-lc-sys = "*"
 tls-harness = { path = "../tls-harness" }
 

bindings/rust/standard/integration/src/mtls/async_verify_and_offload.rs

@@ -0,0 +1,140 @@
+// Copyright Amazon.com, Inc. or its affiliates.
+// SPDX-License-Identifier: Apache-2.0
+
+// mTLS async verify + async offload test
+// This is a "stress test" of our mTLS async callback which configures:
+// - rustls client (TLS 1.3)
+// - s2n-tls server
+// - async certificate validation callback
+// - async offload operation (pkey verify)
+
+use super::*;
+use s2n_tls_sys::{
+    s2n_async_offload_op, s2n_async_offload_op_perform, s2n_async_offload_op_type,
+    s2n_config_set_async_offload_callback,
+};
+use std::ffi::c_void;
+
+/// A wrapper around a raw pointer to `s2n_async_offload_op` that can be sent across threads.
+///
+/// This is used in tests to simulate async offload operations where the operation
+/// is deferred and performed on a different thread or after some async operation.
+struct SendableAsyncOffloadOp(*mut s2n_async_offload_op);
+
+// SAFETY: The pointer is owned by s2n-tls and remains valid for the duration of the
+// pending async offload operation (until perform() is called, or freed).
+// The test mimics the intended usage pattern where an application hands off the pointer
+// to a worker thread that later performs the operation.
+unsafe impl Send for SendableAsyncOffloadOp {}
+
+// Async offload context for C FFI
+struct AsyncOffloadCtx {
+    invoked: Arc<AtomicU64>,
+    sender: Sender<SendableAsyncOffloadOp>,
+}
+
+// C-style async offload callback
+unsafe extern "C" fn test_async_offload_cb(
+    _conn: *mut s2n_connection,
+    op: *mut s2n_async_offload_op,
+    ctx: *mut c_void,
+) -> i32 {
+    let ctx = unsafe { Box::from_raw(ctx as *mut AsyncOffloadCtx) };
+
+    ctx.invoked.fetch_add(1, Ordering::SeqCst);
+    ctx.sender.send(SendableAsyncOffloadOp(op)).unwrap();
+
+    s2n_status_code::SUCCESS
+}
+
+/// Registers an async pkey verify offload callback and returns (invoked_counter, operation_receiver).
+fn register_async_pkey_verify_offload(
+    s2n_cfg: &mut S2NConfig,
+) -> (Arc<AtomicU64>, Receiver<SendableAsyncOffloadOp>) {
+    let invoked = Arc::new(AtomicU64::new(0));
+    let (tx, rx) = std::sync::mpsc::channel();
+
+    let ctx = Box::new(AsyncOffloadCtx {
+        invoked: Arc::clone(&invoked),
+        sender: tx,
+    });
+    let ctx_ptr = Box::into_raw(ctx) as *mut c_void;
+
+    // SAFETY: s2n stores this context pointer and later returns it in the async
+    // callback. The callback will reclaim ownership using Box::from_raw to prevent leaks.
+    unsafe {
+        let raw = raw_config(s2n_cfg);
+        let allowed_types = s2n_async_offload_op_type::OFFLOAD_PKEY_VERIFY;
+
+        let result = s2n_config_set_async_offload_callback(
+            raw,
+            allowed_types,
+            Some(test_async_offload_cb),
+            ctx_ptr,
+        );
+        assert_eq!(
+            result,
+            s2n_status_code::SUCCESS,
+            "s2n_config_set_async_offload_callback failed"
+        );
+    }
+
+    (invoked, rx)
+}
+
+/// rustls client and s2n-tls server with async cert validation and async pkey
+/// verify offload over TLS 1.3.
+#[test]
+fn s2n_server_tls13() {
+    crate::capability_check::required_capability(
+        &[crate::capability_check::Capability::Tls13],
+        || {
+            let client = rustls_mtls_client(SigType::Rsa2048, &rustls::version::TLS13);
+
+            let (server, cert_invoked, cert_rx, offload_invoked, offload_rx) = {
+                let builder = s2n_mtls_base_builder(SigType::Rsa2048);
+                let mut s2n_cfg = S2NConfig::from(builder.build().unwrap());
+
+                let (cert_invoked, cert_rx) = register_async_cert_callback(&mut s2n_cfg);
+                let (offload_invoked, offload_rx) =
+                    register_async_pkey_verify_offload(&mut s2n_cfg);
+
+                (s2n_cfg, cert_invoked, cert_rx, offload_invoked, offload_rx)
+            };
+
+            let mut pair =
+                TlsConnPair::<RustlsConnection, S2NConnection>::from_configs(&client, &server);
+
+            // Drive early handshake messages (ClientHello, ServerHello, EncryptedExtensions)
+            pair.client.handshake().unwrap();
+            pair.server.handshake().unwrap();
+
+            // Progress until just before client Certificate is processed:
+            pair.client.handshake().unwrap();
+            assert_eq!(cert_invoked.load(Ordering::SeqCst), 0);
+            assert_eq!(offload_invoked.load(Ordering::SeqCst), 0);
+
+            // Server processes client Certificate -> async cert validation fires
+            pair.server.handshake().unwrap();
+            assert_eq!(cert_invoked.load(Ordering::SeqCst), 1);
+
+            let cert_ptr = cert_rx.recv().unwrap().0;
+            unsafe {
+                assert_eq!(s2n_cert_validation_accept(cert_ptr), 0);
+            }
+
+            // Continue server handshake: async offload for pkey verify fires
+            pair.server.handshake().unwrap();
+            assert_eq!(offload_invoked.load(Ordering::SeqCst), 1);
+
+            let SendableAsyncOffloadOp(offload_op_ptr) = offload_rx.recv().unwrap();
+            unsafe {
+                assert_eq!(s2n_async_offload_op_perform(offload_op_ptr), 0);
+            }
+
+            pair.handshake().unwrap();
+            pair.round_trip_assert(10).unwrap();
+            pair.shutdown().unwrap();
+        },
+    );
+}

bindings/rust/standard/integration/src/mtls.rs → bindings/rust/standard/integration/src/mtls/mod.rs

@@ -1,6 +1,10 @@
 // Copyright Amazon.com, Inc. or its affiliates.
 // SPDX-License-Identifier: Apache-2.0
 
+// async_verify_and_offload contains a “stress” test that wires the async cert-validation
+// and async pkey-verify callbacks together. It’s kept separate from the main mTLS test matrix.
+mod async_verify_and_offload;
+
 // This test suite exercises mTLS interoperability between s2n-tls and rustls,
 // including:
 //   - basic mTLS handshakes (TLS 1.2 and 1.3)
@@ -105,7 +109,7 @@ extern "C" fn test_async_cert_cb(
     info: *mut s2n_cert_validation_info,
     ctx: *mut c_void,
 ) -> i32 {
-    let ctx = unsafe { &*(ctx as *mut AsyncCertCtx) };
+    let ctx = unsafe { Box::from_raw(ctx as *mut AsyncCertCtx) };
 
     ctx.invoked.fetch_add(1, Ordering::SeqCst);
     ctx.sender
@@ -152,8 +156,7 @@ fn register_async_cert_callback(
     let ctx_ptr = Box::into_raw(ctx) as *mut c_void;
 
     // SAFETY: s2n stores this context pointer and later returns it in the async
-    // callback. Because s2n never frees it, we intentionally leak the Box so the
-    // memory stays valid for the lifetime of the config (test-only).
+    // callback. The callback will reclaim ownership using Box::from_raw to prevent leaks.
     unsafe {
         let raw = raw_config(s2n_cfg);
         let result = s2n_config_set_cert_validation_cb(raw, Some(test_async_cert_cb), ctx_ptr);