Introduce async callbacks

We introduce tokio_boring::SslContextBuilderExt, with 2 methods:

* set_async_select_certificate_callback
* set_async_private_key_method
This commit is contained in:
Anthony Ramine 2023-07-28 15:15:11 +02:00 committed by Alessandro Ghedini
parent 0ffbdb030f
commit 907eaf079c
8 changed files with 580 additions and 3 deletions

View File

@ -482,6 +482,9 @@ pub struct SelectCertError(ffi::ssl_select_cert_result_t);
impl SelectCertError { impl SelectCertError {
/// A fatal error occured and the handshake should be terminated. /// A fatal error occured and the handshake should be terminated.
pub const ERROR: Self = Self(ffi::ssl_select_cert_result_t::ssl_select_cert_error); pub const ERROR: Self = Self(ffi::ssl_select_cert_result_t::ssl_select_cert_error);
/// The operation could not be completed and should be retried later.
pub const RETRY: Self = Self(ffi::ssl_select_cert_result_t::ssl_select_cert_retry);
} }
/// Extension types, to be used with `ClientHello::get_extension`. /// Extension types, to be used with `ClientHello::get_extension`.
@ -3280,6 +3283,11 @@ impl<S> MidHandshakeSslStream<S> {
self.stream.ssl() self.stream.ssl()
} }
/// Returns a mutable reference to the `Ssl` of the stream.
pub fn ssl_mut(&mut self) -> &mut SslRef {
self.stream.ssl_mut()
}
/// Returns the underlying error which interrupted this handshake. /// Returns the underlying error which interrupted this handshake.
pub fn error(&self) -> &Error { pub fn error(&self) -> &Error {
&self.error &self.error
@ -3585,6 +3593,11 @@ impl<S> SslStream<S> {
pub fn ssl(&self) -> &SslRef { pub fn ssl(&self) -> &SslRef {
&self.ssl &self.ssl
} }
/// Returns a mutable reference to the `Ssl` object associated with this stream.
pub fn ssl_mut(&mut self) -> &mut SslRef {
&mut self.ssl
}
} }
impl<S: Read + Write> Read for SslStream<S> { impl<S: Read + Write> Read for SslStream<S> {

View File

@ -189,7 +189,8 @@ fn test_sign_retry_complete_failure() {
ErrorCode::WANT_PRIVATE_KEY_OPERATION ErrorCode::WANT_PRIVATE_KEY_OPERATION
); );
let HandshakeError::WouldBlock(mid_handshake) = mid_handshake.handshake().unwrap_err() else { let HandshakeError::WouldBlock(mid_handshake) = mid_handshake.handshake().unwrap_err()
else {
panic!("should be WouldBlock"); panic!("should be WouldBlock");
}; };

View File

@ -39,6 +39,7 @@ no-patches = ["boring/no-patches"]
[dependencies] [dependencies]
boring = { workspace = true } boring = { workspace = true }
boring-sys = { workspace = true } boring-sys = { workspace = true }
once_cell = { workspace = true }
tokio = { workspace = true } tokio = { workspace = true }
[dev-dependencies] [dev-dependencies]

View File

@ -0,0 +1,263 @@
use boring::ex_data::Index;
use boring::ssl::{self, ClientHello, PrivateKeyMethod, Ssl, SslContextBuilder};
use once_cell::sync::Lazy;
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>>;
/// 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 + Sync>>;
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, BoxSelectCertFuture>> =
Lazy::new(|| Ssl::new_ex_index().unwrap());
pub(crate) static SELECT_PRIVATE_KEY_METHOD_FUTURE_INDEX: Lazy<
Index<Ssl, BoxPrivateKeyMethodFuture>,
> = Lazy::new(|| Ssl::new_ex_index().unwrap());
/// Extensions to [`SslContextBuilder`].
///
/// This trait provides additional methods to use async callbacks with boring.
pub trait SslContextBuilderExt: private::Sealed {
/// 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.
///
/// See [`SslContextBuilder::set_select_certificate_callback`] for the sync
/// setter of this callback.
fn set_async_select_certificate_callback<F>(&mut self, callback: F)
where
F: Fn(&mut ClientHello<'_>) -> Result<BoxSelectCertFuture, AsyncSelectCertError>
+ Send
+ Sync
+ 'static;
/// Configures a custom private key method on the context.
///
/// See [`AsyncPrivateKeyMethod`] for more details.
fn set_async_private_key_method(&mut self, method: impl AsyncPrivateKeyMethod);
}
impl SslContextBuilderExt for SslContextBuilder {
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,
);
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)));
}
}
/// 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),
);
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, T, E>(
ssl_handle: &mut H,
index: Index<ssl::Ssl, ExDataFuture<Result<T, E>>>,
get_ssl_mut: impl Fn(&mut H) -> &mut ssl::SslRef,
create_fut: impl FnOnce(&mut H) -> Result<ExDataFuture<Result<T, E>>, 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);
match ssl.ex_data_mut(index) {
Some(fut) => {
let fut_result = ready!(fut.as_mut().poll(&mut ctx));
// NOTE(nox): For memory usage concerns, maybe we should implement
// a way to remove the stored future from the `Ssl` value here?
Poll::Ready(fut_result)
}
None => {
let mut fut = create_fut(ssl_handle)?;
match fut.as_mut().poll(&mut ctx) {
Poll::Ready(fut_result) => Poll::Ready(fut_result),
Poll::Pending => {
get_ssl_mut(ssl_handle).set_ex_data(index, fut);
Poll::Pending
}
}
}
}
}
mod private {
pub trait Sealed {}
}
impl private::Sealed for SslContextBuilder {}

