server: extend handshake timeout to LazyConfigAcceptor + into_stream

Author: cpu

src/server.rs

@@ -12,9 +12,11 @@ use std::time::Duration;
 use rustls::server::AcceptedAlert;
 use rustls::{ServerConfig, ServerConnection};
 use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
+use tokio::time::Instant;
 
 use crate::common::{
-    HandshakeFuture, IoSession, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, TlsState,
+    HandshakeFuture, IoSession, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, Timeout,
+    TlsState,
 };
 
 /// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method.
@@ -39,9 +41,11 @@ impl TlsAcceptor {
     /// `None` disables the handshake timeout.
     ///
     /// The timeout applies to handshakes started via [`TlsAcceptor::accept`] and
-    /// [`TlsAcceptor::accept_with`]. It does not apply to handshakes resumed via
-    /// [`LazyConfigAcceptor`] / [`StartHandshake::into_stream`], which proceed
-    /// without a timeout.
+    /// [`TlsAcceptor::accept_with`]. For the lazy-acceptor flow, set the timeout
+    /// on the [`LazyConfigAcceptor`] itself via
+    /// [`LazyConfigAcceptor::with_handshake_timeout`]. The deadline established
+    /// there is inherited by the [`Accept`] returned from
+    /// [`StartHandshake::into_stream`], so a single timeout covers both phases.
     pub fn with_handshake_timeout(mut self, timeout: Option<Duration>) -> Self {
         self.handshake_timeout = timeout;
         self
@@ -97,6 +101,7 @@ pub struct LazyConfigAcceptor<IO> {
     acceptor: rustls::server::Acceptor,
     io: Option<IO>,
     alert: Option<(rustls::Error, AcceptedAlert)>,
+    timeout: Option<Timeout>,
 }
 
 impl<IO> LazyConfigAcceptor<IO>
@@ -109,9 +114,30 @@ where
             acceptor,
             io: Some(io),
             alert: None,
+            timeout: None,
         }
     }
 
+    /// Set the maximum amount of time to allow for the TLS handshake.
+    ///
+    /// `None` disables the handshake timeout.
+    ///
+    /// The deadline is fixed to `Instant::now() + duration` when this builder
+    /// is called and spans both phases of the lazy-acceptor flow:
+    ///
+    /// 1. The ClientHello phase driven by this `LazyConfigAcceptor`.
+    /// 2. The post-ClientHello phase driven by the [`Accept`] returned from
+    ///    [`StartHandshake::into_stream`].
+    ///
+    /// Both phases race against the same wall-clock deadline, so a single
+    /// timeout covers the full handshake. To set a timeout on handshakes
+    /// driven via [`TlsAcceptor::accept`] instead, use
+    /// [`TlsAcceptor::with_handshake_timeout`].
+    pub fn with_handshake_timeout(mut self, timeout: Option<Duration>) -> Self {
+        self.timeout = timeout.map(Timeout::from_duration);
+        self
+    }
+
     /// Takes back the client connection. Will return `None` if called more than once or if the
     /// connection has been accepted.
     ///
@@ -181,7 +207,7 @@ where
                 match alert.write(&mut SyncWriteAdapter { io, cx }) {
                     Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                         this.alert = Some((err, alert));
-                        return Poll::Pending;
+                        return poll_pending_or_timeout(&mut this.timeout, cx);
                     }
                     Ok(0) | Err(_) => {
                         return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err)));
@@ -197,14 +223,20 @@ where
             match this.acceptor.read_tls(&mut reader) {
                 Ok(0) => return Err(io::ErrorKind::UnexpectedEof.into()).into(),
                 Ok(_) => {}
-                Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Poll::Pending,
+                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                    return poll_pending_or_timeout(&mut this.timeout, cx);
+                }
                 Err(e) => return Err(e).into(),
             }
 
             match this.acceptor.accept() {
                 Ok(Some(accepted)) => {
                     let io = this.io.take().unwrap();
-                    return Poll::Ready(Ok(StartHandshake { accepted, io }));
+                    return Poll::Ready(Ok(StartHandshake {
+                        accepted,
+                        io,
+                        deadline: this.timeout.as_ref().map(Timeout::deadline),
+                    }));
                 }
                 Ok(None) => {}
                 Err((err, alert)) => {
@@ -215,6 +247,25 @@ where
     }
 }
 
+/// Returns `Poll::Pending`, unless the timeout has elapsed, in which case a `TimedOut`
+/// error is returned. With no timeout configured, always returns `Poll::Pending`.
+fn poll_pending_or_timeout<T>(
+    timeout: &mut Option<Timeout>,
+    cx: &mut Context<'_>,
+) -> Poll<Result<T, io::Error>> {
+    let timeout = match timeout {
+        Some(timeout) => timeout,
+        None => return Poll::Pending,
+    };
+    if timeout.poll_deadline(cx).is_pending() {
+        return Poll::Pending;
+    }
+    Poll::Ready(Err(io::Error::new(
+        io::ErrorKind::TimedOut,
+        "TLS handshake timed out",
+    )))
+}
+
 /// An incoming connection received through [`LazyConfigAcceptor`].
 ///
 /// This contains the generic `IO` asynchronous transport,
@@ -226,6 +277,10 @@ where
 pub struct StartHandshake<IO> {
     pub accepted: rustls::server::Accepted,
     pub io: IO,
+    // Handshake deadline inherited from the LazyConfigAcceptor, if any.
+    // Propagated into the Accept returned by into_stream so the timeout
+    // spans both phases of the handshake.
+    pub(crate) deadline: Option<Instant>,
 }
 
 impl<IO> StartHandshake<IO>
