feat(bindings): add renegotiate to the rust bindings (#4668)

bindings/rust/s2n-tls/Cargo.toml

@@ -12,6 +12,7 @@ license = "Apache-2.0"
 default = []
 unstable-fingerprint = ["s2n-tls-sys/unstable-fingerprint"]
 unstable-ktls = ["s2n-tls-sys/unstable-ktls"]
+unstable-renegotiate = ["s2n-tls-sys/unstable-renegotiate"]
 quic = ["s2n-tls-sys/quic"]
 fips = ["s2n-tls-sys/fips"]
 pq = ["s2n-tls-sys/pq"]
@@ -27,6 +28,8 @@ hex = "0.4"
 [dev-dependencies]
 futures-test = "0.3"
 openssl = "0.10"
+openssl-sys = "0.9"
+foreign-types = "0.3"
 temp-env = "0.3"
 checkers = "0.6"
 # newer versions require rust 1.66, see https://github.com/aws/s2n-tls/issues/4241

bindings/rust/s2n-tls/src/callbacks/pkey.rs

@@ -73,7 +73,7 @@ impl PrivateKeyOperation {
     /// The size of the slice returned by [`input()`]
     pub fn input_size(&self) -> Result<usize, Error> {
         let mut size = 0;
-        unsafe { s2n_async_pkey_op_get_input_size(self.raw.as_ptr(), &mut size) }.into_result()?;
+        unsafe { s2n_async_pkey_op_get_input_size(self.as_ptr(), &mut size) }.into_result()?;
         size.try_into().map_err(|_| Error::INVALID_INPUT)
     }
 
@@ -84,8 +84,7 @@ impl PrivateKeyOperation {
     pub fn input(&self, buf: &mut [u8]) -> Result<(), Error> {
         let buf_len: u32 = buf.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
         let buf_ptr = buf.as_ptr() as *mut u8;
-        unsafe { s2n_async_pkey_op_get_input(self.raw.as_ptr(), buf_ptr, buf_len) }
-            .into_result()?;
+        unsafe { s2n_async_pkey_op_get_input(self.as_ptr(), buf_ptr, buf_len) }.into_result()?;
         Ok(())
     }
 
@@ -94,11 +93,15 @@ impl PrivateKeyOperation {
         let buf_len: u32 = buf.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
         let buf_ptr = buf.as_ptr();
         unsafe {
-            s2n_async_pkey_op_set_output(self.raw.as_ptr(), buf_ptr, buf_len).into_result()?;
-            s2n_async_pkey_op_apply(self.raw.as_ptr(), conn.as_ptr()).into_result()?;
+            s2n_async_pkey_op_set_output(self.as_ptr(), buf_ptr, buf_len).into_result()?;
+            s2n_async_pkey_op_apply(self.as_ptr(), conn.as_ptr()).into_result()?;
         }
         Ok(())
     }
+
+    pub(crate) fn as_ptr(&self) -> *mut s2n_async_pkey_op {
+        self.raw.as_ptr()
+    }
 }
 
 impl Drop for PrivateKeyOperation {

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

@@ -1,6 +1,8 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+#[cfg(feature = "unstable-renegotiate")]
+use crate::renegotiate::RenegotiateCallback;
 use crate::{
     callbacks::*,
     enums::*,
@@ -436,12 +438,12 @@ impl Builder {
             verify_host(host_name, host_name_len, handler)
         }
 
-        self.config.context_mut().verify_host_callback = Some(Box::new(handler));
+        self.context_mut().verify_host_callback = Some(Box::new(handler));
         unsafe {
             s2n_config_set_verify_host_callback(
                 self.as_mut_ptr(),
                 Some(verify_host_cb_fn),
-                self.config.context_mut() as *mut Context as *mut c_void,
+                self.context_mut() as *mut Context as *mut c_void,
             )
             .into_result()?;
         }
@@ -490,7 +492,7 @@ impl Builder {
         }
 
         let handler = Box::new(handler);
-        let context = self.config.context_mut();
+        let context = self.context_mut();
         context.client_hello_callback = Some(handler);
 
         unsafe {
@@ -511,7 +513,7 @@ impl Builder {
     ) -> Result<&mut Self, Error> {
         // Store callback in config context
         let handler = Box::new(handler);
-        let context = self.config.context_mut();
+        let context = self.context_mut();
         context.connection_initializer = Some(handler);
         Ok(self)
     }
@@ -540,14 +542,14 @@ impl Builder {
 
         // Store callback in context
         let handler = Box::new(handler);
-        let context = self.config.context_mut();
+        let context = self.context_mut();
         context.session_ticket_callback = Some(handler);
 
         unsafe {
             s2n_config_set_session_ticket_cb(
                 self.as_mut_ptr(),
                 Some(session_ticket_cb),
-                self.config.context_mut() as *mut Context as *mut c_void,
+                self.context_mut() as *mut Context as *mut c_void,
             )
             .into_result()
         }?;
@@ -577,7 +579,7 @@ impl Builder {
         }
 
         let handler = Box::new(handler);
-        let context = self.config.context_mut();
+        let context = self.context_mut();
         context.private_key_callback = Some(handler);
 
         unsafe {
@@ -611,13 +613,13 @@ impl Builder {
         }
 
         let handler = Box::new(handler);
-        let context = self.config.context_mut();
+        let context = self.context_mut();
         context.wall_clock = Some(handler);
         unsafe {
             s2n_config_set_wall_clock(
                 self.as_mut_ptr(),
                 Some(clock_cb),
-                self.config.context_mut() as *mut _ as *mut c_void,
+                self.context_mut() as *mut _ as *mut c_void,
             )
             .into_result()?;
         }
@@ -648,13 +650,13 @@ impl Builder {
         }
 
         let handler = Box::new(handler);
-        let context = self.config.context_mut();
+        let context = self.context_mut();
         context.monotonic_clock = Some(handler);
         unsafe {
             s2n_config_set_monotonic_clock(
                 self.as_mut_ptr(),
                 Some(clock_cb),
-                self.config.context_mut() as *mut _ as *mut c_void,
+                self.context_mut() as *mut _ as *mut c_void,
             )
             .into_result()?;
         }
@@ -762,9 +764,20 @@ impl Builder {
         Ok(self.config)
     }
 
-    fn as_mut_ptr(&mut self) -> *mut s2n_config {
+    pub(crate) fn as_mut_ptr(&mut self) -> *mut s2n_config {
         self.config.as_mut_ptr()
     }
+
+    /// Retrieve a mutable reference to the [`Context`] stored on the config.
+    pub(crate) fn context_mut(&mut self) -> &mut Context {
+        let mut ctx = core::ptr::null_mut();
+        unsafe {
+            s2n_config_get_ctx(self.as_mut_ptr(), &mut ctx)
+                .into_result()
+                .unwrap();
+            &mut *(ctx as *mut Context)
+        }
+    }
 }
 
 #[cfg(feature = "quic")]
@@ -797,6 +810,8 @@ pub(crate) struct Context {
     pub(crate) connection_initializer: Option<Box<dyn ConnectionInitializer>>,
     pub(crate) wall_clock: Option<Box<dyn WallClock>>,
     pub(crate) monotonic_clock: Option<Box<dyn MonotonicClock>>,
+    #[cfg(feature = "unstable-renegotiate")]
+    pub(crate) renegotiate: Option<Box<dyn RenegotiateCallback>>,
 }
 
 impl Default for Context {
@@ -814,6 +829,8 @@ impl Default for Context {
             connection_initializer: None,
             wall_clock: None,
             monotonic_clock: None,
+            #[cfg(feature = "unstable-renegotiate")]
+            renegotiate: None,
         }
     }
 }

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

@@ -394,6 +394,14 @@ impl Connection {
         Ok(self)
     }
 
+    /// # Safety
+    ///
+    /// The `context` pointer must live at least as long as the connection
+    pub unsafe fn set_send_context(&mut self, context: *mut c_void) -> Result<&mut Self, Error> {
+        s2n_connection_set_send_ctx(self.connection.as_ptr(), context).into_result()?;
+        Ok(self)
+    }
+
     /// Sets the callback to use for verifying that a hostname from an X.509 certificate is
     /// trusted.
     ///
@@ -428,14 +436,6 @@ impl Connection {
         Ok(self)
     }
 
-    /// # Safety
-    ///
-    /// The `context` pointer must live at least as long as the connection
-    pub unsafe fn set_send_context(&mut self, context: *mut c_void) -> Result<&mut Self, Error> {
-        s2n_connection_set_send_ctx(self.connection.as_ptr(), context).into_result()?;
-        Ok(self)
-    }
-
     /// Connections prefering low latency will be encrypted using small record sizes that
     /// can be decrypted sooner by the recipient.
     pub fn prefer_low_latency(&mut self) -> Result<&mut Self, Error> {
@@ -463,38 +463,36 @@ impl Connection {
         Ok(self)
     }
 
-    /// wipes an existing connection and allows it to be reused.
-    ///
-    /// This method erases all data associated with a connection including pending reads.
-    /// This function should be called after all I/O is completed and s2n_shutdown has been
-    /// called. Reusing the same connection handle(s) is more performant than repeatedly
-    /// calling s2n_connection_new and s2n_connection_free
-    pub fn wipe(&mut self) -> Result<&mut Self, Error> {
+    pub(crate) fn wipe_method<F, T>(&mut self, wipe: F) -> Result<(), Error>
+    where
+        F: FnOnce(&mut Self) -> Result<T, Error>,
+    {
         let mode = self.mode();
         unsafe {
             // Wiping the connection will wipe the pointer to the context,
             // so retrieve and drop that memory first.
             let ctx = self.context_mut();
             drop(Box::from_raw(ctx));
 
-            s2n_connection_wipe(self.connection.as_ptr()).into_result()
+            wipe(self)
         }?;
 
         self.init_context(mode);
-        Ok(self)
+        Ok(())
     }
 
-    /// Performs the TLS handshake to completion
-    ///
-    /// Multiple callbacks can be configured for a connection and config, but
-    /// [`Self::poll_negotiate()`] can only execute and block on one callback at a time.
-    /// The handshake is sequential, not concurrent, and stops execution when
-    /// it encounters an async callback.
+    /// wipes an existing connection and allows it to be reused.
     ///
-    /// The handshake does not continue execution (and therefore can't call
-    /// any other callbacks) until the blocking async task reports completion.
-    pub fn poll_negotiate(&mut self) -> Poll<Result<&mut Self, Error>> {
-        let mut blocked = s2n_blocked_status::NOT_BLOCKED;
+    /// This method erases all data associated with a connection including pending reads.
+    /// This function should be called after all I/O is completed and s2n_shutdown has been
+    /// called. Reusing the same connection handle(s) is more performant than repeatedly
+    /// calling s2n_connection_new and s2n_connection_free
+    pub fn wipe(&mut self) -> Result<&mut Self, Error> {
+        self.wipe_method(|conn| unsafe { s2n_connection_wipe(conn.as_ptr()).into_result() })?;
+        Ok(self)
+    }
+
+    fn trigger_initializer(&mut self) {
         if !core::mem::replace(&mut self.context_mut().connection_initialized, true) {
             if let Some(config) = self.config() {
                 if let Some(callback) = config.context().connection_initializer.as_ref() {
@@ -503,40 +501,6 @@ impl Connection {
                 }
             }
         }
-
-        loop {
-            // check if an async task exists and poll it to completion
-            if let Some(fut) = self.poll_async_task() {
-                match fut {
-                    Poll::Ready(Ok(())) => {
-                        // happy case:
-                        // continue and call s2n_negotiate to make progress on the handshake
-                    }
-                    Poll::Ready(Err(err)) => {
-                        // error case:
-                        // if the callback returned an error then abort the handshake
-                        return Poll::Ready(Err(err));
-                    }
-                    Poll::Pending => return Poll::Pending,
-                }
-            }
-
-            let res = unsafe { s2n_negotiate(self.connection.as_ptr(), &mut blocked).into_poll() };
-
-            match res {
-                Poll::Ready(res) => {
-                    let res = res.map(|_| self);
-                    return Poll::Ready(res);
-                }
-                Poll::Pending => {
-                    // if there is no connection_future then return, otherwise continue
-                    // looping and polling the future
-                    if self.context_mut().async_callback.is_none() {
-                        return Poll::Pending;
-                    }
-                }
-            }
-        }
     }
 
     // Poll the connection future if it exists.
@@ -557,6 +521,39 @@ impl Connection {
         })
     }
 
+    pub(crate) fn poll_negotiate_method<F, T>(&mut self, negotiate: F) -> Poll<Result<(), Error>>
+    where
+        F: FnOnce(&mut Connection) -> Poll<Result<T, Error>>,
+    {
+        self.trigger_initializer();
+
+        // Check whether renegotiate is blocked by any async callbacks
+        match self.poll_async_task().unwrap_or(Poll::Ready(Ok(()))) {
+            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+            Poll::Pending => return Poll::Pending,
+            Poll::Ready(Ok(_)) => {}
+        };
+
+        negotiate(self).map_ok(|_| ())
+    }
+
+    /// Performs the TLS handshake to completion
+    ///
+    /// Multiple callbacks can be configured for a connection and config, but
+    /// [`Self::poll_negotiate()`] can only execute and block on one callback at a time.
+    /// The handshake is sequential, not concurrent, and stops execution when
+    /// it encounters an async callback.
+    ///
+    /// The handshake does not continue execution (and therefore can't call
+    /// any other callbacks) until the blocking async task reports completion.
+    pub fn poll_negotiate(&mut self) -> Poll<Result<&mut Self, Error>> {
+        let mut blocked = s2n_blocked_status::NOT_BLOCKED;
+        self.poll_negotiate_method(|conn| unsafe {
+            s2n_negotiate(conn.as_ptr(), &mut blocked).into_poll()
+        })
+        .map_ok(|_| self)
+    }
+
     /// Encrypts and sends data on a connection where
     /// [negotiate](`Self::poll_negotiate`) has succeeded.
     ///

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

@@ -23,6 +23,8 @@ pub mod enums;
 pub mod fingerprint;
 pub mod init;
 pub mod pool;
+#[cfg(feature = "unstable-renegotiate")]
+pub mod renegotiate;
 pub mod security;
 
 pub use s2n_tls_sys as ffi;