View File

@ -1,5 +1,4 @@
//! Bridge between sync IO traits and async tokio IO traits. //! Bridge between sync IO traits and async tokio IO traits.
use std::fmt; use std::fmt;
use std::io; use std::io;
use std::pin::Pin; use std::pin::Pin;
@ -35,7 +34,7 @@ impl<S> AsyncStreamBridge<S> {
F: FnOnce(&mut Context<'_>, Pin<&mut S>) -> R, F: FnOnce(&mut Context<'_>, Pin<&mut S>) -> R,
{ {
let mut ctx = let mut ctx =
Context::from_waker(self.waker.as_ref().expect("missing task context pointer")); Context::from_waker(self.waker.as_ref().expect("BUG: missing waker in bridge"));
f(&mut ctx, Pin::new(&mut self.stream)) f(&mut ctx, Pin::new(&mut self.stream))
} }

View File

@ -27,8 +27,15 @@ use std::pin::Pin;
use std::task::{Context, Poll}; use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
mod async_callbacks;
mod bridge; mod bridge;
use self::async_callbacks::TASK_WAKER_INDEX;
pub use self::async_callbacks::{
AsyncPrivateKeyMethod, AsyncPrivateKeyMethodError, AsyncSelectCertError,
BoxPrivateKeyMethodFinish, BoxPrivateKeyMethodFuture, BoxSelectCertFinish, BoxSelectCertFuture,
ExDataFuture, SslContextBuilderExt,
};
use self::bridge::AsyncStreamBridge; use self::bridge::AsyncStreamBridge;
/// Asynchronously performs a client-side TLS handshake over the provided stream. /// Asynchronously performs a client-side TLS handshake over the provided stream.
@ -90,6 +97,11 @@ impl<S> SslStream<S> {
self.0.ssl() self.0.ssl()
} }
/// Returns a mutable reference to the `Ssl` object associated with this stream.
pub fn ssl_mut(&mut self) -> &mut SslRef {
self.0.ssl_mut()
}
/// Returns a shared reference to the underlying stream. /// Returns a shared reference to the underlying stream.
pub fn get_ref(&self) -> &S { pub fn get_ref(&self) -> &S {
&self.0.get_ref().stream &self.0.get_ref().stream
@ -285,15 +297,20 @@ where
let mut mid_handshake = self.0.take().expect("future polled after completion"); let mut mid_handshake = self.0.take().expect("future polled after completion");
mid_handshake.get_mut().set_waker(Some(ctx)); mid_handshake.get_mut().set_waker(Some(ctx));
mid_handshake
.ssl_mut()
.set_ex_data(*TASK_WAKER_INDEX, Some(ctx.waker().clone()));
match mid_handshake.handshake() { match mid_handshake.handshake() {
Ok(mut stream) => { Ok(mut stream) => {
stream.get_mut().set_waker(None); stream.get_mut().set_waker(None);
stream.ssl_mut().set_ex_data(*TASK_WAKER_INDEX, None);
Poll::Ready(Ok(SslStream(stream))) Poll::Ready(Ok(SslStream(stream)))
} }
Err(ssl::HandshakeError::WouldBlock(mut mid_handshake)) => { Err(ssl::HandshakeError::WouldBlock(mut mid_handshake)) => {
mid_handshake.get_mut().set_waker(None); mid_handshake.get_mut().set_waker(None);
mid_handshake.ssl_mut().set_ex_data(*TASK_WAKER_INDEX, None);
self.0 = Some(mid_handshake); self.0 = Some(mid_handshake);

View File

@ -0,0 +1,187 @@
use boring::hash::MessageDigest;
use boring::pkey::PKey;
use boring::rsa::Padding;
use boring::sign::{RsaPssSaltlen, Signer};
use boring::ssl::{SslRef, SslSignatureAlgorithm};
use futures::future;
use tokio::task::yield_now;
use tokio_boring::{
AsyncPrivateKeyMethod, AsyncPrivateKeyMethodError, BoxPrivateKeyMethodFuture,
SslContextBuilderExt,
};
mod common;
use self::common::{connect, create_server, with_trivial_client_server_exchange};
#[allow(clippy::type_complexity)]
struct Method {
sign: Box<
dyn Fn(
&mut SslRef,
&[u8],
SslSignatureAlgorithm,
&mut [u8],
) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>
+ Send
+ Sync
+ 'static,
>,
decrypt: Box<
dyn Fn(
&mut SslRef,
&[u8],
&mut [u8],
) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>
+ Send
+ Sync
+ 'static,
>,
}
impl Method {
fn new() -> Self {
Self {
sign: Box::new(|_, _, _, _| unreachable!("called sign")),
decrypt: Box::new(|_, _, _| unreachable!("called decrypt")),
}
}
fn sign(
mut self,
sign: impl Fn(
&mut SslRef,
&[u8],
SslSignatureAlgorithm,
&mut [u8],
) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>
+ Send
+ Sync
+ 'static,
) -> Self {
self.sign = Box::new(sign);
self
}
#[allow(dead_code)]
fn decrypt(
mut self,
decrypt: impl Fn(
&mut SslRef,
&[u8],
&mut [u8],
) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError>
+ Send
+ Sync
+ 'static,
) -> Self {
self.decrypt = Box::new(decrypt);
self
}
}
impl AsyncPrivateKeyMethod for Method {
fn sign(
&self,
ssl: &mut SslRef,
input: &[u8],
signature_algorithm: SslSignatureAlgorithm,
output: &mut [u8],
) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError> {
(self.sign)(ssl, input, signature_algorithm, output)
}
fn decrypt(
&self,
ssl: &mut SslRef,
input: &[u8],
output: &mut [u8],
) -> Result<BoxPrivateKeyMethodFuture, AsyncPrivateKeyMethodError> {
(self.decrypt)(ssl, input, output)
}
}
#[tokio::test]
async fn test_sign_failure() {
with_async_private_key_method_error(
Method::new().sign(|_, _, _, _| Err(AsyncPrivateKeyMethodError)),
)
.await;
}
#[tokio::test]
async fn test_sign_future_failure() {
with_async_private_key_method_error(
Method::new().sign(|_, _, _, _| Ok(Box::pin(async { Err(AsyncPrivateKeyMethodError) }))),
)
.await;
}
#[tokio::test]
async fn test_sign_future_yield_failure() {
with_async_private_key_method_error(Method::new().sign(|_, _, _, _| {
Ok(Box::pin(async {
yield_now().await;
Err(AsyncPrivateKeyMethodError)
}))
}))
.await;
}
#[tokio::test]
async fn test_sign_ok() {
with_trivial_client_server_exchange(|builder| {
builder.set_async_private_key_method(Method::new().sign(
|_, input, signature_algorithm, _| {
assert_eq!(
signature_algorithm,
SslSignatureAlgorithm::RSA_PSS_RSAE_SHA256,
);
let input = input.to_owned();
Ok(Box::pin(async move {
Ok(Box::new(move |_: &mut SslRef, output: &mut [u8]| {
Ok(sign_with_default_config(&input, output))
}) as Box<_>)
}))
},
));
})
.await;
}
fn sign_with_default_config(input: &[u8], output: &mut [u8]) -> usize {
let pkey = PKey::private_key_from_pem(include_bytes!("key.pem")).unwrap();
let mut signer = Signer::new(MessageDigest::sha256(), &pkey).unwrap();
signer.set_rsa_padding(Padding::PKCS1_PSS).unwrap();
signer
.set_rsa_pss_saltlen(RsaPssSaltlen::DIGEST_LENGTH)
.unwrap();
signer.update(input).unwrap();
signer.sign(output).unwrap()
}
async fn with_async_private_key_method_error(method: Method) {
let (stream, addr) = create_server(move |builder| {
builder.set_async_private_key_method(method);
});
let server = async {
let _err = stream.await.unwrap_err();
};
let client = async {
let _err = connect(addr, |builder| builder.set_ca_file("tests/cert.pem"))
.await
.unwrap_err();
};
future::join(server, client).await;
}

View File

@ -0,0 +1,96 @@
use boring::ssl::ClientHello;
use futures::future;
use tokio::task::yield_now;
use tokio_boring::{
AsyncSelectCertError, BoxSelectCertFinish, BoxSelectCertFuture, SslContextBuilderExt,
};
mod common;
use self::common::{connect, create_server, with_trivial_client_server_exchange};
#[tokio::test]
async fn test_async_select_certificate_callback_trivial() {
with_trivial_client_server_exchange(|builder| {
builder.set_async_select_certificate_callback(|_| {
Ok(Box::pin(async {
Ok(Box::new(|_: ClientHello<'_>| Ok(())) as BoxSelectCertFinish)
}))
});
})
.await;
}
#[tokio::test]
async fn test_async_select_certificate_callback_yield() {
with_trivial_client_server_exchange(|builder| {
builder.set_async_select_certificate_callback(|_| {
Ok(Box::pin(async {
yield_now().await;
Ok(Box::new(|_: ClientHello<'_>| Ok(())) as BoxSelectCertFinish)
}))
});
})
.await;
}
#[tokio::test]
async fn test_async_select_certificate_callback_return_error() {
with_async_select_certificate_callback_error(|_| Err(AsyncSelectCertError)).await;
}
#[tokio::test]
async fn test_async_select_certificate_callback_future_error() {
with_async_select_certificate_callback_error(|_| {
Ok(Box::pin(async move { Err(AsyncSelectCertError) }))
})
.await;
}
#[tokio::test]
async fn test_async_select_certificate_callback_future_yield_error() {
with_async_select_certificate_callback_error(|_| {
Ok(Box::pin(async move {
yield_now().await;
Err(AsyncSelectCertError)
}))
})
.await;
}
#[tokio::test]
async fn test_async_select_certificate_callback_finish_error() {
with_async_select_certificate_callback_error(|_| {
Ok(Box::pin(async move {
yield_now().await;
Ok(Box::new(|_: ClientHello<'_>| Err(AsyncSelectCertError)) as BoxSelectCertFinish)
}))
})
.await;
}
async fn with_async_select_certificate_callback_error(
callback: impl Fn(&mut ClientHello<'_>) -> Result<BoxSelectCertFuture, AsyncSelectCertError>
+ Send
+ Sync
+ 'static,
) {
let (stream, addr) = create_server(|builder| {
builder.set_async_select_certificate_callback(callback);
});
let server = async {
let _err = stream.await.unwrap_err();
};
let client = async {
let _err = connect(addr, |builder| builder.set_ca_file("tests/cert.pem"))
.await
.unwrap_err();
};
future::join(server, client).await;
}