@@ -237,6 +292,7 @@ where
         Self {
             accepted,
             io: transport,
+            deadline: None,
         }
     }
 
@@ -255,28 +311,28 @@ where
         let mut conn = match self.accepted.into_connection(config) {
             Ok(conn) => conn,
             Err((error, alert)) => {
-                return Accept(HandshakeFuture::new(
+                return Accept(HandshakeFuture::from_deadline(
                     MidHandshake::SendAlert {
                         io: self.io,
                         alert,
                         // TODO(eliza): should this really return an `io::Error`?
                         // Probably not...
                         error: io::Error::new(io::ErrorKind::InvalidData, error),
                     },
-                    None,
+                    self.deadline,
                 ));
             }
         };
         f(&mut conn);
 
-        Accept(HandshakeFuture::new(
+        Accept(HandshakeFuture::from_deadline(
             MidHandshake::Handshaking(TlsStream {
                 session: conn,
                 io: self.io,
                 state: TlsState::Stream,
                 need_flush: false,
             }),
-            None,
+            self.deadline,
         ))
     }
 }

tests/test.rs

@@ -206,6 +206,58 @@ async fn accept_handshake_timeout_fallible() {
     assert_eq!(err.kind(), ErrorKind::TimedOut);
 }
 
+#[tokio::test(start_paused = true)]
+async fn lazy_config_acceptor_handshake_timeout() {
+    // The client end is held but never written to. The lazy acceptor will block
+    // waiting on the ClientHello and with all tasks pending, the start_paused
+    // runtime advances to the handshake timeout's deadline.
+    let (_client, server) = tokio::io::duplex(4096);
+    let acceptor = LazyConfigAcceptor::new(rustls::server::Acceptor::default(), server)
+        .with_handshake_timeout(Some(HANDSHAKE_TIMEOUT));
+
+    let err = acceptor.await.unwrap_err();
+    assert_eq!(err.kind(), ErrorKind::TimedOut);
+}
+
+/// lazy_config_acceptor_handshake_timeout_into_stream confirms the deadline established
+/// on a LazyConfigAcceptor propagates through StartHandshake::into_stream.
+///
+/// A valid ClientHello is pre-loaded into the duplex and the lazy acceptor
+/// parses it instantly. Afterwards, into_stream returns an Accept that blocks trying to
+/// read the next client message which will never arrive. With all tasks pending,
+/// the start_paused runtime advances to the inherited deadline.
+#[tokio::test(start_paused = true)]
+async fn lazy_config_acceptor_handshake_timeout_into_stream() {
+    // Real TLS 1.2 ClientHello, sourced from https://tls12.xargs.org/#client-hello/annotated
+    // (the same bytes used by acceptor_alert, without the minor-version tweak).
+    let client_hello = [
+        0x16, 0x03, 0x01, 0x00, 0xa5, 0x01, 0x00, 0x00, 0xa1, 0x03, 0x03, 0x00, 0x01, 0x02, 0x03,
+        0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12,
+        0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x00, 0x00,
+        0x20, 0xcc, 0xa8, 0xcc, 0xa9, 0xc0, 0x2f, 0xc0, 0x30, 0xc0, 0x2b, 0xc0, 0x2c, 0xc0, 0x13,
+        0xc0, 0x09, 0xc0, 0x14, 0xc0, 0x0a, 0x00, 0x9c, 0x00, 0x9d, 0x00, 0x2f, 0x00, 0x35, 0xc0,
+        0x12, 0x00, 0x0a, 0x01, 0x00, 0x00, 0x58, 0x00, 0x00, 0x00, 0x18, 0x00, 0x16, 0x00, 0x00,
+        0x13, 0x65, 0x78, 0x61, 0x6d, 0x70, 0x6c, 0x65, 0x2e, 0x75, 0x6c, 0x66, 0x68, 0x65, 0x69,
+        0x6d, 0x2e, 0x6e, 0x65, 0x74, 0x00, 0x05, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x0a, 0x00, 0x0a, 0x00, 0x08, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x18, 0x00, 0x19, 0x00, 0x0b,
+        0x00, 0x02, 0x01, 0x00, 0x00, 0x0d, 0x00, 0x12, 0x00, 0x10, 0x04, 0x01, 0x04, 0x03, 0x05,
+        0x01, 0x05, 0x03, 0x06, 0x01, 0x06, 0x03, 0x02, 0x01, 0x02, 0x03, 0xff, 0x01, 0x00, 0x01,
+        0x00, 0x00, 0x12, 0x00, 0x00,
+    ];
+
+    let (sconfig, _) = utils::make_configs();
+    // Size generously so the server's ServerHello + cert chain + ServerHelloDone fit
+    // without backpressuring the write side.
+    let (mut cstream, sstream) = tokio::io::duplex(16 * 1024);
+    cstream.write_all(&client_hello).await.unwrap();
+
+    let acceptor = LazyConfigAcceptor::new(rustls::server::Acceptor::default(), sstream)
+        .with_handshake_timeout(Some(HANDSHAKE_TIMEOUT));
+    let start = acceptor.await.unwrap();
+    let err = start.into_stream(Arc::new(sconfig)).await.unwrap_err();
+    assert_eq!(err.kind(), ErrorKind::TimedOut);
+}
+
 const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);
 
 #[tokio::test]