feat: add additional application context into Connection

boquan-fang · boquan-fang · commit adfff1adbf93 · 2025-11-25T19:19:14.000Z
diff --git a/bindings/rust/extended/s2n-tls/src/callbacks/cert_validation.rs b/bindings/rust/extended/s2n-tls/src/callbacks/cert_validation.rs
@@ -53,6 +53,7 @@ pub trait CertValidationCallbackSync: 'static + Send + Sync {
 mod tests {
     use super::*;
     use crate::{connection::Connection, security, testing::*};
+    use std::any::TypeId;
 
     struct ValidationContext {
         accept: bool,
@@ -66,7 +67,10 @@ mod tests {
             _info: &mut CertValidationInfo,
         ) -> Result<bool, Error> {
             self.0.increment();
-            let context = conn.application_context::<ValidationContext>().unwrap();
+            let application_context_type_id = TypeId::of::<ValidationContext>();
+            let context = conn
+                .application_context::<ValidationContext>(application_context_type_id)
+                .unwrap();
             Ok(context.accept)
         }
     }
diff --git a/bindings/rust/extended/s2n-tls/src/connection.rs b/bindings/rust/extended/s2n-tls/src/connection.rs
@@ -28,7 +28,11 @@ use core::{
 };
 use libc::c_void;
 use s2n_tls_sys::*;
-use std::{any::Any, ffi::CStr};
+use std::{
+    any::{Any, TypeId},
+    collections::HashMap,
+    ffi::CStr,
+};
 
 mod builder;
 pub use builder::*;
@@ -1408,27 +1412,41 @@ impl Connection {
         Ok(())
     }
 
-    /// Associates an arbitrary application context with the Connection to be later retrieved via
+    /// Associates arbitrary application contexts with the Connection to be later retrieved via
     /// the [`Self::application_context()`] and [`Self::application_context_mut()`] APIs.
     ///
-    /// This API will override an existing application context set on the Connection.
+    /// This API will add additional application context on top of existing contexts set on the Connection.
     ///
     /// Corresponds to [s2n_connection_set_ctx].
     pub fn set_application_context<T: Send + Sync + 'static>(&mut self, app_context: T) {
-        self.context_mut().app_context = Some(Box::new(app_context));
+        let context_type_id = TypeId::of::<T>();
+        self.context_mut()
+            .app_context
+            .insert(context_type_id, Box::new(app_context));
+    }
+
+    /// Remove a application context set on the Connection.
+    pub fn remove_application_context(
+        &mut self,
+        context_type_id: TypeId,
+    ) -> Option<Box<dyn Any + Send + Sync>> {
+        self.context_mut().app_context.remove(&context_type_id)
     }
 
     /// Retrieves a reference to the application context associated with the Connection.
     ///
-    /// If an application context hasn't already been set on the Connection, or if the set
-    /// application context isn't of type T, None will be returned.
+    /// If application context hasn't already been set on the Connection, or if the set
+    /// application context doesn't match the type id of type T, None will be returned.
     ///
     /// To set a context on the connection, use [`Self::set_application_context()`]. To retrieve a
     /// mutable reference to the context, use [`Self::application_context_mut()`].
     ///
     /// Corresponds to [s2n_connection_get_ctx].
