aws-crt 0.1.4 → 0.1.5

data/aws-crt-ffi/crt/s2n/bindings/rust/s2n-tls/src/raw/config.rs

@@ -2,70 +2,159 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::raw::{
+    connection::Connection,
     error::{Error, Fallible},
     security,
 };
-use alloc::sync::Arc;
-use core::{convert::TryInto, ptr::NonNull};
+use core::{convert::TryInto, ptr::NonNull, task::Poll};
 use s2n_tls_sys::*;
-use std::ffi::CString;
+use std::{
+    ffi::{c_void, CStr, CString},
+    mem::ManuallyDrop,
+    sync::atomic::{AtomicUsize, Ordering},
+};
 
-struct Owned(NonNull<s2n_config>);
+pub struct Config(NonNull<s2n_config>);
 
+/// # Safety
+///
 /// Safety: s2n_config objects can be sent across threads
-unsafe impl Send for Owned {}
+unsafe impl Send for Config {}
 
-impl Default for Owned {
-    fn default() -> Self {
-        Self::new()
+impl Config {
+    /// Returns a Config object with pre-defined defaults.
+    ///
+    /// Use the [`Builder`] if custom configuration is desired.
+    pub fn new() -> Self {
+        Self::default()
     }
-}
 
-impl Owned {
-    fn new() -> Self {
-        crate::raw::init::init();
-        let config = unsafe { s2n_config_new().into_result() }.unwrap();
+    /// Returns a Builder which can be used to configure the Config
+    pub fn builder() -> Builder {
+        Builder::default()
+    }
+
+    /// # Safety
+    ///
+    /// This config _MUST_ have been initialized with a [`Builder`].
+    pub(crate) unsafe fn from_raw(config: NonNull<s2n_config>) -> Self {
         Self(config)
     }
 
     pub(crate) fn as_mut_ptr(&mut self) -> *mut s2n_config {
         self.0.as_ptr()
     }
-}
 
-impl Drop for Owned {
-    fn drop(&mut self) {
-        let _ = unsafe { s2n_config_free(self.0.as_ptr()).into_result() };
+    /// Retrieve a reference to the [`Context`] stored on the config.
+    ///
+    /// # Safety
+    ///
+    /// This function assumes that this Config was created using a [`Builder`]
+    /// and has a [`Context`] stored on its context.
+    unsafe fn context(&self) -> &Context {
+        let mut ctx = core::ptr::null_mut();
+        s2n_config_get_ctx(self.0.as_ptr(), &mut ctx)
+            .into_result()
+            .unwrap();
+        &*(ctx as *const Context)
     }
-}
 
-#[derive(Clone, Default)]
-pub struct Config(Arc<Owned>);
+    /// Retrieve a mutable reference to the [`Context`] stored on the config.
+    ///
+    /// # Safety
+    ///
+    /// This function assumes that this Config was created using a [`Builder`]
+    /// and has a [`Context`] stored on its context.
+    unsafe fn context_mut(&mut self) -> &mut Context {
+        let mut ctx = core::ptr::null_mut();
+        s2n_config_get_ctx(self.as_mut_ptr(), &mut ctx)
+            .into_result()
+            .unwrap();
+        &mut *(ctx as *mut Context)
+    }
 
-/// Safety: s2n_config objects can be sent across threads
-#[allow(unknown_lints, clippy::non_send_fields_in_send_ty)]
-unsafe impl Send for Config {}
+    #[cfg(test)]
+    /// Get the refcount associated with the config
+    pub fn test_get_refcount(&self) -> Result<usize, Error> {
+        let context = unsafe { self.context() };
+        Ok(context.refcount.load(Ordering::SeqCst))
+    }
+}
 
-impl Config {
-    pub fn new() -> Self {
-        Self::default()
+impl Default for Config {
+    fn default() -> Self {
+        Builder::new().build().unwrap()
     }
+}
 
-    pub fn builder() -> Builder {
-        Builder::default()
+impl Clone for Config {
+    fn clone(&self) -> Self {
+        let context = unsafe { self.context() };
+
+        // Safety
+        //
+        // Using a relaxed ordering is alright here, as knowledge of the
+        // original reference prevents other threads from erroneously deleting
+        // the object.
+        // https://github.com/rust-lang/rust/blob/e012a191d768adeda1ee36a99ef8b92d51920154/library/alloc/src/sync.rs#L1329
+        let _count = context.refcount.fetch_add(1, Ordering::Relaxed);
+        Self(self.0)
     }
+}
 
-    pub(crate) fn as_mut_ptr(&mut self) -> *mut s2n_config {
-        (self.0).0.as_ptr()
+impl Drop for Config {
+    fn drop(&mut self) {
+        let context = unsafe { self.context_mut() };
+        let count = context.refcount.fetch_sub(1, Ordering::Release);
+        debug_assert!(count > 0, "refcount should not drop below 1 instance");
+
+        // only free the config if this is the last instance
+        if count != 1 {
+            return;
+        }
+
+        // Safety
+        //
+        // The use of Ordering and fence mirrors the `Arc` implementation in
+        // the standard library.
+        //
+        // This fence is needed to prevent reordering of use of the data and
+        // deletion of the data.  Because it is marked `Release`, the decreasing
+        // of the reference count synchronizes with this `Acquire` fence. This
+        // means that use of the data happens before decreasing the reference
+        // count, which happens before this fence, which happens before the
+        // deletion of the data.
+        // https://github.com/rust-lang/rust/blob/e012a191d768adeda1ee36a99ef8b92d51920154/library/alloc/src/sync.rs#L1637
+        std::sync::atomic::fence(Ordering::Acquire);
+
+        unsafe {
+            // This is the last instance so free the context.
+            let context = Box::from_raw(context);
+            drop(context);
+
+            let _ = s2n_config_free(self.0.as_ptr()).into_result();
+        }
     }
 }
 
 #[derive(Default)]
-pub struct Builder(Owned);
+pub struct Builder(Config);
 
 impl Builder {
     pub fn new() -> Self {
-        Default::default()
+        crate::raw::init::init();
+        let config = unsafe { s2n_config_new().into_result() }.unwrap();
+
+        let context = Box::new(Context::default());
+        let context = Box::into_raw(context) as *mut c_void;
+
+        unsafe {
+            s2n_config_set_ctx(config.as_ptr(), context)
+                .into_result()
+                .unwrap();
+        }
+
+        Self(Config(config))
     }
 
     pub fn set_alert_behavior(
@@ -156,15 +245,57 @@ impl Builder {
         Ok(self)
     }
 
-    /// # Safety
+    pub fn wipe_trust_store(&mut self) -> Result<&mut Self, Error> {
+        unsafe { s2n_config_wipe_trust_store(self.as_mut_ptr()).into_result()? };
+        Ok(self)
+    }
+
+    /// Sets whether or not a client certificate should be required to complete the TLS connection.
     ///
-    /// The `context` pointer must live at least as long as the config
-    pub unsafe fn set_verify_host_callback(
+    /// See the [Usage Guide](https://github.com/aws/s2n-tls/blob/main/docs/USAGE-GUIDE.md#client-auth-related-calls) for more details.
+    pub fn set_client_auth_type(
         &mut self,
-        callback: s2n_verify_host_fn,
-        context: *mut core::ffi::c_void,
+        auth_type: s2n_cert_auth_type::Type,
     ) -> Result<&mut Self, Error> {
-        s2n_config_set_verify_host_callback(self.as_mut_ptr(), callback, context).into_result()?;
+        unsafe { s2n_config_set_client_auth_type(self.as_mut_ptr(), auth_type).into_result() }?;
+        Ok(self)
+    }
+
+    /// Set a custom callback function which is run during client certificate validation during
+    /// a mutual TLS handshake.
+    ///
+    /// The callback may be called more than once during certificate validation as each SAN on
+    /// the certificate will be checked.
+    pub fn set_verify_host_handler<T: 'static + VerifyClientCertificateHandler>(
+        &mut self,
+        handler: T,
+    ) -> Result<&mut Self, Error> {
+        unsafe extern "C" fn verify_host_cb(
+            host_name: *const ::libc::c_char,
+            host_name_len: usize,
+            context: *mut ::libc::c_void,
+        ) -> u8 {
+            let host_name = host_name as *const u8;
+            let host_name = core::slice::from_raw_parts(host_name, host_name_len);
+            if let Ok(host_name_str) = core::str::from_utf8(host_name) {
+                let context = &mut *(context as *mut Context);
+                let handler = context.verify_host_handler.as_mut().unwrap();
+                return handler.verify_host_name(host_name_str) as u8;
+            }
+            0 // If the host name can't be parsed, fail closed.
+        }
+
+        let handler = Box::new(handler);
+        let context = unsafe { self.0.context_mut() };
+        context.verify_host_handler = Some(handler);
+        unsafe {
+            s2n_config_set_verify_host_callback(
+                self.as_mut_ptr(),
+                Some(verify_host_cb),
+                self.0.context_mut() as *mut _ as *mut c_void,
+            )
+            .into_result()?;
+        }
         Ok(self)
     }
 
@@ -185,8 +316,68 @@ impl Builder {
         Ok(self)
     }
 
+    /// Set a custom callback function which is run after parsing the client hello.
+    ///
+    /// The callback can be called more than once. The application is response for calling
+    /// [`Connection::server_name_extension_used()`] if connection properties are modified
+    /// on the server name extension.
+    pub fn set_client_hello_handler<T: 'static + ClientHelloHandler>(
+        &mut self,
+        handler: T,
+    ) -> Result<&mut Self, Error> {
+        unsafe extern "C" fn client_hello_cb(
+            connection_ptr: *mut s2n_connection,
+            context: *mut core::ffi::c_void,
+        ) -> libc::c_int {
+            let context = &mut *(context as *mut Context);
+            let handler = context.client_hello_handler.as_mut().unwrap();
+
+            let connection_ptr =
+                NonNull::new(connection_ptr).expect("connection should not be null");
+
+            // Since this is a callback, which receives a pointer to the connection, do
+            // not call drop on the connection object.
+            let mut connection = ManuallyDrop::new(Connection::from_raw(connection_ptr));
+
+            match handler.poll_client_hello(&mut connection) {
+                Poll::Ready(Ok(())) => {
+                    s2n_client_hello_cb_done(connection_ptr.as_ptr());
+                    s2n_status_code::SUCCESS
+                }
+                Poll::Ready(Err(_)) => {
+                    s2n_client_hello_cb_done(connection_ptr.as_ptr());
+                    s2n_status_code::FAILURE
+                }
+                Poll::Pending => s2n_status_code::SUCCESS,
+            }
+        }
+
+        let handler = Box::new(handler);
+        let context = unsafe { self.0.context_mut() };
+        context.client_hello_handler = Some(handler);
+
+        unsafe {
+            // Enable client_hello_callback polling to model the polling behavior of Futures in
+            // Rust
+            s2n_config_client_hello_cb_enable_poll(self.as_mut_ptr()).into_result()?;
+            s2n_config_set_client_hello_cb_mode(
+                self.as_mut_ptr(),
+                s2n_client_hello_cb_mode::NONBLOCKING,
+            )
+            .into_result()?;
+            s2n_config_set_client_hello_cb(
+                self.as_mut_ptr(),
+                Some(client_hello_cb),
+                self.0.context_mut() as *mut _ as *mut c_void,
+            )
+            .into_result()?;
+        }
+
+        Ok(self)
+    }
+
     pub fn build(self) -> Result<Config, Error> {
-        Ok(Config(Arc::new(self.0)))
+        Ok(self.0)
     }
 
     fn as_mut_ptr(&mut self) -> *mut s2n_config {
@@ -201,3 +392,38 @@ impl Builder {
         Ok(self)
     }
 }
+
+pub(crate) struct Context {
+    refcount: AtomicUsize,
+    client_hello_handler: Option<Box<dyn ClientHelloHandler>>,
+    verify_host_handler: Option<Box<dyn VerifyClientCertificateHandler>>,
+}
+
+impl Default for Context {
+    fn default() -> Self {
+        // The AtomicUsize is used to manually track the reference count of the Config.
+        // This mechanism is used to track when the Config object should be freed.
+        let refcount = AtomicUsize::new(1);
+
+        Self {
+            refcount,
+            client_hello_handler: None,
+            verify_host_handler: None,
+        }
+    }
+}
+
+/// This trait represents the client_hello callback which is run after parsing the client_hello.
+///
+/// Use in conjunction with [`Builder::set_client_hello_handler()`].
+pub trait ClientHelloHandler: Send + Sync {
+    fn poll_client_hello(&self, connection: &mut Connection) -> core::task::Poll<Result<(), ()>>;
+}
+
+/// Trait which a user must implement to verify host name(s) during X509 verification when using
+/// mutual TLS.
+pub trait VerifyClientCertificateHandler: Send + Sync {
+    /// The implementation shall verify the host name by returning `true` if the certificate host name is valid,
+    /// and `false` otherwise.
+    fn verify_host_name(&self, _host_name: &str) -> bool;
+}

data/aws-crt-ffi/crt/s2n/bindings/rust/s2n-tls/src/raw/connection.rs

@@ -8,17 +8,20 @@ use crate::raw::{
     error::{Error, Fallible},
     security,
 };
-use core::{convert::TryInto, fmt, ptr::NonNull, task::Poll};
+use core::{
+    convert::TryInto,
+    fmt,
+    ptr::NonNull,
+    task::{Poll, Waker},
+};
 use libc::c_void;
 use s2n_tls_sys::*;
+use std::{ffi::CStr, mem};
 
 pub use s2n_tls_sys::s2n_mode;
 
 pub struct Connection {
     connection: NonNull<s2n_connection>,
-    // The config needs to be stored so the reference count is accurate
-    #[allow(dead_code)]
-    config: Option<Config>,
 }
 
 impl fmt::Debug for Connection {
@@ -29,6 +32,8 @@ impl fmt::Debug for Connection {
     }
 }
 
+/// # Safety
+///
 /// Safety: s2n_connection objects can be sent across threads
 unsafe impl Send for Connection {}
 
@@ -38,16 +43,24 @@ impl Connection {
         let connection = unsafe { s2n_connection_new(mode).into_result() }.unwrap();
 
         unsafe {
-            let mut config = core::ptr::null_mut();
             debug_assert! {
-                s2n_connection_get_config(connection.as_ptr(), &mut config).into_result().is_err()
-            };
+                s2n_connection_get_config(connection.as_ptr(), &mut core::ptr::null_mut())
+                    .into_result()
+                    .is_err()
+            }
         }
-        Self {
-            connection,
-            config: None,
+        let context = Box::new(Context::default());
+        let context = Box::into_raw(context) as *mut c_void;
+        // allocate a new context object
+        unsafe {
+            s2n_connection_set_ctx(connection.as_ptr(), context)
+                .into_result()
+                .unwrap();
         }
+
+        Self { connection }
     }
+
     pub fn new_client() -> Self {
         Self::new(s2n_mode::CLIENT)
     }
@@ -56,6 +69,13 @@ impl Connection {
         Self::new(s2n_mode::SERVER)
     }
 
+    /// # Safety
+    ///
+    /// Caller must ensure s2n_connection is a valid reference to a [`s2n_connection`] object
+    pub(crate) unsafe fn from_raw(connection: NonNull<s2n_connection>) -> Self {
+        Self { connection }
+    }
+
     /// can be used to configure s2n to either use built-in blinding (set blinding
     /// to S2N_BUILT_IN_BLINDING) or self-service blinding (set blinding to
     /// S2N_SELF_SERVICE_BLINDING).
@@ -80,18 +100,49 @@ impl Connection {
         Ok(self)
     }
 
+    /// Attempts to drop the config on the connection.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure the config associated with the connection was created
+    /// with a [`config::Builder`].
+    unsafe fn drop_config(&mut self) -> Result<(), Error> {
+        let mut prev_config = core::ptr::null_mut();
+
+        // A valid non-null pointer is returned only if the application previously called
+        // [`Self::set_config()`].
+        if s2n_connection_get_config(self.connection.as_ptr(), &mut prev_config)
+            .into_result()
+            .is_ok()
+        {
+            let prev_config = NonNull::new(prev_config).expect(
+                "config should exist since the call to s2n_connection_get_config was successful",
+            );
+            drop(Config::from_raw(prev_config));
+        }
+
+        Ok(())
+    }
+
     /// Associates a configuration object with a connection.
     pub fn set_config(&mut self, mut config: Config) -> Result<&mut Self, Error> {
         unsafe {
-            s2n_connection_set_config(self.connection.as_ptr(), config.as_mut_ptr()).into_result()
-        }?;
-        unsafe {
-            let mut config = core::ptr::null_mut();
+            // attempt to drop the currently set config
+            self.drop_config()?;
+
+            s2n_connection_set_config(self.connection.as_ptr(), config.as_mut_ptr())
+                .into_result()?;
+
             debug_assert! {
-                s2n_connection_get_config(self.connection.as_ptr(), &mut config).into_result().is_ok()
+                s2n_connection_get_config(self.connection.as_ptr(), &mut core::ptr::null_mut()).into_result().is_ok(),
+                "s2n_connection_set_config was successful"
             };
+
+            // Setting the config on the connection creates one additional reference to the config
+            // so do not drop so prevent Rust from calling `drop()` at the end of this function.
+            mem::forget(config);
         }
-        self.config = Some(config);
+
         Ok(self)
     }
 
@@ -251,11 +302,98 @@ impl Connection {
     }
 
     /// Sets the server name value for the connection
-    pub fn set_server_name(&mut self, sni: &[u8]) -> Result<&mut Self, Error> {
-        let sni = std::ffi::CString::new(sni).map_err(|_| Error::InvalidInput)?;
-        unsafe { s2n_set_server_name(self.connection.as_ptr(), sni.as_ptr()).into_result() }?;
+    pub fn set_server_name(&mut self, server_name: &str) -> Result<&mut Self, Error> {
+        let server_name = std::ffi::CString::new(server_name).map_err(|_| Error::InvalidInput)?;
+        unsafe {
+            s2n_set_server_name(self.connection.as_ptr(), server_name.as_ptr()).into_result()
+        }?;
+        Ok(self)
+    }
+
+    /// Get the server name associated with the connection client hello.
+    pub fn server_name(&self) -> Option<&str> {
+        unsafe {
+            let server_name = s2n_get_server_name(self.connection.as_ptr());
+            match server_name.into_result() {
+                Ok(server_name) => CStr::from_ptr(server_name).to_str().ok(),
+                Err(_) => None,
+            }
+        }
+    }
+
+    /// Sets a Waker on the connection context or clears it if `None` is passed.
+    pub fn set_waker(&mut self, waker: Option<&Waker>) -> Result<&mut Self, Error> {
+        let ctx = self.context_mut();
+
+        if let Some(waker) = waker {
+            if let Some(prev_waker) = ctx.waker.as_mut() {
+                // only replace the Waker if they dont reference the same task
+                if !prev_waker.will_wake(waker) {
+                    *prev_waker = waker.clone();
+                }
+            } else {
+                ctx.waker = Some(waker.clone());
+            }
+        } else {
+            ctx.waker = None;
+        }
         Ok(self)
     }
+
+    /// Returns the Waker set on the connection context.
+    pub fn waker(&self) -> Option<&Waker> {
+        let ctx = self.context();
+        ctx.waker.as_ref()
+    }
+
+    /// Retrieve a mutable reference to the [`Context`] stored on the connection.
+    fn context_mut(&mut self) -> &mut Context {
+        unsafe {
+            let ctx = s2n_connection_get_ctx(self.connection.as_ptr())
+                .into_result()
+                .unwrap();
+            &mut *(ctx.as_ptr() as *mut Context)
+        }
+    }
+
+    /// Retrieve a reference to the [`Context`] stored on the connection.
+    fn context(&self) -> &Context {
+        unsafe {
+            let ctx = s2n_connection_get_ctx(self.connection.as_ptr())
+                .into_result()
+                .unwrap();
+            &*(ctx.as_ptr() as *mut Context)
+        }
+    }
+
+    /// Mark that the server_name extension was used to configure the connection.
+    pub fn server_name_extension_used(&mut self) {
+        // TODO: requiring the application to call this method is a pretty sharp edge.
+        // Figure out if its possible to automatically call this from the Rust bindings.
+        unsafe {
+            s2n_connection_server_name_extension_used(self.connection.as_ptr())
+                .into_result()
+                .unwrap();
+        }
+    }
+
+    #[cfg(test)]
+    /// Test if a config has been set on the connection.
+    ///
+    /// `s2n_connection_get_config` should return a NULL pointer if the `s2n_connection_set_config`
+    /// has not been called by the application.
+    pub fn test_config_exists(&mut self) -> Result<(), Error> {
+        let mut config = core::ptr::null_mut();
+        let _config = unsafe {
+            s2n_connection_get_config(self.connection.as_ptr(), &mut config).into_result()?
+        };
+        Ok(())
+    }
+}
+
+#[derive(Default)]
+struct Context {
+    waker: Option<Waker>,
 }
 
 #[cfg(feature = "quic")]
@@ -309,6 +447,18 @@ impl Connection {
 impl Drop for Connection {
     fn drop(&mut self) {
         // ignore failures since there's not much we can do about it
-        let _ = unsafe { s2n_connection_free(self.connection.as_ptr()).into_result() };
+        unsafe {
+            // clean up context
+            let prev_ctx = self.context_mut();
+            drop(Box::from_raw(prev_ctx));
+            let _ = s2n_connection_set_ctx(self.connection.as_ptr(), core::ptr::null_mut())
+                .into_result();
+
+            // cleanup config
+            let _ = self.drop_config();
+
+            // cleanup connection
+            let _ = s2n_connection_free(self.connection.as_ptr()).into_result();
+        }
     }
 }