feat(bindings): enable multiple cert chains on config

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

@@ -4,8 +4,10 @@
 use crate::error::{Error, Fallible};
 use s2n_tls_sys::*;
 use std::{
+    ffi::{c_void, CString},
     marker::PhantomData,
     ptr::{self, NonNull},
+    sync::atomic::{AtomicUsize, Ordering},
 };
 
 /// A CertificateChain represents a chain of X.509 certificates.
@@ -18,16 +20,26 @@ pub struct CertificateChain<'a> {
 impl CertificateChain<'_> {
     /// This allocates a new certificate chain from s2n.
     pub(crate) fn new() -> Result<CertificateChain<'static>, Error> {
+        crate::init::init();
+        let ptr = unsafe { s2n_cert_chain_and_key_new().into_result()? };
+
+        let context = Box::<CertificateChainContext>::default();
+        let context = Box::into_raw(context) as *mut c_void;
+
         unsafe {
-            let ptr = s2n_cert_chain_and_key_new().into_result()?;
-            Ok(CertificateChain {
-                ptr,
-                is_owned: true,
-                _lifetime: PhantomData,
-            })
+            s2n_cert_chain_and_key_set_ctx(ptr.as_ptr(), context)
+                .into_result()
+                .unwrap();
         }
+        Ok(CertificateChain {
+            ptr,
+            is_owned: true,
+            _lifetime: PhantomData,
+        })
     }
 
+    /// This CertificateChain is not owned and will not increment the reference count.
+    /// When the rust instance is dropped it will not drop the pointer.
     pub(crate) unsafe fn from_ptr_reference<'a>(
         ptr: NonNull<s2n_cert_chain_and_key>,
     ) -> CertificateChain<'a> {
@@ -38,6 +50,26 @@ impl CertificateChain<'_> {
         }
     }
 
+    /// # Safety
+    ///
+    /// This CertificateChain _MUST_ have been initialized with the constructor.
+    /// Additionally, this does NOT increment the reference count,
+    /// so consider cloning the result if the source pointer is still valid and usable afterwards.
+    pub(crate) unsafe fn from_owned_ptr_reference<'a>(
+        ptr: NonNull<s2n_cert_chain_and_key>,
+    ) -> CertificateChain<'a> {
+        let cert_chain = CertificateChain {
+            ptr,
+            is_owned: true,
+            _lifetime: PhantomData,
+        };
+
+        // Check if the context can be retrieved.
+        // If it can't, this is not an owned CertificateChain created through constructor.
+        cert_chain.context();
+        cert_chain
+    }
+
     pub fn iter(&self) -> CertificateChainIter<'_> {
         CertificateChainIter {
             idx: 0,
@@ -75,20 +107,115 @@ impl CertificateChain<'_> {
     pub(crate) fn as_mut_ptr(&mut self) -> NonNull<s2n_cert_chain_and_key> {
         self.ptr
     }
+
+    /// Retrieve a reference to the [`CertificateChainContext`] stored on the CertificateChain.
+    pub(crate) fn context(&self) -> &CertificateChainContext {
+        unsafe {
+            let ctx = s2n_cert_chain_and_key_get_ctx(self.ptr.as_ptr())
+                .into_result()
+                .unwrap();
+            &*(ctx.as_ptr() as *const CertificateChainContext)
+        }
+    }
+
+    /// Retrieve a mutable reference to the [`CertificateChainContext`] stored on the CertificateChain.
+    pub(crate) fn context_mut(&mut self) -> &mut CertificateChainContext {
+        unsafe {
+            let ctx = s2n_cert_chain_and_key_get_ctx(self.ptr.as_ptr())
+                .into_result()
+                .unwrap();
+            &mut *(ctx.as_ptr() as *mut CertificateChainContext)
+        }
+    }
+
+    pub fn load_pem(&mut self, certificate: &[u8], private_key: &[u8]) -> Result<&mut Self, Error> {
+        let certificate = CString::new(certificate).map_err(|_| Error::INVALID_INPUT)?;
+        let private_key = CString::new(private_key).map_err(|_| Error::INVALID_INPUT)?;
+        unsafe {
+            s2n_cert_chain_and_key_load_pem(
+                self.ptr.as_ptr(),
+                certificate.as_ptr(),
+                private_key.as_ptr(),
+            )
+            .into_result()
+        }?;
+        Ok(self)
+    }
+
+    pub fn set_ocsp_data(&mut self, data: &[u8]) -> Result<&mut Self, Error> {
+        let size: u32 = data.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
+        unsafe {
+            s2n_cert_chain_and_key_set_ocsp_data(self.ptr.as_ptr(), data.as_ptr(), size)
+                .into_result()
+        }?;
+        Ok(self)
+    }
+}
+
+impl Clone for CertificateChain<'_> {
+    fn clone(&self) -> Self {
+        let context = 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 {
+            ptr: self.ptr,
+            is_owned: true, // clone only makes sense for owned
+            _lifetime: PhantomData,
+        }
+    }
 }
 
 // # Safety
 //
 // s2n_cert_chain_and_key objects can be sent across threads.
 unsafe impl Send for CertificateChain<'_> {}
 
