From 1a55d51289da37a1adbd5f5581f2c66babc14d20 Mon Sep 17 00:00:00 2001 From: Yuyi Wang Date: Sat, 14 Feb 2026 11:54:08 +0800 Subject: [PATCH] feat(tls): add LazyConfigAcceptor for rustls --- compio-tls/src/lib.rs | 5 +++ compio-tls/src/rtls.rs | 70 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+) create mode 100644 compio-tls/src/rtls.rs diff --git a/compio-tls/src/lib.rs b/compio-tls/src/lib.rs index ea6624f3..e95eecaf 100644 --- a/compio-tls/src/lib.rs +++ b/compio-tls/src/lib.rs @@ -23,3 +23,8 @@ mod stream; pub use adapter::*; pub use maybe::*; pub use stream::*; + +#[cfg(feature = "rustls")] +mod rtls; +#[cfg(feature = "rustls")] +pub use rtls::*; diff --git a/compio-tls/src/rtls.rs b/compio-tls/src/rtls.rs new file mode 100644 index 00000000..db6bcc57 --- /dev/null +++ b/compio-tls/src/rtls.rs @@ -0,0 +1,70 @@ +use std::{ + io, + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; + +use compio_io::{AsyncRead, AsyncWrite, compat::AsyncStream}; +use futures_util::FutureExt; +use rustls::{ + ServerConfig, ServerConnection, + server::{Acceptor, ClientHello}, +}; + +use crate::TlsStream; + +/// A lazy TLS acceptor that performs the initial handshake and allows access to +/// the [`ClientHello`] message before completing the handshake. +pub struct LazyConfigAcceptor(futures_rustls::LazyConfigAcceptor>); + +impl LazyConfigAcceptor { + /// Create a new [`LazyConfigAcceptor`] with the given acceptor and stream. + pub fn new(acceptor: Acceptor, s: S) -> Self { + Self(futures_rustls::LazyConfigAcceptor::new( + acceptor, + AsyncStream::new(s), + )) + } +} + +impl Future for LazyConfigAcceptor { + type Output = Result, io::Error>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.0.poll_unpin(cx).map_ok(StartHandshake) + } +} + +/// A TLS acceptor that has completed the initial handshake and allows access to +/// the [`ClientHello`] message. +pub struct StartHandshake(futures_rustls::StartHandshake>); + +impl StartHandshake { + /// Get the [`ClientHello`] message from the initial handshake. + pub fn client_hello(&self) -> ClientHello<'_> { + self.0.client_hello() + } + + /// Complete the TLS handshake and return a [`TlsStream`] if successful. + pub fn into_stream( + self, + config: Arc, + ) -> impl Future>> { + self.into_stream_with(config, |_| ()) + } + + /// Complete the TLS handshake and return a [`TlsStream`] if successful. + pub fn into_stream_with( + self, + config: Arc, + f: F, + ) -> impl Future>> + where + F: FnOnce(&mut ServerConnection), + { + self.0 + .into_stream_with(config, f) + .map(|res| res.map(TlsStream::from)) + } +}