Move async callbacks infra to boring itself

This helps drive async callbacks from outside tokio-boring, such as in quiche.

Not a breaking change because every public item in tokio-boring is preserved as is.

boring/src/ssl/async_callbacks.rs

@@ -0,0 +1,308 @@
+use super::mut_only::MutOnly;
+use super::{
+    ClientHello, GetSessionPendingError, PrivateKeyMethod, PrivateKeyMethodError, SelectCertError,
+    Ssl, SslContextBuilder, SslRef, SslSession, SslSignatureAlgorithm,
+};
+use crate::ex_data::Index;
+use once_cell::sync::Lazy;
+use std::convert::identity;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{ready, Context, Poll, Waker};
+
+/// The type of futures to pass to [`SslContextBuilderExt::set_async_select_certificate_callback`].
+pub type BoxSelectCertFuture = ExDataFuture<Result<BoxSelectCertFinish, AsyncSelectCertError>>;
+
+/// The type of callbacks returned by [`BoxSelectCertFuture`] methods.
+pub type BoxSelectCertFinish = Box<dyn FnOnce(ClientHello<'_>) -> Result<(), AsyncSelectCertError>>;
+
+/// The type of futures returned by [`AsyncPrivateKeyMethod`] methods.
+pub type BoxPrivateKeyMethodFuture =
+    ExDataFuture<Result<BoxPrivateKeyMethodFinish, AsyncPrivateKeyMethodError>>;
+
+/// The type of callbacks returned by [`BoxPrivateKeyMethodFuture`].
+pub type BoxPrivateKeyMethodFinish =
+    Box<dyn FnOnce(&mut SslRef, &mut [u8]) -> Result<usize, AsyncPrivateKeyMethodError>>;
+
+/// The type of futures to pass to [`SslContextBuilderExt::set_async_get_session_callback`].
+pub type BoxGetSessionFuture = ExDataFuture<Option<BoxGetSessionFinish>>;
+
+/// The type of callbacks returned by [`BoxSelectCertFuture`] methods.
+pub type BoxGetSessionFinish = Box<dyn FnOnce(&mut SslRef, &[u8]) -> Option<SslSession>>;
+
+/// Convenience alias for futures stored in [`Ssl`] ex data by [`SslContextBuilderExt`] methods.
+///
+/// Public for documentation purposes.
+pub type ExDataFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
+
+pub(crate) static TASK_WAKER_INDEX: Lazy<Index<Ssl, Option<Waker>>> =
+    Lazy::new(|| Ssl::new_ex_index().unwrap());
+pub(crate) static SELECT_CERT_FUTURE_INDEX: Lazy<Index<Ssl, MutOnly<Option<BoxSelectCertFuture>>>> =
+    Lazy::new(|| Ssl::new_ex_index().unwrap());
+pub(crate) static SELECT_PRIVATE_KEY_METHOD_FUTURE_INDEX: Lazy<
+    Index<Ssl, MutOnly<Option<BoxPrivateKeyMethodFuture>>>,
+> = Lazy::new(|| Ssl::new_ex_index().unwrap());
+pub(crate) static SELECT_GET_SESSION_FUTURE_INDEX: Lazy<
+    Index<Ssl, MutOnly<Option<BoxGetSessionFuture>>>,
+> = Lazy::new(|| Ssl::new_ex_index().unwrap());
+
+impl SslContextBuilder {
+    /// Sets a callback that is called before most [`ClientHello`] processing
+    /// and before the decision whether to resume a session is made. The
+    /// callback may inspect the [`ClientHello`] and configure the connection.
+    ///
+    /// This method uses a function that returns a future whose output is
+    /// itself a closure that will be passed [`ClientHello`] to configure
+    /// the connection based on the computations done in the future.
+    ///
+    /// A task waker must be set on `Ssl` values associated with the resulting
+    /// `SslContext` with [`SslRef::set_task_waker`].
+    ///
+    /// See [`SslContextBuilder::set_select_certificate_callback`] for the sync
+    /// setter of this callback.
+    pub fn set_async_select_certificate_callback<F>(&mut self, callback: F)
+    where
+        F: Fn(&mut ClientHello<'_>) -> Result<BoxSelectCertFuture, AsyncSelectCertError>
+            + Send
+            + Sync
+            + 'static,
+    {
+        self.set_select_certificate_callback(move |mut client_hello| {
+            let fut_poll_result = with_ex_data_future(
+                &mut client_hello,
+                *SELECT_CERT_FUTURE_INDEX,
+                ClientHello::ssl_mut,
+                &callback,
+                identity,
+            );
+
+            let fut_result = match fut_poll_result {
+                Poll::Ready(fut_result) => fut_result,
+                Poll::Pending => return Err(SelectCertError::RETRY),
+            };
+
+            let finish = fut_result.or(Err(SelectCertError::ERROR))?;
+
+            finish(client_hello).or(Err(SelectCertError::ERROR))
+        })
+    }
+
+    /// Configures a custom private key method on the context.
+    ///
+    /// A task waker must be set on `Ssl` values associated with the resulting
+    /// `SslContext` with [`SslRef::set_task_waker`].
+    ///
+    /// See [`AsyncPrivateKeyMethod`] for more details.
+    pub fn set_async_private_key_method(&mut self, method: impl AsyncPrivateKeyMethod) {
+        self.set_private_key_method(AsyncPrivateKeyMethodBridge(Box::new(method)));
+    }
+
+    /// Sets a callback that is called when a client proposed to resume a session
+    /// but it was not found in the internal cache.
+    ///
+    /// The callback is passed a reference to the session ID provided by the client.
+    /// It should return the session corresponding to that ID if available. This is
+    /// only used for servers, not clients.
+    ///
+    /// A task waker must be set on `Ssl` values associated with the resulting
+    /// `SslContext` with [`SslRef::set_task_waker`].
+    ///
+    /// See [`SslContextBuilder::set_get_session_callback`] for the sync setter
+    /// of this callback.
+    ///
+    /// # Safety
+    ///
+    /// The returned [`SslSession`] must not be associated with a different [`SslContext`].
+    pub unsafe fn set_async_get_session_callback<F>(&mut self, callback: F)
+    where
+        F: Fn(&mut SslRef, &[u8]) -> Option<BoxGetSessionFuture> + Send + Sync + 'static,
+    {
+        let async_callback = move |ssl: &mut SslRef, id: &[u8]| {
+            let fut_poll_result = with_ex_data_future(
+                &mut *ssl,
+                *SELECT_GET_SESSION_FUTURE_INDEX,
+                |ssl| ssl,
+                |ssl| callback(ssl, id).ok_or(()),
+                |option| option.ok_or(()),
+            );
+
+            match fut_poll_result {
+                Poll::Ready(Err(())) => Ok(None),
+                Poll::Ready(Ok(finish)) => Ok(finish(ssl, id)),
+                Poll::Pending => Err(GetSessionPendingError),
+            }
+        };
+
+        self.set_get_session_callback(async_callback)
+    }
+}
+
+impl SslRef {
+    /// Sets the task waker to be used in async callbacks installed on this `Ssl`.
+    pub fn set_task_waker(&mut self, waker: Option<Waker>) {
+        self.replace_ex_data(*TASK_WAKER_INDEX, waker);
+    }
+}
+
+/// A fatal error to be returned from async select certificate callbacks.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct AsyncSelectCertError;
+
+/// Describes async private key hooks. This is used to off-load signing
+/// operations to a custom, potentially asynchronous, backend. Metadata about the
+/// key such as the type and size are parsed out of the certificate.
+///
+/// See [`PrivateKeyMethod`] for the sync version of those hooks.
+///
+/// [`ssl_private_key_method_st`]: https://commondatastorage.googleapis.com/chromium-boringssl-docs/ssl.h.html#ssl_private_key_method_st
+pub trait AsyncPrivateKeyMethod: Send + Sync + 'static {
+    /// Signs the message `input` using the specified signature algorithm.
+    ///
+    /// This method uses a function that returns a future whose output is
+    /// itself a closure that will be passed `ssl` and `output`
+    /// to finish writing the signature.
+    ///
+    /// See [`PrivateKeyMethod::sign`] for the sync version of this method.
+    fn sign(
+        &self,
+        ssl: &mut SslRef,
+        input: &[u8],
+        signature_algorithm: SslSignatureAlgorithm,
+        output: &mut [u8],
+    ) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>;
+
+    /// Decrypts `input`.
+    ///
+    /// This method uses a function that returns a future whose output is
+    /// itself a closure that will be passed `ssl` and `output`
+    /// to finish decrypting the input.
+    ///
+    /// See [`PrivateKeyMethod::decrypt`] for the sync version of this method.
+    fn decrypt(
+        &self,
+        ssl: &mut SslRef,
+        input: &[u8],
+        output: &mut [u8],
+    ) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>;
+}
+
+/// A fatal error to be returned from async private key methods.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct AsyncPrivateKeyMethodError;
+
+struct AsyncPrivateKeyMethodBridge(Box<dyn AsyncPrivateKeyMethod>);
+
+impl PrivateKeyMethod for AsyncPrivateKeyMethodBridge {
+    fn sign(
+        &self,
+        ssl: &mut SslRef,
+        input: &[u8],
+        signature_algorithm: SslSignatureAlgorithm,
+        output: &mut [u8],
+    ) -> Result<usize, PrivateKeyMethodError> {
+        with_private_key_method(ssl, output, |ssl, output| {
+            <dyn AsyncPrivateKeyMethod>::sign(&*self.0, ssl, input, signature_algorithm, output)
+        })
+    }
+
+    fn decrypt(
+        &self,
+        ssl: &mut SslRef,
+        input: &[u8],
+        output: &mut [u8],
+    ) -> Result<usize, PrivateKeyMethodError> {
+        with_private_key_method(ssl, output, |ssl, output| {
+            <dyn AsyncPrivateKeyMethod>::decrypt(&*self.0, ssl, input, output)
+        })
+    }
+
+    fn complete(
+        &self,
+        ssl: &mut SslRef,
+        output: &mut [u8],
+    ) -> Result<usize, PrivateKeyMethodError> {
+        with_private_key_method(ssl, output, |_, _| {
+            // This should never be reached, if it does, that's a bug on boring's side,
+            // which called `complete` without having been returned to with a pending
+            // future from `sign` or `decrypt`.
+
+            if cfg!(debug_assertions) {
+                panic!("BUG: boring called complete without a pending operation");
+            }
+
+            Err(AsyncPrivateKeyMethodError)
+        })
+    }
+}
+
+/// Creates and drives a private key method future.
+///
+/// This is a convenience function for the three methods of impl `PrivateKeyMethod``
+/// for `dyn AsyncPrivateKeyMethod`. It relies on [`with_ex_data_future`] to
+/// drive the future and then immediately calls the final [`BoxPrivateKeyMethodFinish`]
+/// when the future is ready.
+fn with_private_key_method(
+    ssl: &mut SslRef,
+    output: &mut [u8],
+    create_fut: impl FnOnce(
+        &mut SslRef,
+        &mut [u8],
+    ) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>,
+) -> Result<usize, PrivateKeyMethodError> {
+    let fut_poll_result = with_ex_data_future(
+        ssl,
+        *SELECT_PRIVATE_KEY_METHOD_FUTURE_INDEX,
+        |ssl| ssl,
+        |ssl| create_fut(ssl, output),
+        identity,
+    );
+
+    let fut_result = match fut_poll_result {
+        Poll::Ready(fut_result) => fut_result,
+        Poll::Pending => return Err(PrivateKeyMethodError::RETRY),
+    };
+
+    let finish = fut_result.or(Err(PrivateKeyMethodError::FAILURE))?;
+
+    finish(ssl, output).or(Err(PrivateKeyMethodError::FAILURE))
+}
+
+/// Creates and drives a future stored in `ssl_handle`'s `Ssl` at ex data index `index`.
+///
+/// This function won't even bother storing the future in `index` if the future
+/// created by `create_fut` returns `Poll::Ready(_)` on the first poll call.
+fn with_ex_data_future<H, R, T, E>(
+    ssl_handle: &mut H,
+    index: Index<Ssl, MutOnly<Option<ExDataFuture<R>>>>,
+    get_ssl_mut: impl Fn(&mut H) -> &mut SslRef,
+    create_fut: impl FnOnce(&mut H) -> Result<ExDataFuture<R>, E>,
+    into_result: impl Fn(R) -> Result<T, E>,
+) -> Poll<Result<T, E>> {
+    let ssl = get_ssl_mut(ssl_handle);
+    let waker = ssl
+        .ex_data(*TASK_WAKER_INDEX)
+        .cloned()
+        .flatten()
+        .expect("task waker should be set");
+
+    let mut ctx = Context::from_waker(&waker);
+
+    if let Some(data @ Some(_)) = ssl.ex_data_mut(index).map(MutOnly::get_mut) {
+        let fut_result = into_result(ready!(data.as_mut().unwrap().as_mut().poll(&mut ctx)));
+
+        *data = None;
+
+        Poll::Ready(fut_result)
+    } else {
+        let mut fut = create_fut(ssl_handle)?;
+
+        match fut.as_mut().poll(&mut ctx) {
+            Poll::Ready(fut_result) => Poll::Ready(into_result(fut_result)),
+            Poll::Pending => {
+                get_ssl_mut(ssl_handle).replace_ex_data(index, MutOnly::new(Some(fut)));
+
+                Poll::Pending
+            }
+        }
+    }
+}

boring/src/ssl/mod.rs

@@ -97,15 +97,22 @@ use crate::x509::{
 };
 use crate::{cvt, cvt_0i, cvt_n, cvt_p, init};
 
-pub use crate::ssl::connector::{
+pub use self::async_callbacks::{
+    AsyncPrivateKeyMethod, AsyncPrivateKeyMethodError, AsyncSelectCertError, BoxGetSessionFinish,
+    BoxGetSessionFuture, BoxPrivateKeyMethodFinish, BoxPrivateKeyMethodFuture, BoxSelectCertFinish,
+    BoxSelectCertFuture, ExDataFuture,
+};
+pub use self::connector::{
     ConnectConfiguration, SslAcceptor, SslAcceptorBuilder, SslConnector, SslConnectorBuilder,
 };
-pub use crate::ssl::error::{Error, ErrorCode, HandshakeError};
+pub use self::error::{Error, ErrorCode, HandshakeError};
 
+mod async_callbacks;
 mod bio;
 mod callbacks;
 mod connector;
 mod error;
+mod mut_only;
 #[cfg(test)]
 mod test;
 

tokio-boring/src/async_callbacks.rs

@@ -1,47 +1,7 @@
-use crate::mut_only::MutOnly;
+use boring::ssl::{
-use boring::ex_data::Index;
+    AsyncPrivateKeyMethod, AsyncSelectCertError, BoxGetSessionFuture, BoxSelectCertFuture,
-use boring::ssl::{self, ClientHello, PrivateKeyMethod, Ssl, SslContextBuilder};
+    ClientHello, SslContextBuilder, SslRef,
-use once_cell::sync::Lazy;
+};
-use std::convert::identity;
-use std::future::Future;
-use std::pin::Pin;
-use std::task::{ready, Context, Poll, Waker};
-
-/// The type of futures to pass to [`SslContextBuilderExt::set_async_select_certificate_callback`].
-pub type BoxSelectCertFuture = ExDataFuture<Result<BoxSelectCertFinish, AsyncSelectCertError>>;
-
-/// The type of callbacks returned by [`BoxSelectCertFuture`] methods.
-pub type BoxSelectCertFinish = Box<dyn FnOnce(ClientHello<'_>) -> Result<(), AsyncSelectCertError>>;
-
-/// The type of futures returned by [`AsyncPrivateKeyMethod`] methods.
-pub type BoxPrivateKeyMethodFuture =
-    ExDataFuture<Result<BoxPrivateKeyMethodFinish, AsyncPrivateKeyMethodError>>;
-
-/// The type of callbacks returned by [`BoxPrivateKeyMethodFuture`].
-pub type BoxPrivateKeyMethodFinish =
-    Box<dyn FnOnce(&mut ssl::SslRef, &mut [u8]) -> Result<usize, AsyncPrivateKeyMethodError>>;
-
-/// The type of futures to pass to [`SslContextBuilderExt::set_async_get_session_callback`].
-pub type BoxGetSessionFuture = ExDataFuture<Option<BoxGetSessionFinish>>;
-
-/// The type of callbacks returned by [`BoxSelectCertFuture`] methods.
-pub type BoxGetSessionFinish = Box<dyn FnOnce(&mut ssl::SslRef, &[u8]) -> Option<ssl::SslSession>>;
-
-/// Convenience alias for futures stored in [`Ssl`] ex data by [`SslContextBuilderExt`] methods.
-///
-/// Public for documentation purposes.
-pub type ExDataFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
-
-pub(crate) static TASK_WAKER_INDEX: Lazy<Index<Ssl, Option<Waker>>> =
-    Lazy::new(|| Ssl::new_ex_index().unwrap());
-pub(crate) static SELECT_CERT_FUTURE_INDEX: Lazy<Index<Ssl, MutOnly<Option<BoxSelectCertFuture>>>> =
-    Lazy::new(|| Ssl::new_ex_index().unwrap());
-pub(crate) static SELECT_PRIVATE_KEY_METHOD_FUTURE_INDEX: Lazy<
-    Index<Ssl, MutOnly<Option<BoxPrivateKeyMethodFuture>>>,
-> = Lazy::new(|| Ssl::new_ex_index().unwrap());
-pub(crate) static SELECT_GET_SESSION_FUTURE_INDEX: Lazy<
-    Index<Ssl, MutOnly<Option<BoxGetSessionFuture>>>,
-> = Lazy::new(|| Ssl::new_ex_index().unwrap());
 
 /// Extensions to [`SslContextBuilder`].
 ///
@@ -84,7 +44,7 @@ pub trait SslContextBuilderExt: private::Sealed {
     /// The returned [`SslSession`] must not be associated with a different [`SslContext`].
     unsafe fn set_async_get_session_callback<F>(&mut self, callback: F)
     where
-        F: Fn(&mut ssl::SslRef, &[u8]) -> Option<BoxGetSessionFuture> + Send + Sync + 'static;
+        F: Fn(&mut SslRef, &[u8]) -> Option<BoxGetSessionFuture> + Send + Sync + 'static;
 }
 
 impl SslContextBuilderExt for SslContextBuilder {
@@ -95,214 +55,18 @@ impl SslContextBuilderExt for SslContextBuilder {
             + Sync
             + 'static,
     {
-        self.set_select_certificate_callback(move |mut client_hello| {
+        self.set_async_select_certificate_callback(callback);
-            let fut_poll_result = with_ex_data_future(
-                &mut client_hello,
-                *SELECT_CERT_FUTURE_INDEX,
-                ClientHello::ssl_mut,
-                &callback,
-                identity,
-            );
-
-            let fut_result = match fut_poll_result {
-                Poll::Ready(fut_result) => fut_result,
-                Poll::Pending => return Err(ssl::SelectCertError::RETRY),
-            };
-
-            let finish = fut_result.or(Err(ssl::SelectCertError::ERROR))?;
-
-            finish(client_hello).or(Err(ssl::SelectCertError::ERROR))
-        })
     }
 
     fn set_async_private_key_method(&mut self, method: impl AsyncPrivateKeyMethod) {
-        self.set_private_key_method(AsyncPrivateKeyMethodBridge(Box::new(method)));
+        self.set_async_private_key_method(method);
     }
 
     unsafe fn set_async_get_session_callback<F>(&mut self, callback: F)
     where
-        F: Fn(&mut ssl::SslRef, &[u8]) -> Option<BoxGetSessionFuture> + Send + Sync + 'static,
+        F: Fn(&mut SslRef, &[u8]) -> Option<BoxGetSessionFuture> + Send + Sync + 'static,
     {
-        let async_callback = move |ssl: &mut ssl::SslRef, id: &[u8]| {
+        self.set_async_get_session_callback(callback);
-            let fut_poll_result = with_ex_data_future(
-                &mut *ssl,
-                *SELECT_GET_SESSION_FUTURE_INDEX,
-                |ssl| ssl,
-                |ssl| callback(ssl, id).ok_or(()),
-                |option| option.ok_or(()),
-            );
-
-            match fut_poll_result {
-                Poll::Ready(Err(())) => Ok(None),
-                Poll::Ready(Ok(finish)) => Ok(finish(ssl, id)),
-                Poll::Pending => Err(ssl::GetSessionPendingError),
-            }
-        };
-
-        self.set_get_session_callback(async_callback)
-    }
-}
-
-/// A fatal error to be returned from async select certificate callbacks.
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub struct AsyncSelectCertError;
-
-/// Describes async private key hooks. This is used to off-load signing
-/// operations to a custom, potentially asynchronous, backend. Metadata about the
-/// key such as the type and size are parsed out of the certificate.
-///
-/// See [`PrivateKeyMethod`] for the sync version of those hooks.
-///
-/// [`ssl_private_key_method_st`]: https://commondatastorage.googleapis.com/chromium-boringssl-docs/ssl.h.html#ssl_private_key_method_st
-pub trait AsyncPrivateKeyMethod: Send + Sync + 'static {
-    /// Signs the message `input` using the specified signature algorithm.
-    ///
-    /// This method uses a function that returns a future whose output is
-    /// itself a closure that will be passed `ssl` and `output`
-    /// to finish writing the signature.
-    ///
-    /// See [`PrivateKeyMethod::sign`] for the sync version of this method.
-    fn sign(
-        &self,
-        ssl: &mut ssl::SslRef,
-        input: &[u8],
-        signature_algorithm: ssl::SslSignatureAlgorithm,
-        output: &mut [u8],
-    ) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>;
-
-    /// Decrypts `input`.
-    ///
-    /// This method uses a function that returns a future whose output is
-    /// itself a closure that will be passed `ssl` and `output`
-    /// to finish decrypting the input.
-    ///
-    /// See [`PrivateKeyMethod::decrypt`] for the sync version of this method.
-    fn decrypt(
-        &self,
-        ssl: &mut ssl::SslRef,
-        input: &[u8],
-        output: &mut [u8],
-    ) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>;
-}
-
-/// A fatal error to be returned from async private key methods.
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub struct AsyncPrivateKeyMethodError;
-
-struct AsyncPrivateKeyMethodBridge(Box<dyn AsyncPrivateKeyMethod>);
-
-impl PrivateKeyMethod for AsyncPrivateKeyMethodBridge {
-    fn sign(
-        &self,
-        ssl: &mut ssl::SslRef,
-        input: &[u8],
-        signature_algorithm: ssl::SslSignatureAlgorithm,
-        output: &mut [u8],
-    ) -> Result<usize, ssl::PrivateKeyMethodError> {
-        with_private_key_method(ssl, output, |ssl, output| {
-            <dyn AsyncPrivateKeyMethod>::sign(&*self.0, ssl, input, signature_algorithm, output)
-        })
-    }
-
-    fn decrypt(
-        &self,
-        ssl: &mut ssl::SslRef,
-        input: &[u8],
-        output: &mut [u8],
-    ) -> Result<usize, ssl::PrivateKeyMethodError> {
-        with_private_key_method(ssl, output, |ssl, output| {
-            <dyn AsyncPrivateKeyMethod>::decrypt(&*self.0, ssl, input, output)
-        })
-    }
-
-    fn complete(
-        &self,
-        ssl: &mut ssl::SslRef,
-        output: &mut [u8],
-    ) -> Result<usize, ssl::PrivateKeyMethodError> {
-        with_private_key_method(ssl, output, |_, _| {
-            // This should never be reached, if it does, that's a bug on boring's side,
-            // which called `complete` without having been returned to with a pending
-            // future from `sign` or `decrypt`.
-
-            if cfg!(debug_assertions) {
-                panic!("BUG: boring called complete without a pending operation");
-            }
-
-            Err(AsyncPrivateKeyMethodError)
-        })
-    }
-}
-
-/// Creates and drives a private key method future.
-///
-/// This is a convenience function for the three methods of impl `PrivateKeyMethod``
-/// for `dyn AsyncPrivateKeyMethod`. It relies on [`with_ex_data_future`] to
-/// drive the future and then immediately calls the final [`BoxPrivateKeyMethodFinish`]
-/// when the future is ready.
-fn with_private_key_method(
-    ssl: &mut ssl::SslRef,
-    output: &mut [u8],
-    create_fut: impl FnOnce(
-        &mut ssl::SslRef,
-        &mut [u8],
-    ) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>,
-) -> Result<usize, ssl::PrivateKeyMethodError> {
-    let fut_poll_result = with_ex_data_future(
-        ssl,
-        *SELECT_PRIVATE_KEY_METHOD_FUTURE_INDEX,
-        |ssl| ssl,
-        |ssl| create_fut(ssl, output),
-        identity,
-    );
-
-    let fut_result = match fut_poll_result {
-        Poll::Ready(fut_result) => fut_result,
-        Poll::Pending => return Err(ssl::PrivateKeyMethodError::RETRY),
-    };
-
-    let finish = fut_result.or(Err(ssl::PrivateKeyMethodError::FAILURE))?;
-
-    finish(ssl, output).or(Err(ssl::PrivateKeyMethodError::FAILURE))
-}
-
-/// Creates and drives a future stored in `ssl_handle`'s `Ssl` at ex data index `index`.
-///
-/// This function won't even bother storing the future in `index` if the future
-/// created by `create_fut` returns `Poll::Ready(_)` on the first poll call.
-fn with_ex_data_future<H, R, T, E>(
-    ssl_handle: &mut H,
-    index: Index<ssl::Ssl, MutOnly<Option<ExDataFuture<R>>>>,
-    get_ssl_mut: impl Fn(&mut H) -> &mut ssl::SslRef,
-    create_fut: impl FnOnce(&mut H) -> Result<ExDataFuture<R>, E>,
-    into_result: impl Fn(R) -> Result<T, E>,
-) -> Poll<Result<T, E>> {
-    let ssl = get_ssl_mut(ssl_handle);
-    let waker = ssl
-        .ex_data(*TASK_WAKER_INDEX)
-        .cloned()
-        .flatten()
-        .expect("task waker should be set");
-
-    let mut ctx = Context::from_waker(&waker);
-
-    if let Some(data @ Some(_)) = ssl.ex_data_mut(index).map(MutOnly::get_mut) {
-        let fut_result = into_result(ready!(data.as_mut().unwrap().as_mut().poll(&mut ctx)));
-
-        *data = None;
-
-        Poll::Ready(fut_result)
-    } else {
-        let mut fut = create_fut(ssl_handle)?;
-
-        match fut.as_mut().poll(&mut ctx) {
-            Poll::Ready(fut_result) => Poll::Ready(into_result(fut_result)),
-            Poll::Pending => {
-                get_ssl_mut(ssl_handle).replace_ex_data(index, MutOnly::new(Some(fut)));
-
-                Poll::Pending
-            }
-        }
     }
 }
 

tokio-boring/src/lib.rs

@@ -29,17 +29,20 @@ use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 
 mod async_callbacks;
 mod bridge;
-mod mut_only;
 
-use self::async_callbacks::TASK_WAKER_INDEX;
-pub use self::async_callbacks::{
-    AsyncPrivateKeyMethod, AsyncPrivateKeyMethodError, AsyncSelectCertError, BoxGetSessionFinish,
-    BoxGetSessionFuture, BoxPrivateKeyMethodFinish, BoxPrivateKeyMethodFuture, BoxSelectCertFinish,
-    BoxSelectCertFuture, ExDataFuture, SslContextBuilderExt,
-};
 use self::bridge::AsyncStreamBridge;
 
+pub use crate::async_callbacks::SslContextBuilderExt;
+pub use boring::ssl::{
+    AsyncPrivateKeyMethod, AsyncPrivateKeyMethodError, AsyncSelectCertError, BoxGetSessionFinish,
+    BoxGetSessionFuture, BoxPrivateKeyMethodFinish, BoxPrivateKeyMethodFuture, BoxSelectCertFinish,
+    BoxSelectCertFuture, ExDataFuture,
+};
+
 /// Asynchronously performs a client-side TLS handshake over the provided stream.
+///
+/// This function automatically sets the task waker on the `Ssl` from `config` to
+/// allow to make use of async callbacks provided by the boring crate.
 pub async fn connect<S>(
     config: ConnectConfiguration,
     domain: &str,
@@ -52,6 +55,9 @@ where
 }
 
 /// Asynchronously performs a server-side TLS handshake over the provided stream.
+///
+/// This function automatically sets the task waker on the `Ssl` from `config` to
+/// allow to make use of async callbacks provided by the boring crate.
 pub async fn accept<S>(acceptor: &SslAcceptor, stream: S) -> Result<SslStream<S>, HandshakeError<S>>
 where
     S: AsyncRead + AsyncWrite + Unpin,
@@ -300,20 +306,18 @@ where
         mid_handshake.get_mut().set_waker(Some(ctx));
         mid_handshake
             .ssl_mut()
-            .replace_ex_data(*TASK_WAKER_INDEX, Some(ctx.waker().clone()));
+            .set_task_waker(Some(ctx.waker().clone()));
 
         match mid_handshake.handshake() {
             Ok(mut stream) => {
                 stream.get_mut().set_waker(None);
-                stream.ssl_mut().replace_ex_data(*TASK_WAKER_INDEX, None);
+                stream.ssl_mut().set_task_waker(None);
 
                 Poll::Ready(Ok(SslStream(stream)))
             }
             Err(ssl::HandshakeError::WouldBlock(mut mid_handshake)) => {
                 mid_handshake.get_mut().set_waker(None);
-                mid_handshake
+                mid_handshake.ssl_mut().set_task_waker(None);
-                    .ssl_mut()
-                    .replace_ex_data(*TASK_WAKER_INDEX, None);
 
                 self.0 = Some(mid_handshake);
 

tokio-boring/tests/async_get_session.rs

@@ -4,7 +4,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::OnceLock;
 use tokio::net::TcpStream;
 use tokio::task::yield_now;
-use tokio_boring::{BoxGetSessionFinish, SslContextBuilderExt};
+use tokio_boring::BoxGetSessionFinish;
 
 mod common;
 

tokio-boring/tests/async_private_key_method.rs

@@ -5,10 +5,7 @@ use boring::sign::{RsaPssSaltlen, Signer};
 use boring::ssl::{SslRef, SslSignatureAlgorithm};
 use futures::future;
 use tokio::task::yield_now;
-use tokio_boring::{
+use tokio_boring::{AsyncPrivateKeyMethod, AsyncPrivateKeyMethodError, BoxPrivateKeyMethodFuture};
-    AsyncPrivateKeyMethod, AsyncPrivateKeyMethodError, BoxPrivateKeyMethodFuture,
-    SslContextBuilderExt,
-};
 
 mod common;
 

tokio-boring/tests/async_select_certificate.rs

@@ -1,9 +1,7 @@
 use boring::ssl::ClientHello;
 use futures::future;
 use tokio::task::yield_now;
-use tokio_boring::{
+use tokio_boring::{AsyncSelectCertError, BoxSelectCertFinish, BoxSelectCertFuture};
-    AsyncSelectCertError, BoxSelectCertFinish, BoxSelectCertFuture, SslContextBuilderExt,
-};
 
 mod common;