+/// # Safety
+///
+/// Safety: All C methods that mutate the s2n_cert_chain are wrapped
+/// in Rust methods that require a mutable reference.
+unsafe impl Sync for CertificateChain<'_> {}
+
 impl Drop for CertificateChain<'_> {
     fn drop(&mut self) {
-        if self.is_owned {
-            // ignore failures since there's not much we can do about it
-            unsafe {
-                let _ = s2n_cert_chain_and_key_free(self.ptr.as_ptr()).into_result();
-            }
+        if !self.is_owned {
+            // not ours to cleanup
+            return;
+        }
+        let context = 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 cert 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_cert_chain_and_key_free(self.ptr.as_ptr()).into_result();
         }
     }
 }
@@ -152,3 +279,40 @@ impl<'a> Certificate<'a> {
 //
 // Certificates just reference data in the chain, so share the Send-ness of the chain.
 unsafe impl Send for Certificate<'_> {}
+
+pub(crate) struct CertificateChainContext {
+    refcount: AtomicUsize,
+}
+
+impl Default for CertificateChainContext {
+    fn default() -> Self {
+        // The AtomicUsize is used to manually track the reference count of the CertificateChain.
+        // This mechanism is used to track when the CertificateChain object should be freed.
+        Self {
+            refcount: AtomicUsize::new(1),
+        }
+    }
+}
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn clone_and_drop_update_ref_count() {
+        let original_cert = CertificateChain::new().unwrap();
+        assert_eq!(original_cert.context().refcount.load(Ordering::Relaxed), 1);
+
+        let second_cert = original_cert.clone();
+        assert_eq!(original_cert.context().refcount.load(Ordering::Relaxed), 2);
+
+        drop(second_cert);
+        assert_eq!(original_cert.context().refcount.load(Ordering::Relaxed), 1);
+    }
+
+    // ensure the CertificateChainContext is send and sync
+    #[test]
+    fn context_send_sync_test() {
+        fn assert_send_sync<T: 'static + Send + Sync>() {}
+        assert_send_sync::<CertificateChainContext>();
+    }
+}

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

@@ -5,10 +5,12 @@
 use crate::renegotiate::RenegotiateCallback;
 use crate::{
     callbacks::*,
+    cert_chain::CertificateChain,
     enums::*,
     error::{Error, Fallible},
     security,
 };
+use core::mem::{self};
 use core::{convert::TryInto, ptr::NonNull};
 use s2n_tls_sys::*;
 use std::{
@@ -156,6 +158,26 @@ impl Drop for Config {
             let context = Box::from_raw(context);
             drop(context);
 
+            // Clean up the certificate chains
+            let mut cert_chains: *mut *mut s2n_cert_chain_and_key = std::ptr::null_mut();
+            let mut chain_count: u32 = 0;
+            if s2n_config_get_cert_chains(self.0.as_ptr(), &mut cert_chains, &mut chain_count)
+                .into_result()
+                .is_ok()
+            {
+                if !cert_chains.is_null() && chain_count > 0 {
+                    let cert_slice = std::slice::from_raw_parts(cert_chains, chain_count as usize);
+                    for &cert_ptr in cert_slice {
+                        if let Some(non_null_ptr) = NonNull::new(cert_ptr) {
+                            drop(CertificateChain::from_owned_ptr_reference(non_null_ptr));
+                        }
+                    }
+
+                    let _ =
+                        s2n_free_config_cert_chains(&mut cert_chains, chain_count).into_result();
+                }
+            }
+
             let _ = s2n_config_free(self.0.as_ptr()).into_result();
         }
     }
@@ -277,6 +299,26 @@ impl Builder {
         Ok(self)
     }
 
+    /// Adds the CertificateChain to the Config.
+    /// Prefer to use this function over load_pem.
+    /// This function is not compatible with load_pem and will error if both are used
+    pub fn add_cert_chain(&mut self, mut cert_chain: CertificateChain) -> Result<&mut Self, Error> {
+        unsafe {
+            s2n_config_add_cert_chain_and_key_to_store(
+                self.as_mut_ptr(),
+                cert_chain.as_mut_ptr().as_ptr(),
+            )
+            .into_result()?;
+        }
+        // Setting the cert chain on the config creates one additional reference
+        // so do not drop so prevent Rust from calling `drop()` at the end of this function.
+        mem::forget(cert_chain);
+        Ok(self)
+    }
+
+    /// Creates a certificate chain and binds it to the config.
+    /// Prefer to use the newer add_cert_chain which enables multiple certs over this function.
+    /// This function is not compatible with add_cert_chain and will error if both are used
     pub fn load_pem(&mut self, certificate: &[u8], private_key: &[u8]) -> Result<&mut Self, Error> {
         let certificate = CString::new(certificate).map_err(|_| Error::INVALID_INPUT)?;
         let private_key = CString::new(private_key).map_err(|_| Error::INVALID_INPUT)?;
@@ -395,6 +437,8 @@ impl Builder {
     /// Sets the OCSP data for the default certificate chain associated with the Config.
     ///
     /// Servers will send the data in response to OCSP stapling requests from clients.
+    ///
+    /// Prefer to use add_cert_chain with a CertificateChain that has OCSP data set over this function.
     //
     // NOTE: this modifies a certificate chain, NOT the Config itself. This is currently safe
     // because the certificate chain is set with s2n_config_add_cert_chain_and_key, which