feat(bindings): Associate an application context with a Connection

bindings/rust/s2n-tls/src/connection.rs

@@ -23,7 +23,7 @@ use core::{
 };
 use libc::c_void;
 use s2n_tls_sys::*;
-use std::ffi::CStr;
+use std::{any::Any, ffi::CStr};
 
 mod builder;
 pub use builder::*;
@@ -1049,6 +1049,42 @@ impl Connection {
     pub fn resumed(&self) -> bool {
         unsafe { s2n_connection_is_session_resumed(self.connection.as_ptr()) == 1 }
     }
+
+    /// Associates an arbitrary application context 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.
+    pub fn set_application_context<T: Send + Sync + 'static>(&mut self, app_context: T) {
+        self.context_mut().app_context = Some(Box::new(app_context));
+    }
+
+    /// Retrieves a reference to the application context associated with the Connection.
+    ///
+    /// To set a context on the connection, use the [`Self::set_application_context()`] API. If an
+    /// application context of the specified type hasn't already been set on the Connection, None
+    /// will be returned.
+    ///
+    /// To retrieve a mutable reference to the context, use [`Self::application_context_mut()`].
+    pub fn application_context<T: Send + Sync + 'static>(&self) -> Option<&T> {
+        match self.context().app_context.as_ref() {
+            None => None,
+            Some(app_context) => app_context.downcast_ref::<T>(),
+        }
+    }
+
+    /// Retrieves a mutable reference to the application context associated with the Connection.
+    ///
+    /// To set a context on the connection, use the [`Self::set_application_context()`] API. If an
+    /// application context of the specified type hasn't already been set on the Connection, None
+    /// will be returned.
+    ///
+    /// To retrieve an immutable reference to the context, use [`Self::application_context()`].
+    pub fn application_context_mut<T: Send + Sync + 'static>(&mut self) -> Option<&mut T> {
+        match self.context_mut().app_context.as_mut() {
+            None => None,
+            Some(app_context) => app_context.downcast_mut::<T>(),
+        }
+    }
 }
 
 struct Context {
@@ -1057,6 +1093,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>>,
 }
 
 impl Context {
@@ -1067,6 +1104,7 @@ impl Context {
             async_callback: None,
             verify_host_callback: None,
             connection_initialized: false,
+            app_context: None,
         }
     }
 }

bindings/rust/s2n-tls/src/testing/s2n_tls.rs

@@ -233,7 +233,7 @@ impl<'a, T: 'a + Context> Callback<'a, T> {
 #[cfg(test)]
 mod tests {
     use crate::{
-        callbacks::{ClientHelloCallback, ConnectionFuture},
+        callbacks::{ClientHelloCallback, ConnectionFuture, ConnectionFutureResult},
         enums::ClientAuthType,
         error::ErrorType,
         testing::{client_hello::*, s2n_tls::*, *},
@@ -970,4 +970,162 @@ mod tests {
         init::init();
         assert!(init::fips_mode().unwrap().is_enabled());
     }
+
+    /// Test that an application context can be set and retrieved.
+    #[test]
+    fn test_app_context_set_and_retrieve() {
+        struct TestApplicationContext {
+            test_value: u32,
+        }
+
+        let mut connection = connection::Connection::new_server();
+        connection.set_application_context(TestApplicationContext { test_value: 1142 });
+
+        let context = connection
+            .application_context::<TestApplicationContext>()
+            .unwrap();
+        assert_eq!(context.test_value, 1142);
+    }
+
+    /// Test that an application context can be modified.
+    #[test]
+    fn test_app_context_modify() {
+        struct TestApplicationContext {
+            test_value: u32,
+        }
+
+        let mut connection = connection::Connection::new_server();
+        connection.set_application_context(TestApplicationContext { test_value: 0 });
+
+        let context = connection
+            .application_context_mut::<TestApplicationContext>()
+            .unwrap();
+        context.test_value += 1;
+
+        let context = connection
+            .application_context::<TestApplicationContext>()
+            .unwrap();
+        assert_eq!(context.test_value, 1);
+    }
+
+    /// Test that an application context can be overridden.
+    #[test]
+    fn test_app_context_override() {
+        struct TestApplicationContext {
+            test_value: u32,
+        }
+
+        let mut connection = connection::Connection::new_server();
+        connection.set_application_context(TestApplicationContext { test_value: 1142 });
+
+        {
+            let context = connection
+                .application_context::<TestApplicationContext>()
+                .unwrap();
+            assert_eq!(context.test_value, 1142);
+        }
+
+        // Override the context with a new value.
+        connection.set_application_context(TestApplicationContext { test_value: 10 });
+
+        {
+            let context = connection
+                .application_context::<TestApplicationContext>()
+                .unwrap();
+            assert_eq!(context.test_value, 10);
+        }
+
+        // Test overriding to a different context type.
+        struct AnotherApplicationContext {
+            another_test_value: i16,
+        }
+        connection.set_application_context(AnotherApplicationContext {
+            another_test_value: -20,
+        });
+
+        let context = connection
+            .application_context::<AnotherApplicationContext>()
+            .unwrap();
+        assert_eq!(context.another_test_value, -20);
+    }
+
+    /// Test that a context of another type can't be retrieved.
+    #[test]
+    fn test_app_context_invalid_type() {
+        struct TestApplicationContext {}
+        struct AnotherApplicationContext {}
+
+        let mut connection = connection::Connection::new_server();
+        connection.set_application_context(TestApplicationContext {});
+
+        // A context type that wasn't set shouldn't be returned.
+        assert!(connection
+            .application_context::<AnotherApplicationContext>()
+            .is_none());
+
+        // Retrieving the correct type succeeds.
+        assert!(connection
+            .application_context::<TestApplicationContext>()
+            .is_some());
+    }
+
+    /// Test that a context can be used from within a callback.
+    #[test]
+    fn test_app_context_callback() {
+        struct TestApplicationContext {
+            invoked_count: u32,
+        }
+
+        struct TestClientHelloHandler {}
+
+        impl ClientHelloCallback for TestClientHelloHandler {
+            fn on_client_hello(
+                &self,
+                connection: &mut connection::Connection,
+            ) -> ConnectionFutureResult {
+                let app_context = connection
+                    .application_context_mut::<TestApplicationContext>()
+                    .unwrap();
+                app_context.invoked_count += 1;
+                Ok(None)
+            }
+        }
+
+        let config = {
+            let keypair = CertKeyPair::default();
+            let mut builder = Builder::new();
+            builder
+                .set_verify_host_callback(InsecureAcceptAllCertificatesHandler {})
+                .unwrap();
+            builder
+                .set_client_hello_callback(TestClientHelloHandler {})
+                .unwrap();
+            builder.load_pem(keypair.cert, keypair.key).unwrap();
+            builder.trust_pem(keypair.cert).unwrap();
+            builder.build().unwrap()
+        };
+
+        let mut pair = tls_pair(config);
+        pair.server
+            .0
+            .connection_mut()
+            .set_waker(Some(&noop_waker()))
+            .unwrap();
+
+        let context = TestApplicationContext { invoked_count: 0 };
+        pair.server
+            .0
+            .connection_mut()
+            .set_application_context(context);
+
+        assert!(poll_tls_pair_result(&mut pair).is_ok());
+
+        let context = pair
+            .server
+            .0
+            .connection()
+            .application_context::<TestApplicationContext>()
+            .unwrap();
+        assert_eq!(context.invoked_count, 1);
+    }
 }