-    pub fn application_context<T: Send + Sync + 'static>(&self) -> Option<&T> {
-        match self.context().app_context.as_ref() {
+    pub fn application_context<T: Send + Sync + 'static>(
+        &self,
+        context_type_id: TypeId,
+    ) -> Option<&T> {
+        match self.context().app_context.get(&context_type_id) {
             None => None,
             // The Any trait keeps track of the application context's type. downcast_ref() returns
             // Some only if the correct type is provided:
@@ -1439,15 +1457,18 @@ impl Connection {
 
     /// Retrieves a mutable reference to the application context associated with the Connection.
     ///
-    /// If an application context hasn't already been set on the Connection, or if the set
-    /// application context isn't of type T, None will be returned.
+    /// If application context hasn't already been set on the Connection, or if the set
+    /// application context doesn't match the type id of type T, None will be returned.
     ///
     /// To set a context on the connection, use [`Self::set_application_context()`]. To retrieve an
     /// immutable reference to the context, use [`Self::application_context()`].
     ///
     /// Corresponds to [s2n_connection_get_ctx].
-    pub fn application_context_mut<T: Send + Sync + 'static>(&mut self) -> Option<&mut T> {
-        match self.context_mut().app_context.as_mut() {
+    pub fn application_context_mut<T: Send + Sync + 'static>(
+        &mut self,
+        context_type_id: TypeId,
+    ) -> Option<&mut T> {
+        match self.context_mut().app_context.get_mut(&context_type_id) {
             None => None,
             Some(app_context) => app_context.downcast_mut::<T>(),
         }
@@ -1475,7 +1496,7 @@ struct Context {
     async_callback: Option<AsyncCallback>,
     verify_host_callback: Option<Box<dyn VerifyHostNameCallback>>,
     connection_initialized: bool,
-    app_context: Option<Box<dyn Any + Send + Sync>>,
+    app_context: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
     #[cfg(feature = "unstable-renegotiate")]
     pub(crate) renegotiate_state: RenegotiateState,
     #[cfg(feature = "unstable-cert_authorities")]
@@ -1490,7 +1511,7 @@ impl Context {
             async_callback: None,
             verify_host_callback: None,
             connection_initialized: false,
-            app_context: None,
+            app_context: HashMap::new(),
             #[cfg(feature = "unstable-renegotiate")]
             renegotiate_state: RenegotiateState::default(),
             #[cfg(feature = "unstable-cert_authorities")]
@@ -1602,6 +1623,7 @@ impl Drop for Connection {
 mod tests {
     use super::*;
     use crate::testing::{build_config, SniTestCerts, TestPair};
+    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
     // ensure the connection context is send
     #[test]
@@ -1622,14 +1644,23 @@ mod tests {
     fn test_app_context_set_and_retrieve() {
         let mut connection = Connection::new_server();
 
+        let test_value: u32 = 1142;
+        let test_context_type_id = TypeId::of::<u32>();
+
         // Before a context is set, None is returned.
-        assert!(connection.application_context::<u32>().is_none());
+        assert!(connection
+            .application_context::<u32>(test_context_type_id)
+            .is_none());
 
-        let test_value: u32 = 1142;
         connection.set_application_context(test_value);
 
         // After a context is set, the application data is returned.
-        assert_eq!(*connection.application_context::<u32>().unwrap(), 1142);
+        assert_eq!(
+            *connection
+                .application_context::<u32>(test_context_type_id)
+                .unwrap(),
+            1142
+        );
     }
 
     /// Test that an application context can be modified.
@@ -1640,33 +1671,62 @@ mod tests {
         let mut connection = Connection::new_server();
         connection.set_application_context(test_value);
 
-        let context_value = connection.application_context_mut::<u64>().unwrap();
+        let context_type_id = TypeId::of::<u64>();
+        let context_value = connection
+            .application_context_mut::<u64>(context_type_id)
+            .unwrap();
         *context_value += 1;
 
-        assert_eq!(*connection.application_context::<u64>().unwrap(), 1);
+        assert_eq!(
+            *connection
+                .application_context::<u64>(context_type_id)
+                .unwrap(),
+            1
+        );
     }
 
-    /// Test that an application context can be overridden.
+    /// Test that multiple application contexts can be set in a connection
     #[test]
-    fn test_app_context_override() {
+    fn test_multiple_app_contexts() {
         let mut connection = Connection::new_server();
 
-        let test_value: u16 = 1142;
-        connection.set_application_context(test_value);
-
-        assert_eq!(*connection.application_context::<u16>().unwrap(), 1142);
-
-        // Override the context with a new value.
-        let test_value: u16 = 10;
-        connection.set_application_context(test_value);
-
-        assert_eq!(*connection.application_context::<u16>().unwrap(), 10);
-
-        // Override the context with a new type.
-        let test_value: i16 = -20;
-        connection.set_application_context(test_value);
-
-        assert_eq!(*connection.application_context::<i16>().unwrap(), -20);
+        let first_test_value: u16 = 1142;
+        connection.set_application_context(first_test_value);
+
+        let first_test_value_type_id = TypeId::of::<u16>();
+        assert_eq!(
+            *connection
+                .application_context::<u16>(first_test_value_type_id)
+                .unwrap(),
+            1142
+        );
+
+        // Insert the second application context to the connection
+        let second_test_value: SocketAddr =
+            SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);
+        connection.set_application_context(second_test_value);
+
+        let second_test_value_type_id = TypeId::of::<SocketAddr>();
+        assert_eq!(
+            *connection
+                .application_context::<SocketAddr>(second_test_value_type_id)
+                .unwrap(),
+            second_test_value
+        );
+
+        // Remove the second application context
+        assert_eq!(
+            second_test_value,
+            *connection
+                .remove_application_context(second_test_value_type_id)
+                .unwrap()
+                .downcast::<SocketAddr>()
+                .unwrap()
+        );
+
+        assert!(connection
+            .application_context::<SocketAddr>(second_test_value_type_id)
+            .is_none());
     }
 
     /// Test that a context of another type can't be retrieved.
@@ -1677,11 +1737,17 @@ mod tests {
         let test_value: u32 = 0;
         connection.set_application_context(test_value);
 
+        let invalid_context_type_id = TypeId::of::<i16>();
         // A context type that wasn't set shouldn't be returned.
-        assert!(connection.application_context::<i16>().is_none());
+        assert!(connection
+            .application_context::<i16>(invalid_context_type_id)
+            .is_none());
 
+        let valid_context_type_id = TypeId::of::<u32>();
         // Retrieving the correct type succeeds.
-        assert!(connection.application_context::<u32>().is_some());
+        assert!(connection
+            .application_context::<u32>(valid_context_type_id)
+            .is_some());
     }
 
     /// Test that the `certificate_match` Rust wrapper returns expected enum variant
diff --git a/bindings/rust/extended/s2n-tls/src/renegotiate.rs b/bindings/rust/extended/s2n-tls/src/renegotiate.rs
@@ -440,6 +440,7 @@ mod tests {
         ErrorCode, Ssl, SslContext, SslFiletype, SslMethod, SslStream, SslVerifyMode, SslVersion,
     };
     use std::{
+        any::TypeId,
         error::Error,
         io::{Read, Write},
         pin::Pin,
@@ -1087,7 +1088,10 @@ mod tests {
                 ctx.waker().wake_by_ref();
                 let this = self.get_mut();
                 // Assert that nothing is currently set
-                assert!(conn.application_context::<String>().is_none());
+                let application_context_type_id = TypeId::of::<String>();
+                assert!(conn
+                    .application_context::<String>(application_context_type_id)
+                    .is_none());
                 if this.count > 1 {
                     // Repeatedly block the handshake in order to verify
                     // that renegotiate can handle Pending callbacks.
@@ -1122,7 +1126,8 @@ mod tests {
         pair.assert_renegotiate()?;
         assert_eq!(wake_count, count_per_handshake * 2);
 
-        let context: Option<&String> = pair.client.application_context();
+        let application_context_type_id = TypeId::of::<String>();
+        let context: Option<&String> = pair.client.application_context(application_context_type_id);
         assert_eq!(Some(&expected_context), context);
 
         Ok(())
diff --git a/bindings/rust/extended/s2n-tls/src/testing/s2n_tls.rs b/bindings/rust/extended/s2n-tls/src/testing/s2n_tls.rs
@@ -13,7 +13,7 @@ mod tests {
     use core::sync::atomic::Ordering;
     use futures_test::task::{new_count_waker, noop_waker};
     use security::Policy;
-    use std::{fs, path::Path, pin::Pin, sync::atomic::AtomicUsize};
+    use std::{any::TypeId, fs, path::Path, pin::Pin, sync::atomic::AtomicUsize};
 
     #[test]
     fn handshake_default() {
@@ -634,8 +634,9 @@ mod tests {
                 &self,
                 connection: &mut connection::Connection,
             ) -> ConnectionFutureResult {
+                let context_type_id = TypeId::of::<TestApplicationContext>();
                 let app_context = connection
-                    .application_context_mut::<TestApplicationContext>()
+                    .application_context_mut::<TestApplicationContext>(context_type_id)
                     .unwrap();
                 app_context.invoked_count += 1;
                 Ok(None)
@@ -663,9 +664,10 @@ mod tests {
 
         pair.handshake()?;
 
+        let context_type_id = TypeId::of::<TestApplicationContext>();
         let context = pair
             .server
-            .application_context::<TestApplicationContext>()
+            .application_context::<TestApplicationContext>(context_type_id)
             .unwrap();
         assert_eq!(context.invoked_count, 1);
 
diff --git a/bindings/rust/standard/integration/tests/memory.rs b/bindings/rust/standard/integration/tests/memory.rs
@@ -274,12 +274,12 @@ mod memory_test {
         /// lifecycle. The static memory row is an absolute measurement, not a diff.
         fn assert_expected(&self) {
             const EXPECTED_MEMORY: &[(Lifecycle, usize)] = &[
-                (Lifecycle::ConnectionInit, 61_482),
-                (Lifecycle::AfterClientHello, 88_302),
-                (Lifecycle::AfterServerHello, 116_669),
-                (Lifecycle::AfterClientFinished, 107_976),
-                (Lifecycle::HandshakeComplete, 90_563),
-                (Lifecycle::ApplicationData, 90_563),
+                (Lifecycle::ConnectionInit, 61_578),
+                (Lifecycle::AfterClientHello, 88_406),
+                (Lifecycle::AfterServerHello, 116_773),
+                (Lifecycle::AfterClientFinished, 108_080),
+                (Lifecycle::HandshakeComplete, 90_667),
+                (Lifecycle::ApplicationData, 90_667),
             ];
             let actual_memory: Vec<(Lifecycle, usize)> = Lifecycle::all_stages()
                 .into_iter()
