Skip to main content

compio_tls/
rtls.rs

1use std::{
2    io,
3    pin::Pin,
4    sync::Arc,
5    task::{Context, Poll},
6};
7
8use futures_util::{AsyncRead, AsyncWrite, FutureExt};
9use rustls::{
10    ServerConfig, ServerConnection,
11    server::{Acceptor, ClientHello},
12};
13
14use crate::TlsStream;
15
16/// A lazy TLS acceptor that performs the initial handshake and allows access to
17/// the [`ClientHello`] message before completing the handshake.
18pub struct LazyConfigAcceptor<S>(futures_rustls::LazyConfigAcceptor<S>);
19
20impl<S> LazyConfigAcceptor<S>
21where
22    S: AsyncRead + AsyncWrite + Unpin,
23{
24    /// Create a new [`LazyConfigAcceptor`] with the given acceptor and stream.
25    pub fn new(acceptor: Acceptor, s: S) -> Self {
26        Self(futures_rustls::LazyConfigAcceptor::new(acceptor, s))
27    }
28}
29
30impl<S> Future for LazyConfigAcceptor<S>
31where
32    S: AsyncRead + AsyncWrite + Unpin,
33{
34    type Output = Result<StartHandshake<S>, io::Error>;
35
36    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
37        self.0.poll_unpin(cx).map_ok(StartHandshake)
38    }
39}
40
41/// A TLS acceptor that has completed the initial handshake and allows access to
42/// the [`ClientHello`] message.
43pub struct StartHandshake<S>(futures_rustls::StartHandshake<S>);
44
45impl<S> StartHandshake<S>
46where
47    S: AsyncRead + AsyncWrite + Unpin,
48{
49    /// Get the [`ClientHello`] message from the initial handshake.
50    pub fn client_hello(&self) -> ClientHello<'_> {
51        self.0.client_hello()
52    }
53
54    /// Complete the TLS handshake and return a [`TlsStream`] if successful.
55    pub fn into_stream(
56        self,
57        config: Arc<ServerConfig>,
58    ) -> impl Future<Output = io::Result<TlsStream<S>>> {
59        self.into_stream_with(config, |_| ())
60    }
61
62    /// Complete the TLS handshake and return a [`TlsStream`] if successful.
63    pub fn into_stream_with<F>(
64        self,
65        config: Arc<ServerConfig>,
66        f: F,
67    ) -> impl Future<Output = io::Result<TlsStream<S>>>
68    where
69        F: FnOnce(&mut ServerConnection),
70    {
71        self.0
72            .into_stream_with(config, f)
73            .map(|res| res.map(TlsStream::from))
74    }
75}