From a001e27a638c5bc07cd70b484ed8de93ffae2d40 Mon Sep 17 00:00:00 2001
From: Lindsay Stewart <slindsay@amazon.com>
Date: Tue, 24 Sep 2024 21:36:52 -0700
Subject: [PATCH] fix(bindings): handle failures from wipe (#4798)

Co-authored-by: Wesley Rosenblum <55108558+WesleyRosenblum@users.noreply.github.com>
---
 .github/workflows/ci_rust.yml            |  2 +-
 bindings/rust/s2n-tls/src/connection.rs  | 37 +++++++++++++++++-------
 bindings/rust/s2n-tls/src/renegotiate.rs | 11 +++++++
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/ci_rust.yml b/.github/workflows/ci_rust.yml
index 3c49b90c9cb..cfaf30318ba 100644
--- a/.github/workflows/ci_rust.yml
+++ b/.github/workflows/ci_rust.yml
@@ -47,7 +47,7 @@ jobs:
       - name: Tests
         working-directory: ${{env.ROOT_PATH}}
         # Test all features except for FIPS, which is tested separately.
-        run: cargo test --features unstable-fingerprint,unstable-ktls,quic,pq
+        run: cargo test --features unstable-renegotiate,unstable-fingerprint,unstable-ktls,quic,pq
 
       # Ensure that all tests pass with the default feature set
       - name: Default Tests
diff --git a/bindings/rust/s2n-tls/src/connection.rs b/bindings/rust/s2n-tls/src/connection.rs
index 9e70bcd8fee..8482d0f70fa 100644
--- a/bindings/rust/s2n-tls/src/connection.rs
+++ b/bindings/rust/s2n-tls/src/connection.rs
@@ -468,16 +468,17 @@ impl Connection {
         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));
 
-            wipe(self)
-        }?;
+        // Safety:
+        // We re-init the context after the wipe
+        unsafe { self.drop_context()? };
 
+        let result = wipe(self);
+        // We must initialize the context again whether or not wipe succeeds.
+        // A connection without a context is invalid and has undefined behavior.
         self.init_context(mode);
+        result?;
+
         Ok(())
     }
 
@@ -793,6 +794,23 @@ impl Connection {
         }
     }
 
+    /// Drop the context
+    ///
+    /// SAFETY:
+    /// A connection without a context is invalid. After calling this method
+    /// from anywhere other than Drop, you must reinitialize the context.
+    unsafe fn drop_context(&mut self) -> Result<(), Error> {
+        let ctx = s2n_connection_get_ctx(self.connection.as_ptr()).into_result();
+        if let Ok(ctx) = ctx {
+            drop(Box::from_raw(ctx.as_ptr() as *mut Context));
+        }
+        // Setting a NULL context is important: if we don't also remove the context
+        // from the connection, then the invalid memory is still accessible and
+        // may even be double-freed.
+        s2n_connection_set_ctx(self.connection.as_ptr(), core::ptr::null_mut()).into_result()?;
+        Ok(())
+    }
+
     /// 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.
@@ -1253,10 +1271,7 @@ impl Drop for Connection {
         // ignore failures since there's not much we can do about it
         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();
+            let _ = self.drop_context();
 
             // cleanup config
             let _ = self.drop_config();
diff --git a/bindings/rust/s2n-tls/src/renegotiate.rs b/bindings/rust/s2n-tls/src/renegotiate.rs
index a1832618eeb..ba382f174a4 100644
--- a/bindings/rust/s2n-tls/src/renegotiate.rs
+++ b/bindings/rust/s2n-tls/src/renegotiate.rs
@@ -148,6 +148,7 @@ mod tests {
             ConnectionFuture, ConnectionFutureResult, PrivateKeyCallback, PrivateKeyOperation,
         },
         config::ConnectionInitializer,
+        error::ErrorSource,
         testing::{CertKeyPair, InsecureAcceptAllCertificatesHandler, TestPair, TestPairIO},
     };
     use foreign_types::ForeignTypeRef;
@@ -615,4 +616,14 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn wipe_for_renegotiate_failure() -> Result<(), Box<dyn Error>> {
+        let mut connection = Connection::new_server();
+        // Servers can't renegotiate
+        let error = connection.wipe_for_renegotiate().unwrap_err();
+        assert_eq!(error.source(), ErrorSource::Library);
+        assert_eq!(error.name(), "S2N_ERR_NO_RENEGOTIATION");
+        Ok(())
+    }
 }