From eb644640518d41a30dfe02decc7225e5bdea2c3a Mon Sep 17 00:00:00 2001
From: Ian Campbell <52475242+xv-ian-c@users.noreply.github.com>
Date: Thu, 9 Nov 2023 10:42:26 +0000
Subject: [PATCH] wolfssl: Avoid needing to leak keylogger callback `Box`

We have to jump through a few hoops but it is possible for the callback to work
only in terms of raw pointers and references without ever instantiating an
actual `Box` or `Arc` and therefore having to worry about drop.

In the absence of [strict provenance][] (nightly only unstable feature) we do
still need the `Box` in order to make a thin pointer to the fat `dyn
Tls13SecretCallbacks` which we need to access.

(Aside: I think making `Session` generic over a `CB: Tls13SecretCallbacks`
would avoid the box, as we do with `IOCB`, however for the keylogger debug
facility we don't need to be quite so efficient and the generics get
everywhere)

However the `Box` should be part of `Self` so that the required lifetime is
established. From that we can obtain a (thin) reference to the allocation
within the `Box` which contains the `Arc<dyn...>` which contains a fat
reference to the actual callback object. That thin reference from the `Box` can
then be turned into a raw pointer (which a fat pointer cannot without strict
provenance).

The `&mut **self.secret_cb.as_mut().unwrap()` construct is tricky:

- `self.secret_cb.as_mut().unwrap()` has type `&Box<Arc<dyn...>>`;
- therefore `*...` is the `Box<Arc<dyn...>>` itself;
- therefore `**...` is the allocation within the box i.e. `Arc<dyn...>>`;
- finally `&**...` is a reference to that allocation i.e. the `&Arc<dyn...>`
  which we need.

[strict provenance]: https://github.com/rust-lang/rust/issues/95228
---
 wolfssl/src/ssl.rs | 31 +++++++++++++++++++++----------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/wolfssl/src/ssl.rs b/wolfssl/src/ssl.rs
index 7a8c63f..cb2ff04 100644
--- a/wolfssl/src/ssl.rs
+++ b/wolfssl/src/ssl.rs
@@ -184,7 +184,7 @@ pub struct Session<IOCB: IOCallbacks> {
     io: Box<IOCB>,
 
     #[cfg(feature = "debug")]
-    secret_cb: Option<Tls13SecretCallbacksArg>,
+    secret_cb: Option<Box<Tls13SecretCallbacksArg>>,
 }
 
 /// Error creating a [`Session`] object.
@@ -1131,13 +1131,20 @@ impl<IOCB: IOCallbacks> Session<IOCB> {
     /// Enable TLS1.3 key logging for applications
     #[cfg(feature = "debug")]
     pub(crate) fn enable_tls13_keylog(&mut self, secret_cb: Tls13SecretCallbacksArg) -> Result<()> {
-        self.secret_cb = Some(secret_cb.clone());
+        self.secret_cb = Some(Box::new(secret_cb));
 
-        // SAFETY: `secret_cb` is an Arc pointer and we save one strong reference inside `Session`
-        // This prevents the memory being freed when `secret_cb` goes out of scope.
-        // So it is safe to send this pointer as callback context and `secret_cb` will be valid
-        // as long the connection is valid
-        let secret_cb = Box::into_raw(Box::new(secret_cb)) as *mut c_void;
+        // SAFETY: `secret_cb` is a `Box` pointer so the address is
+        // stable. (The address is the address of the heap allocation
+        // containing the `Tls13Secretcallbacksarg` which is an `Arc`.
+        //
+        // We free `self.ssl` (the `wolfssl_sys::WOLFSSL`) on drop of
+        // `self`, any use of the callback must have stopped before
+        // the drop, since nothing can also be making calls into the session.
+        //
+        // Therefore `secret_cb` here is valid for as long as it needs to be.
+        let secret_cb: &mut Tls13SecretCallbacksArg = &mut **self.secret_cb.as_mut().unwrap();
+        let secret_cb: *mut Tls13SecretCallbacksArg = secret_cb as *mut Tls13SecretCallbacksArg;
+        let secret_cb = secret_cb as *mut c_void;
 
         // SAFETY: No documentation available for [`wolfSSL_KeepArrays`] [`wolfSSL_set_tls13_secret_cb`].
         // But based on api implementation, it expects a valid pointer to `WOLFSSL`.
@@ -1168,9 +1175,13 @@ impl<IOCB: IOCallbacks> Session<IOCB> {
         debug_assert!(!secret.is_null());
         debug_assert!(!ctx.is_null());
 
-        // SAFETY: We have only one strong reference in `self.secret_cb`
-        // Leak the Box pointer, so that it will not be freed while this function return
-        let secret_cb = Box::leak(Box::from_raw(ctx as *mut Tls13SecretCallbacksArg));
+        // SAFETY: We know this pointer is to the contents of the
+        // `Box<Tls13SecretCallbacksArg>` at `self.secret_cb` which is
+        // owned by the `Session`. See `enable_tls13_keylog` above for
+        // an argument to why calls to this callback cannot happen
+        // after the `Session` is dropped.
+        let secret_cb = ctx as *mut Tls13SecretCallbacksArg;
+        let secret_cb: &Tls13SecretCallbacksArg = &*secret_cb;
 
         let mut random: Vec<u8> = vec![0; RANDOM_SIZE];
         let get_random = if 1 == wolfssl_sys::wolfSSL_is_server(ssl) {