Properly drop overwritten ex data

boring/src/ssl/mod.rs

@@ -864,7 +864,9 @@ impl SslContextBuilder {
             // out. When that happens, we wouldn't be able to look up the callback's state in the
             // context's ex data. Instead, pass the pointer directly as the servername arg. It's
             // still stored in ex data to manage the lifetime.
-            let arg = self.set_ex_data_inner(SslContext::cached_ex_index::<F>(), callback);
+            let arg = self
+                .ctx
+                .set_ex_data(SslContext::cached_ex_index::<F>(), callback);
 
             ffi::SSL_CTX_set_tlsext_servername_arg(self.as_ptr(), arg);
             ffi::SSL_CTX_set_tlsext_servername_callback(self.as_ptr(), Some(raw_sni::<F>));
@@ -1655,14 +1657,8 @@ impl SslContextBuilder {
     ///
     /// [`SSL_CTX_set_ex_data`]: https://www.openssl.org/docs/man1.0.2/ssl/SSL_CTX_set_ex_data.html
     pub fn set_ex_data<T>(&mut self, index: Index<SslContext, T>, data: T) {
-        self.set_ex_data_inner(index, data);
-    }
-
-    fn set_ex_data_inner<T>(&mut self, index: Index<SslContext, T>, data: T) -> *mut c_void {
         unsafe {
-            let data = Box::into_raw(Box::new(data)) as *mut c_void;
-            ffi::SSL_CTX_set_ex_data(self.as_ptr(), index.as_raw(), data);
-            data
+            self.ctx.set_ex_data(index, data);
         }
     }
 
@@ -1916,6 +1912,33 @@ impl SslContextRef {
         }
     }
 
+    // Unsafe because SSL contexts are not guaranteed to be unique, we call
+    // this only from SslContextBuilder.
+    unsafe fn ex_data_mut<T>(&mut self, index: Index<SslContext, T>) -> Option<&mut T> {
+        let data = ffi::SSL_CTX_get_ex_data(self.as_ptr(), index.as_raw());
+        if data.is_null() {
+            None
+        } else {
+            Some(&mut *(data as *mut T))
+        }
+    }
+
+    // Unsafe because SSL contexts are not guaranteed to be unique, we call
+    // this only from SslContextBuilder.
+    unsafe fn set_ex_data<T>(&mut self, index: Index<SslContext, T>, data: T) -> *mut c_void {
+        if let Some(old) = self.ex_data_mut(index) {
+            *old = data;
+
+            return old as *mut T as *mut c_void;
+        }
+
+        unsafe {
+            let data = Box::into_raw(Box::new(data)) as *mut c_void;
+            ffi::SSL_CTX_set_ex_data(self.as_ptr(), index.as_raw(), data);
+            data
+        }
+    }
+
     /// Adds a session to the context's cache.
     ///
     /// Returns `true` if the session was successfully added to the cache, and `false` if it was already present.
@@ -3193,6 +3216,12 @@ impl SslRef {
     ///
     /// [`SSL_set_ex_data`]: https://www.openssl.org/docs/manmaster/man3/SSL_set_ex_data.html
     pub fn set_ex_data<T>(&mut self, index: Index<Ssl, T>, data: T) {
+        if let Some(old) = self.ex_data_mut(index) {
+            *old = data;
+
+            return;
+        }
+
         unsafe {
             let data = Box::new(data);
             ffi::SSL_set_ex_data(

boring/src/ssl/test/mod.rs

@@ -5,6 +5,7 @@ use std::mem;
 use std::net::{TcpListener, TcpStream};
 use std::path::Path;
 use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 use std::thread;
 
 use crate::error::ErrorStack;
@@ -1044,3 +1045,56 @@ fn server_set_default_curves_list() {
     // Panics if Kyber768 missing in boringSSL.
     ssl.server_set_default_curves_list();
 }
+
+#[test]
+fn drop_ex_data_in_context() {
+    let index = SslContext::new_ex_index::<Arc<()>>().unwrap();
+    let mut ctx = SslContext::builder(SslMethod::dtls()).unwrap();
+
+    let arc = Arc::new(());
+
+    for _ in 0..5 {
+        ctx.set_ex_data(index, arc.clone());
+    }
+
+    assert_eq!(
+        Arc::strong_count(&arc),
+        2,
+        "should be 2, one for the one we own, one for the ex data"
+    );
+
+    drop(ctx);
+
+    assert_eq!(
+        Arc::strong_count(&arc),
+        1,
+        "should be 1, as we dropped the context and thus its ex data"
+    );
+}
+
+#[test]
+fn drop_ex_data_in_ssl() {
+    let index = Ssl::new_ex_index::<Arc<()>>().unwrap();
+    let ctx = SslContext::builder(SslMethod::dtls()).unwrap().build();
+    let mut ssl = Ssl::new(&ctx).unwrap();
+
+    let arc = Arc::new(());
+
+    for _ in 0..5 {
+        ssl.set_ex_data(index, arc.clone());
+    }
+
+    assert_eq!(
+        Arc::strong_count(&arc),
+        2,
+        "should be 2, one for the one we own, one for the ex data"
+    );
+
+    drop(ssl);
+
+    assert_eq!(
+        Arc::strong_count(&arc),
+        1,
+        "should be 1, as we dropped the context and thus its ex data"
+    );
+}