diff --git a/tokio-boring/src/async_callbacks.rs b/tokio-boring/src/async_callbacks.rs index ee658ea7..b12ad2b2 100644 --- a/tokio-boring/src/async_callbacks.rs +++ b/tokio-boring/src/async_callbacks.rs @@ -26,10 +26,10 @@ pub type ExDataFuture = Pin + Send + Sync>>; pub(crate) static TASK_WAKER_INDEX: Lazy>> = Lazy::new(|| Ssl::new_ex_index().unwrap()); -pub(crate) static SELECT_CERT_FUTURE_INDEX: Lazy> = +pub(crate) static SELECT_CERT_FUTURE_INDEX: Lazy>> = Lazy::new(|| Ssl::new_ex_index().unwrap()); pub(crate) static SELECT_PRIVATE_KEY_METHOD_FUTURE_INDEX: Lazy< - Index, + Index>, > = Lazy::new(|| Ssl::new_ex_index().unwrap()); /// Extensions to [`SslContextBuilder`]. @@ -219,7 +219,7 @@ fn with_private_key_method( /// created by `create_fut` returns `Poll::Ready(_)` on the first poll call. fn with_ex_data_future( ssl_handle: &mut H, - index: Index>>, + index: Index>>>, get_ssl_mut: impl Fn(&mut H) -> &mut ssl::SslRef, create_fut: impl FnOnce(&mut H) -> Result>, E>, ) -> Poll> { @@ -232,25 +232,21 @@ fn with_ex_data_future( 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)); + if let Some(data @ Some(_)) = ssl.ex_data_mut(index) { + let fut_result = ready!(data.as_mut().unwrap().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? + *data = None; - Poll::Ready(fut_result) - } - None => { - let mut fut = create_fut(ssl_handle)?; + 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(fut_result), - Poll::Pending => { - get_ssl_mut(ssl_handle).set_ex_data(index, fut); + 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, Some(fut)); - Poll::Pending - } + Poll::Pending } } }