tokio_rustls/
server.rs

1use std::future::Future;
2use std::io::{self, BufRead as _};
3#[cfg(unix)]
4use std::os::unix::io::{AsRawFd, RawFd};
5#[cfg(windows)]
6use std::os::windows::io::{AsRawSocket, RawSocket};
7use std::pin::Pin;
8use std::sync::Arc;
9use std::task::{Context, Poll};
10
11use rustls::server::AcceptedAlert;
12use rustls::{ServerConfig, ServerConnection};
13use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
14
15use crate::common::{IoSession, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, TlsState};
16
17/// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method.
18#[derive(Clone)]
19pub struct TlsAcceptor {
20    inner: Arc<ServerConfig>,
21}
22
23impl From<Arc<ServerConfig>> for TlsAcceptor {
24    fn from(inner: Arc<ServerConfig>) -> Self {
25        Self { inner }
26    }
27}
28
29impl TlsAcceptor {
30    /// Returns a future for completing a TLS handshake for a client using `stream`.
31    ///
32    /// You likely want to wrap this in a timeout (for example with [`tokio::time::timeout`][])
33    /// to bound the handshake time.
34    ///
35    /// [`tokio::time::timeout`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout.html
36    #[inline]
37    pub fn accept<IO>(&self, stream: IO) -> Accept<IO>
38    where
39        IO: AsyncRead + AsyncWrite + Unpin,
40    {
41        self.accept_with(stream, |_| ())
42    }
43
44    /// Similar to [`Self::accept()`], but calls `f` before performing the handshake.
45    ///
46    /// As with [`Self::accept()`] you likely want to wrap this in a timeout to
47    /// bound the handshake time.
48    ///
49    /// The `f` handler is given a mutable reference to a [`ServerConnection`][] that can be used
50    /// to configure the connection before the handshake, for example, adjusting the buffer limit.
51    ///
52    /// Because no data has been read from `stream` yet when `f` is called ClientHello
53    /// dependent state (like early data) is not yet available.
54    ///
55    /// [`ServerConnection`]: https://docs.rs/rustls/latest/rustls/server/struct.ServerConnection.html
56    pub fn accept_with<IO, F>(&self, stream: IO, f: F) -> Accept<IO>
57    where
58        IO: AsyncRead + AsyncWrite + Unpin,
59        F: FnOnce(&mut ServerConnection),
60    {
61        let mut session = match ServerConnection::new(self.inner.clone()) {
62            Ok(session) => session,
63            Err(error) => {
64                return Accept(MidHandshake::Error {
65                    io: stream,
66                    // TODO(eliza): should this really return an `io::Error`?
67                    // Probably not...
68                    error: io::Error::new(io::ErrorKind::Other, error),
69                });
70            }
71        };
72        f(&mut session);
73
74        Accept(MidHandshake::Handshaking(TlsStream {
75            session,
76            io: stream,
77            state: TlsState::Stream,
78            need_flush: false,
79        }))
80    }
81
82    /// Get a read-only reference to underlying config
83    pub fn config(&self) -> &Arc<ServerConfig> {
84        &self.inner
85    }
86}
87
88/// A future for reading a `ClientHello` from `io` without committing to a [`ServerConfig`][].
89///
90/// Awaiting it yields a [`StartHandshake`], which exposes the
91/// [`ClientHello`][] (for example, to choose a config based on SNI) and performs
92/// the rest of the handshake via [`StartHandshake::into_stream()`].
93///
94/// [`ServerConfig`]: https://docs.rs/rustls/latest/rustls/server/struct.ServerConfig.html
95/// [`ClientHello`]: https://docs.rs/rustls/latest/rustls/server/struct.ClientHello.html
96pub struct LazyConfigAcceptor<IO> {
97    acceptor: rustls::server::Acceptor,
98    io: Option<IO>,
99    alert: Option<(rustls::Error, AcceptedAlert)>,
100}
101
102impl<IO> LazyConfigAcceptor<IO>
103where
104    IO: AsyncRead + AsyncWrite + Unpin,
105{
106    /// Returns a new `LazyConfigAcceptor` that reads a `ClientHello` from `io`.
107    ///
108    /// You likely want to wrap awaiting the acceptor in a timeout to bound how long the
109    /// peer may take to send the `ClientHello`.
110    ///
111    /// Note that awaiting the acceptor is only the first half of the handshake and
112    /// [`StartHandshake::into_stream()`] performs the rest.
113    ///
114    /// To bound the time for the complete handshake, share one deadline across
115    /// both awaits (for example with [`tokio::time::timeout_at`][]) rather than giving each
116    /// its own timeout.
117    ///
118    /// If a timeout elapses before the `ClientHello` arrives, [`Self::take_io()`] can
119    /// recover the `io`, for example to answer the peer in plaintext before closing.
120    ///
121    /// [`tokio::time::timeout_at`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout_at.html
122    #[inline]
123    pub fn new(acceptor: rustls::server::Acceptor, io: IO) -> Self {
124        Self {
125            acceptor,
126            io: Some(io),
127            alert: None,
128        }
129    }
130
131    /// Takes back the client connection. Will return `None` if called more than once or if the
132    /// connection has been accepted.
133    ///
134    /// # Example
135    ///
136    /// ```no_run
137    /// # fn choose_server_config(
138    /// #     _: rustls::server::ClientHello,
139    /// # ) -> std::sync::Arc<rustls::ServerConfig> {
140    /// #     unimplemented!();
141    /// # }
142    /// # #[allow(unused_variables)]
143    /// # async fn listen() {
144    /// use tokio::io::AsyncWriteExt;
145    /// let listener = tokio::net::TcpListener::bind("127.0.0.1:4443").await.unwrap();
146    /// let (stream, _) = listener.accept().await.unwrap();
147    ///
148    /// let acceptor = tokio_rustls::LazyConfigAcceptor::new(rustls::server::Acceptor::default(), stream);
149    /// tokio::pin!(acceptor);
150    ///
151    /// match acceptor.as_mut().await {
152    ///     Ok(start) => {
153    ///         let clientHello = start.client_hello();
154    ///         let config = choose_server_config(clientHello);
155    ///         let stream = start.into_stream(config).await.unwrap();
156    ///         // Proceed with handling the ServerConnection...
157    ///     }
158    ///     Err(err) => {
159    ///         if let Some(mut stream) = acceptor.take_io() {
160    ///             stream
161    ///                 .write_all(
162    ///                     format!("HTTP/1.1 400 Invalid Input\r\n\r\n\r\n{:?}\n", err)
163    ///                         .as_bytes()
164    ///                 )
165    ///                 .await
166    ///                 .unwrap();
167    ///         }
168    ///     }
169    /// }
170    /// # }
171    /// ```
172    pub fn take_io(&mut self) -> Option<IO> {
173        self.io.take()
174    }
175}
176
177impl<IO> Future for LazyConfigAcceptor<IO>
178where
179    IO: AsyncRead + AsyncWrite + Unpin,
180{
181    type Output = Result<StartHandshake<IO>, io::Error>;
182
183    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
184        let this = self.get_mut();
185        loop {
186            let io = match this.io.as_mut() {
187                Some(io) => io,
188                None => {
189                    return Poll::Ready(Err(io::Error::new(
190                        io::ErrorKind::Other,
191                        "acceptor cannot be polled after acceptance",
192                    )));
193                }
194            };
195
196            if let Some((err, mut alert)) = this.alert.take() {
197                match alert.write(&mut SyncWriteAdapter { io, cx }) {
198                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
199                        this.alert = Some((err, alert));
200                        return Poll::Pending;
201                    }
202                    Ok(0) | Err(_) => {
203                        return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err)));
204                    }
205                    Ok(_) => {
206                        this.alert = Some((err, alert));
207                        continue;
208                    }
209                };
210            }
211
212            let mut reader = SyncReadAdapter { io, cx };
213            match this.acceptor.read_tls(&mut reader) {
214                Ok(0) => return Err(io::ErrorKind::UnexpectedEof.into()).into(),
215                Ok(_) => {}
216                Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Poll::Pending,
217                Err(e) => return Err(e).into(),
218            }
219
220            match this.acceptor.accept() {
221                Ok(Some(accepted)) => {
222                    let io = this.io.take().unwrap();
223                    return Poll::Ready(Ok(StartHandshake { accepted, io }));
224                }
225                Ok(None) => {}
226                Err((err, alert)) => {
227                    this.alert = Some((err, alert));
228                }
229            }
230        }
231    }
232}
233
234/// An incoming connection received through [`LazyConfigAcceptor`].
235///
236/// This contains the generic `IO` asynchronous transport,
237/// [`ClientHello`](rustls::server::ClientHello) data,
238/// and all the state required to continue the TLS handshake (e.g. via
239/// [`StartHandshake::into_stream`]).
240#[non_exhaustive]
241#[derive(Debug)]
242pub struct StartHandshake<IO> {
243    pub accepted: rustls::server::Accepted,
244    pub io: IO,
245}
246
247impl<IO> StartHandshake<IO>
248where
249    IO: AsyncRead + AsyncWrite + Unpin,
250{
251    /// Create a new object from an `IO` transport and prior TLS metadata.
252    pub fn from_parts(accepted: rustls::server::Accepted, transport: IO) -> Self {
253        Self {
254            accepted,
255            io: transport,
256        }
257    }
258
259    pub fn client_hello(&self) -> rustls::server::ClientHello<'_> {
260        self.accepted.client_hello()
261    }
262
263    /// Returns a future that performs the rest of the TLS handshake using `config`.
264    ///
265    /// You likely want to wrap this in a timeout to bound the handshake time. Ideally
266    /// with [`tokio::time::timeout_at`][], reusing the deadline that also bounded
267    /// awaiting the [`LazyConfigAcceptor`] so both halves of the handshake share
268    /// one budget. See [`LazyConfigAcceptor::new()`].
269    ///
270    /// [`tokio::time::timeout_at`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout_at.html
271    pub fn into_stream(self, config: Arc<ServerConfig>) -> Accept<IO> {
272        self.into_stream_with(config, |_| ())
273    }
274
275    /// Similar to [`Self::into_stream()`], but calls `f` before performing the handshake.
276    ///
277    /// As with [`Self::into_stream()`] you likely want to wrap this in a timeout to
278    /// bound the handshake time.
279    ///
280    /// The `f` handler is given a mutable reference to a [`ServerConnection`][] that can be
281    /// used to configure the connection before the handshake.
282    ///
283    /// [`ServerConnection`]: https://docs.rs/rustls/latest/rustls/server/struct.ServerConnection.html
284    pub fn into_stream_with<F>(self, config: Arc<ServerConfig>, f: F) -> Accept<IO>
285    where
286        F: FnOnce(&mut ServerConnection),
287    {
288        let mut conn = match self.accepted.into_connection(config) {
289            Ok(conn) => conn,
290            Err((error, alert)) => {
291                return Accept(MidHandshake::SendAlert {
292                    io: self.io,
293                    alert,
294                    // TODO(eliza): should this really return an `io::Error`?
295                    // Probably not...
296                    error: io::Error::new(io::ErrorKind::InvalidData, error),
297                });
298            }
299        };
300        f(&mut conn);
301
302        Accept(MidHandshake::Handshaking(TlsStream {
303            session: conn,
304            io: self.io,
305            state: TlsState::Stream,
306            need_flush: false,
307        }))
308    }
309}
310
311/// Future returned from `TlsAcceptor::accept` which will resolve
312/// once the accept handshake has finished.
313pub struct Accept<IO>(MidHandshake<TlsStream<IO>>);
314
315impl<IO> Accept<IO> {
316    #[inline]
317    pub fn into_fallible(self) -> FallibleAccept<IO> {
318        FallibleAccept(self.0)
319    }
320
321    pub fn get_ref(&self) -> Option<&IO> {
322        match &self.0 {
323            MidHandshake::Handshaking(sess) => Some(sess.get_ref().0),
324            MidHandshake::SendAlert { io, .. } => Some(io),
325            MidHandshake::Error { io, .. } => Some(io),
326            MidHandshake::End => None,
327        }
328    }
329
330    pub fn get_mut(&mut self) -> Option<&mut IO> {
331        match &mut self.0 {
332            MidHandshake::Handshaking(sess) => Some(sess.get_mut().0),
333            MidHandshake::SendAlert { io, .. } => Some(io),
334            MidHandshake::Error { io, .. } => Some(io),
335            MidHandshake::End => None,
336        }
337    }
338}
339
340impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Accept<IO> {
341    type Output = io::Result<TlsStream<IO>>;
342
343    #[inline]
344    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
345        Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err)
346    }
347}
348
349/// Like [Accept], but returns `IO` on failure.
350pub struct FallibleAccept<IO>(MidHandshake<TlsStream<IO>>);
351
352impl<IO: AsyncRead + AsyncWrite + Unpin> Future for FallibleAccept<IO> {
353    type Output = Result<TlsStream<IO>, (io::Error, IO)>;
354
355    #[inline]
356    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
357        Pin::new(&mut self.0).poll(cx)
358    }
359}
360
361/// A wrapper around an underlying raw stream which implements the TLS or SSL
362/// protocol.
363#[derive(Debug)]
364pub struct TlsStream<IO> {
365    pub(crate) io: IO,
366    pub(crate) session: ServerConnection,
367    pub(crate) state: TlsState,
368    pub(crate) need_flush: bool,
369}
370
371impl<IO> TlsStream<IO> {
372    #[inline]
373    pub fn get_ref(&self) -> (&IO, &ServerConnection) {
374        (&self.io, &self.session)
375    }
376
377    #[inline]
378    pub fn get_mut(&mut self) -> (&mut IO, &mut ServerConnection) {
379        (&mut self.io, &mut self.session)
380    }
381
382    #[inline]
383    pub fn into_inner(self) -> (IO, ServerConnection) {
384        (self.io, self.session)
385    }
386}
387
388impl<IO> IoSession for TlsStream<IO> {
389    type Io = IO;
390    type Session = ServerConnection;
391
392    #[inline]
393    fn skip_handshake(&self) -> bool {
394        false
395    }
396
397    #[inline]
398    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session, &mut bool) {
399        (
400            &mut self.state,
401            &mut self.io,
402            &mut self.session,
403            &mut self.need_flush,
404        )
405    }
406
407    #[inline]
408    fn into_io(self) -> Self::Io {
409        self.io
410    }
411}
412
413impl<IO> AsyncRead for TlsStream<IO>
414where
415    IO: AsyncRead + AsyncWrite + Unpin,
416{
417    fn poll_read(
418        mut self: Pin<&mut Self>,
419        cx: &mut Context<'_>,
420        buf: &mut ReadBuf<'_>,
421    ) -> Poll<io::Result<()>> {
422        let data = ready!(self.as_mut().poll_fill_buf(cx))?;
423        let len = data.len().min(buf.remaining());
424        if len == 0 {
425            return Poll::Ready(Ok(()));
426        }
427        buf.put_slice(&data[..len]);
428        self.as_mut().consume(len);
429
430        while buf.remaining() > 0 {
431            let data = match self.as_mut().poll_fill_buf(cx) {
432                Poll::Ready(Ok([])) => break,
433                Poll::Ready(Ok(data)) => data,
434                Poll::Ready(Err(_)) => break, // non-transient error gets re-emitted next poll
435                Poll::Pending => break,
436            };
437            let len = Ord::min(data.len(), buf.remaining());
438            buf.put_slice(&data[..len]);
439            self.as_mut().consume(len);
440        }
441        Poll::Ready(Ok(()))
442    }
443}
444
445impl<IO> AsyncBufRead for TlsStream<IO>
446where
447    IO: AsyncRead + AsyncWrite + Unpin,
448{
449    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
450        match self.state {
451            TlsState::Stream | TlsState::WriteShutdown => {
452                let this = self.get_mut();
453                let stream =
454                    Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
455
456                match stream.poll_fill_buf(cx) {
457                    Poll::Ready(Ok(buf)) => {
458                        if buf.is_empty() {
459                            this.state.shutdown_read();
460                        }
461
462                        Poll::Ready(Ok(buf))
463                    }
464                    Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::ConnectionAborted => {
465                        this.state.shutdown_read();
466                        Poll::Ready(Err(err))
467                    }
468                    output => output,
469                }
470            }
471            TlsState::ReadShutdown | TlsState::FullyShutdown => Poll::Ready(Ok(&[])),
472            #[cfg(feature = "early-data")]
473            ref s => unreachable!("server TLS can not hit this state: {:?}", s),
474        }
475    }
476
477    fn consume(mut self: Pin<&mut Self>, amt: usize) {
478        self.session.reader().consume(amt);
479    }
480}
481
482impl<IO> AsyncWrite for TlsStream<IO>
483where
484    IO: AsyncRead + AsyncWrite + Unpin,
485{
486    /// Note: that it does not guarantee the final data to be sent.
487    /// To be cautious, you must manually call `flush`.
488    fn poll_write(
489        self: Pin<&mut Self>,
490        cx: &mut Context<'_>,
491        buf: &[u8],
492    ) -> Poll<io::Result<usize>> {
493        let this = self.get_mut();
494        let mut stream =
495            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
496        stream.as_mut_pin().poll_write(cx, buf)
497    }
498
499    /// Note: that it does not guarantee the final data to be sent.
500    /// To be cautious, you must manually call `flush`.
501    fn poll_write_vectored(
502        self: Pin<&mut Self>,
503        cx: &mut Context<'_>,
504        bufs: &[io::IoSlice<'_>],
505    ) -> Poll<io::Result<usize>> {
506        let this = self.get_mut();
507        let mut stream =
508            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
509        stream.as_mut_pin().poll_write_vectored(cx, bufs)
510    }
511
512    #[inline]
513    fn is_write_vectored(&self) -> bool {
514        true
515    }
516
517    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
518        let this = self.get_mut();
519        let mut stream =
520            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
521        stream.as_mut_pin().poll_flush(cx)
522    }
523
524    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
525        if self.state.writeable() {
526            self.session.send_close_notify();
527            self.state.shutdown_write();
528        }
529
530        let this = self.get_mut();
531        let mut stream =
532            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
533        stream.as_mut_pin().poll_shutdown(cx)
534    }
535}
536
537#[cfg(unix)]
538impl<IO> AsRawFd for TlsStream<IO>
539where
540    IO: AsRawFd,
541{
542    fn as_raw_fd(&self) -> RawFd {
543        self.get_ref().0.as_raw_fd()
544    }
545}
546
547#[cfg(windows)]
548impl<IO> AsRawSocket for TlsStream<IO>
549where
550    IO: AsRawSocket,
551{
552    fn as_raw_socket(&self) -> RawSocket {
553        self.get_ref().0.as_raw_socket()
554    }
555}
tokio_rustls/server.rs

tokio_rustls/
server.rs
