From a9ba831508f23b936e01f8dc455e1541214afe45 Mon Sep 17 00:00:00 2001 From: Lucas Tembras Date: Mon, 8 Jun 2026 14:00:22 -0400 Subject: [PATCH 1/5] feat: added tls handshake timeout --- Cargo.toml | 4 +- src/client.rs | 77 ++++++++++++++++++++++++--------------- src/common/handshake.rs | 10 +++++ src/common/mod.rs | 2 + src/common/timeout.rs | 81 +++++++++++++++++++++++++++++++++++++++++ tests/test.rs | 38 +++++++++++++++++++ 6 files changed, 180 insertions(+), 32 deletions(-) create mode 100644 src/common/timeout.rs diff --git a/Cargo.toml b/Cargo.toml index d917a5cb..a63dd9aa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,7 @@ exclude = ["/.github", "/examples", "/scripts", "/tests/"] [dependencies] rustls = { version = "0.23.27", default-features = false, features = ["std"] } -tokio = "1.0" +tokio = { version = "1.0", features = ["time"] } [features] default = ["logging", "tls12", "aws_lc_rs"] @@ -33,7 +33,7 @@ argh = "0.1.1" futures-util = "0.3.1" lazy_static = "1.1" rcgen = { version = "0.14", features = ["pem"] } -tokio = { version = "1.0", features = ["full"] } +tokio = { version = "1.0", features = ["full", "test-util"] } webpki-roots = "1" [package.metadata.cargo_check_external_types] diff --git a/src/client.rs b/src/client.rs index a94b64a8..d4383b56 100644 --- a/src/client.rs +++ b/src/client.rs @@ -9,22 +9,32 @@ use std::sync::Arc; #[cfg(feature = "early-data")] use std::task::Waker; use std::task::{Context, Poll}; +use std::time::Duration; use rustls::pki_types::ServerName; use rustls::{ClientConfig, ClientConnection}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; -use crate::common::{IoSession, MidHandshake, Stream, TlsState}; +use crate::common::{HandshakeFuture, IoSession, MidHandshake, Stream, TlsState}; /// A wrapper around a `rustls::ClientConfig`, providing an async `connect` method. #[derive(Clone)] pub struct TlsConnector { inner: Arc, + handshake_timeout: Option, #[cfg(feature = "early-data")] early_data: bool, } impl TlsConnector { + /// Set the maximum amount of time to allow for TLS handshakes. + /// + /// `None` disables the handshake timeout. + pub fn with_handshake_timeout(mut self, timeout: Option) -> Self { + self.handshake_timeout = timeout; + self + } + /// Enable 0-RTT. /// /// If you want to use 0-RTT, @@ -67,36 +77,42 @@ impl TlsConnector { let mut session = match ClientConnection::new_with_alpn(self.inner.clone(), domain, alpn) { Ok(session) => session, Err(error) => { - return Connect(MidHandshake::Error { - io: stream, - // TODO(eliza): should this really return an `io::Error`? - // Probably not... - error: io::Error::new(io::ErrorKind::Other, error), - }); + return Connect(HandshakeFuture::new( + MidHandshake::Error { + io: stream, + // TODO(eliza): should this really return an `io::Error`? + // Probably not... + error: io::Error::new(io::ErrorKind::Other, error), + }, + self.handshake_timeout, + )); } }; f(&mut session); - Connect(MidHandshake::Handshaking(TlsStream { - io: stream, + Connect(HandshakeFuture::new( + MidHandshake::Handshaking(TlsStream { + io: stream, - #[cfg(not(feature = "early-data"))] - state: TlsState::Stream, + #[cfg(not(feature = "early-data"))] + state: TlsState::Stream, - #[cfg(feature = "early-data")] - state: if self.early_data && session.early_data().is_some() { - TlsState::EarlyData(0, Vec::new()) - } else { - TlsState::Stream - }, + #[cfg(feature = "early-data")] + state: if self.early_data && session.early_data().is_some() { + TlsState::EarlyData(0, Vec::new()) + } else { + TlsState::Stream + }, - need_flush: false, + need_flush: false, - #[cfg(feature = "early-data")] - early_waker: None, + #[cfg(feature = "early-data")] + early_waker: None, - session, - })) + session, + }), + self.handshake_timeout, + )) } pub fn with_alpn(&self, alpn_protocols: Vec>) -> TlsConnectorWithAlpn<'_> { @@ -116,6 +132,7 @@ impl From> for TlsConnector { fn from(inner: Arc) -> Self { Self { inner, + handshake_timeout: None, #[cfg(feature = "early-data")] early_data: false, } @@ -150,7 +167,7 @@ impl TlsConnectorWithAlpn<'_> { /// Future returned from `TlsConnector::connect` which will resolve /// once the connection handshake has finished. -pub struct Connect(MidHandshake>); +pub struct Connect(HandshakeFuture>); impl Connect { #[inline] @@ -159,7 +176,7 @@ impl Connect { } pub fn get_ref(&self) -> Option<&IO> { - match &self.0 { + match self.0.handshake() { MidHandshake::Handshaking(sess) => Some(sess.get_ref().0), MidHandshake::SendAlert { io, .. } => Some(io), MidHandshake::Error { io, .. } => Some(io), @@ -168,7 +185,7 @@ impl Connect { } pub fn get_mut(&mut self) -> Option<&mut IO> { - match &mut self.0 { + match self.0.handshake_mut() { MidHandshake::Handshaking(sess) => Some(sess.get_mut().0), MidHandshake::SendAlert { io, .. } => Some(io), MidHandshake::Error { io, .. } => Some(io), @@ -181,8 +198,8 @@ impl Future for Connect { type Output = io::Result>; #[inline] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err) + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.get_mut().0.poll(cx).map_err(|(err, _)| err) } } @@ -190,13 +207,13 @@ impl Future for FallibleConnect { type Output = Result, (io::Error, IO)>; #[inline] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.0).poll(cx) + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.get_mut().0.poll(cx) } } /// Like [Connect], but returns `IO` on failure. -pub struct FallibleConnect(MidHandshake>); +pub struct FallibleConnect(HandshakeFuture>); /// A wrapper around an underlying raw stream which implements the TLS or SSL /// protocol. diff --git a/src/common/handshake.rs b/src/common/handshake.rs index a772c00b..3d203448 100644 --- a/src/common/handshake.rs +++ b/src/common/handshake.rs @@ -33,6 +33,16 @@ pub(crate) enum MidHandshake { }, } +impl MidHandshake { + pub(crate) fn take_io(&mut self) -> Option { + match mem::replace(self, Self::End) { + Self::Handshaking(stream) => Some(stream.into_io()), + Self::SendAlert { io, .. } | Self::Error { io, .. } => Some(io), + Self::End => None, + } + } +} + impl Future for MidHandshake where IS: IoSession + Unpin, diff --git a/src/common/mod.rs b/src/common/mod.rs index ede69d9b..506f7570 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -7,7 +7,9 @@ use rustls::{ConnectionCommon, SideData}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; mod handshake; +mod timeout; pub(crate) use handshake::{IoSession, MidHandshake}; +pub(crate) use timeout::HandshakeFuture; #[derive(Debug)] pub(crate) enum TlsState { diff --git a/src/common/timeout.rs b/src/common/timeout.rs new file mode 100644 index 00000000..60f6944f --- /dev/null +++ b/src/common/timeout.rs @@ -0,0 +1,81 @@ +use std::future::Future; +use std::io; +use std::ops::{Deref, DerefMut}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; + +use rustls::{ConnectionCommon, SideData}; +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::time::{sleep, Sleep}; + +use crate::common::{IoSession, MidHandshake}; + +/// A `MidHandshake` future bundled with an optional deadline. +pub(crate) struct HandshakeFuture { + inner: MidHandshake, + timeout: Option, +} + +impl HandshakeFuture { + pub(crate) fn new(inner: MidHandshake, timeout: Option) -> Self { + Self { + inner, + timeout: timeout.map(|duration| Timeout { + duration, + sleep: None, + }), + } + } + + pub(crate) fn handshake(&self) -> &MidHandshake { + &self.inner + } + + pub(crate) fn handshake_mut(&mut self) -> &mut MidHandshake { + &mut self.inner + } +} + +impl HandshakeFuture +where + IS: IoSession + Unpin, + IS::Io: AsyncRead + AsyncWrite + Unpin, + IS::Session: DerefMut + Deref> + Unpin, + SD: SideData, +{ + pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll> { + if let Poll::Ready(result) = Pin::new(&mut self.inner).poll(cx) { + return Poll::Ready(result); + } + + let timeout = match &mut self.timeout { + Some(timeout) => timeout, + None => return Poll::Pending, + }; + + let sleep = timeout + .sleep + .get_or_insert_with(|| Box::pin(sleep(timeout.duration))); + if sleep.as_mut().poll(cx).is_pending() { + return Poll::Pending; + } + + match self.inner.take_io() { + Some(io) => Poll::Ready(Err(( + io::Error::new(io::ErrorKind::TimedOut, "TLS handshake timed out"), + io, + ))), + // The inner handshake just returned `Pending` above, so it must + // still hold its IO because `take_io()` only returns `None` for the + // `End` state, which `MidHandshake::poll` never leaves behind + // when returning `Pending`. + None => unreachable!("handshake returned Pending but has no IO"), + } + } +} + +struct Timeout { + duration: Duration, + sleep: Option>>, +} diff --git a/tests/test.rs b/tests/test.rs index e64cb05b..5ebbe13b 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -143,6 +143,44 @@ async fn fail() -> io::Result<()> { Ok(()) } +#[tokio::test(start_paused = true)] +async fn handshake_timeout() { + let (_, config) = utils::make_configs(); + let config = + TlsConnector::from(Arc::new(config)).with_handshake_timeout(Some(HANDSHAKE_TIMEOUT)); + let domain = ServerName::try_from(utils::TEST_SERVER_DOMAIN) + .unwrap() + .to_owned(); + // The server end is held but never read from. The connect future will block + // waiting on the ServerHello and with all tasks pending, the start_paused + // runtime advances to the handshake timeout's deadline. + let (client, _server) = tokio::io::duplex(4096); + + let err = config.connect(domain, client).await.unwrap_err(); + assert_eq!(err.kind(), ErrorKind::TimedOut); +} + +/// handshake_timeout_fallible operates equivalently to handshake_timeout, but for FallibleConnect. +#[tokio::test(start_paused = true)] +async fn handshake_timeout_fallible() { + let (_, config) = utils::make_configs(); + let config = + TlsConnector::from(Arc::new(config)).with_handshake_timeout(Some(HANDSHAKE_TIMEOUT)); + let domain = ServerName::try_from(utils::TEST_SERVER_DOMAIN) + .unwrap() + .to_owned(); + let (client, _server) = tokio::io::duplex(4096); + + let (err, _io) = config + .connect(domain, client) + .into_fallible() + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::TimedOut); +} + +const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); + #[tokio::test] async fn test_lazy_config_acceptor() -> io::Result<()> { let (sconfig, cconfig) = utils::make_configs(); From 969aa3ef7b005f005400f99a4e46b57aa4ef8f9e Mon Sep 17 00:00:00 2001 From: Daniel McCarney Date: Fri, 19 Jun 2026 13:28:50 -0400 Subject: [PATCH 2/5] server: add TlsAcceptor/FallibleAccept timeout support Matches the client-side implementation. LazyConfigAcceptor/into_stream() paths not yet supported. --- src/server.rs | 102 +++++++++++++++++++++++++++++++++----------------- tests/test.rs | 27 +++++++++++++ 2 files changed, 94 insertions(+), 35 deletions(-) diff --git a/src/server.rs b/src/server.rs index a43dd8fc..1f530758 100644 --- a/src/server.rs +++ b/src/server.rs @@ -7,26 +7,46 @@ use std::os::windows::io::{AsRawSocket, RawSocket}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use std::time::Duration; use rustls::server::AcceptedAlert; use rustls::{ServerConfig, ServerConnection}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; -use crate::common::{IoSession, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, TlsState}; +use crate::common::{ + HandshakeFuture, IoSession, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, TlsState, +}; /// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method. #[derive(Clone)] pub struct TlsAcceptor { inner: Arc, + handshake_timeout: Option, } impl From> for TlsAcceptor { fn from(inner: Arc) -> Self { - Self { inner } + Self { + inner, + handshake_timeout: None, + } } } impl TlsAcceptor { + /// Set the maximum amount of time to allow for TLS handshakes. + /// + /// `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. + pub fn with_handshake_timeout(mut self, timeout: Option) -> Self { + self.handshake_timeout = timeout; + self + } + #[inline] pub fn accept(&self, stream: IO) -> Accept where @@ -43,22 +63,28 @@ impl TlsAcceptor { let mut session = match ServerConnection::new(self.inner.clone()) { Ok(session) => session, Err(error) => { - return Accept(MidHandshake::Error { - io: stream, - // TODO(eliza): should this really return an `io::Error`? - // Probably not... - error: io::Error::new(io::ErrorKind::Other, error), - }); + return Accept(HandshakeFuture::new( + MidHandshake::Error { + io: stream, + // TODO(eliza): should this really return an `io::Error`? + // Probably not... + error: io::Error::new(io::ErrorKind::Other, error), + }, + self.handshake_timeout, + )); } }; f(&mut session); - Accept(MidHandshake::Handshaking(TlsStream { - session, - io: stream, - state: TlsState::Stream, - need_flush: false, - })) + Accept(HandshakeFuture::new( + MidHandshake::Handshaking(TlsStream { + session, + io: stream, + state: TlsState::Stream, + need_flush: false, + }), + self.handshake_timeout, + )) } /// Get a read-only reference to underlying config @@ -229,29 +255,35 @@ where let mut conn = match self.accepted.into_connection(config) { Ok(conn) => conn, Err((error, alert)) => { - return Accept(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), - }); + return Accept(HandshakeFuture::new( + 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, + )); } }; f(&mut conn); - Accept(MidHandshake::Handshaking(TlsStream { - session: conn, - io: self.io, - state: TlsState::Stream, - need_flush: false, - })) + Accept(HandshakeFuture::new( + MidHandshake::Handshaking(TlsStream { + session: conn, + io: self.io, + state: TlsState::Stream, + need_flush: false, + }), + None, + )) } } /// Future returned from `TlsAcceptor::accept` which will resolve /// once the accept handshake has finished. -pub struct Accept(MidHandshake>); +pub struct Accept(HandshakeFuture>); impl Accept { #[inline] @@ -260,7 +292,7 @@ impl Accept { } pub fn get_ref(&self) -> Option<&IO> { - match &self.0 { + match self.0.handshake() { MidHandshake::Handshaking(sess) => Some(sess.get_ref().0), MidHandshake::SendAlert { io, .. } => Some(io), MidHandshake::Error { io, .. } => Some(io), @@ -269,7 +301,7 @@ impl Accept { } pub fn get_mut(&mut self) -> Option<&mut IO> { - match &mut self.0 { + match self.0.handshake_mut() { MidHandshake::Handshaking(sess) => Some(sess.get_mut().0), MidHandshake::SendAlert { io, .. } => Some(io), MidHandshake::Error { io, .. } => Some(io), @@ -282,20 +314,20 @@ impl Future for Accept { type Output = io::Result>; #[inline] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err) + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.get_mut().0.poll(cx).map_err(|(err, _)| err) } } /// Like [Accept], but returns `IO` on failure. -pub struct FallibleAccept(MidHandshake>); +pub struct FallibleAccept(HandshakeFuture>); impl Future for FallibleAccept { type Output = Result, (io::Error, IO)>; #[inline] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.0).poll(cx) + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.get_mut().0.poll(cx) } } diff --git a/tests/test.rs b/tests/test.rs index 5ebbe13b..7c577d4f 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -179,6 +179,33 @@ async fn handshake_timeout_fallible() { assert_eq!(err.kind(), ErrorKind::TimedOut); } +#[tokio::test(start_paused = true)] +async fn accept_handshake_timeout() { + let (config, _) = utils::make_configs(); + let acceptor = + TlsAcceptor::from(Arc::new(config)).with_handshake_timeout(Some(HANDSHAKE_TIMEOUT)); + // The client end is held but never written to. The accept future 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 err = acceptor.accept(server).await.unwrap_err(); + assert_eq!(err.kind(), ErrorKind::TimedOut); +} + +/// accept_handshake_timeout_fallible operates equivalently to accept_handshake_timeout, but for +/// FallibleAccept. +#[tokio::test(start_paused = true)] +async fn accept_handshake_timeout_fallible() { + let (config, _) = utils::make_configs(); + let acceptor = + TlsAcceptor::from(Arc::new(config)).with_handshake_timeout(Some(HANDSHAKE_TIMEOUT)); + let (_client, server) = tokio::io::duplex(4096); + + let (err, _io) = acceptor.accept(server).into_fallible().await.unwrap_err(); + assert_eq!(err.kind(), ErrorKind::TimedOut); +} + const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); #[tokio::test] From bc4bf5e4272d05c314f4109a8fc6d9d1be6a5800 Mon Sep 17 00:00:00 2001 From: Daniel McCarney Date: Fri, 19 Jun 2026 13:35:57 -0400 Subject: [PATCH 3/5] common/timeout: switch HandshakeFuture to Instant-based deadlines Internally store an absolute `Instant` rather than a relative `Duration`, and `sleep_until()` that deadline. The public `new(Option)` constructor converts eagerly so existing callers see no behavior change. Add a `from_deadline(Option)` constructor for callers that already have an absolute deadline. This is needed so `LazyConfigAcceptor` can establish a deadline in its phase and propagate it through `StartHandshake::into_stream()` to the post-ClientHello Accept, keeping the two phases sharing the same wall-clock deadline. --- src/common/timeout.rs | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/src/common/timeout.rs b/src/common/timeout.rs index 60f6944f..199e07eb 100644 --- a/src/common/timeout.rs +++ b/src/common/timeout.rs @@ -7,7 +7,7 @@ use std::time::Duration; use rustls::{ConnectionCommon, SideData}; use tokio::io::{AsyncRead, AsyncWrite}; -use tokio::time::{sleep, Sleep}; +use tokio::time::{Instant, Sleep, sleep_until}; use crate::common::{IoSession, MidHandshake}; @@ -18,11 +18,23 @@ pub(crate) struct HandshakeFuture { } impl HandshakeFuture { + /// Construct with a relative `Duration` timeout. + /// + /// The deadline is fixed to `Instant::now() + duration` immediately, so the + /// clock starts ticking when this is called, not when the future is first polled. pub(crate) fn new(inner: MidHandshake, timeout: Option) -> Self { + Self::from_deadline(inner, timeout.map(|d| Instant::now() + d)) + } + + /// Construct with an absolute `Instant` deadline. + /// + /// Used when an earlier phase (e.g. `LazyConfigAcceptor`) already established the + /// deadline and the post-ClientHello phase needs to inherit it. + pub(crate) fn from_deadline(inner: MidHandshake, deadline: Option) -> Self { Self { inner, - timeout: timeout.map(|duration| Timeout { - duration, + timeout: deadline.map(|deadline| Timeout { + deadline, sleep: None, }), } @@ -56,7 +68,7 @@ where let sleep = timeout .sleep - .get_or_insert_with(|| Box::pin(sleep(timeout.duration))); + .get_or_insert_with(|| Box::pin(sleep_until(timeout.deadline))); if sleep.as_mut().poll(cx).is_pending() { return Poll::Pending; } @@ -76,6 +88,6 @@ where } struct Timeout { - duration: Duration, + deadline: Instant, sleep: Option>>, } From 3cdbe0d32720c3563b89eb0631a07854d662458a Mon Sep 17 00:00:00 2001 From: Daniel McCarney Date: Fri, 19 Jun 2026 14:16:29 -0400 Subject: [PATCH 4/5] common/timeout: expose Timeout for cross-module reuse --- src/common/mod.rs | 2 +- src/common/timeout.rs | 52 ++++++++++++++++++++++++++++++++++--------- 2 files changed, 42 insertions(+), 12 deletions(-) diff --git a/src/common/mod.rs b/src/common/mod.rs index 506f7570..ce450db1 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -9,7 +9,7 @@ use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; mod handshake; mod timeout; pub(crate) use handshake::{IoSession, MidHandshake}; -pub(crate) use timeout::HandshakeFuture; +pub(crate) use timeout::{HandshakeFuture, Timeout}; #[derive(Debug)] pub(crate) enum TlsState { diff --git a/src/common/timeout.rs b/src/common/timeout.rs index 199e07eb..2da1e47f 100644 --- a/src/common/timeout.rs +++ b/src/common/timeout.rs @@ -23,7 +23,10 @@ impl HandshakeFuture { /// The deadline is fixed to `Instant::now() + duration` immediately, so the /// clock starts ticking when this is called, not when the future is first polled. pub(crate) fn new(inner: MidHandshake, timeout: Option) -> Self { - Self::from_deadline(inner, timeout.map(|d| Instant::now() + d)) + Self { + inner, + timeout: timeout.map(Timeout::from_duration), + } } /// Construct with an absolute `Instant` deadline. @@ -33,10 +36,7 @@ impl HandshakeFuture { pub(crate) fn from_deadline(inner: MidHandshake, deadline: Option) -> Self { Self { inner, - timeout: deadline.map(|deadline| Timeout { - deadline, - sleep: None, - }), + timeout: deadline.map(Timeout::from_deadline), } } @@ -65,11 +65,7 @@ where Some(timeout) => timeout, None => return Poll::Pending, }; - - let sleep = timeout - .sleep - .get_or_insert_with(|| Box::pin(sleep_until(timeout.deadline))); - if sleep.as_mut().poll(cx).is_pending() { + if timeout.poll_deadline(cx).is_pending() { return Poll::Pending; } @@ -87,7 +83,41 @@ where } } -struct Timeout { +/// An absolute deadline paired with a lazily-initialized `Sleep` future. +pub(crate) struct Timeout { deadline: Instant, sleep: Option>>, } + +impl Timeout { + /// Construct from an absolute deadline. + pub(crate) fn from_deadline(deadline: Instant) -> Self { + Self { + deadline, + sleep: None, + } + } + + /// Construct from a relative duration. + /// + /// The deadline is fixed to `Instant::now() + duration` immediately. + pub(crate) fn from_duration(duration: Duration) -> Self { + Self::from_deadline(Instant::now() + duration) + } + + /// Poll the deadline. + /// + /// Returns `Poll::Ready(())` once the deadline has elapsed; `Poll::Pending` + /// otherwise. Lazily initializes the inner `Sleep` on first call. + pub(crate) fn poll_deadline(&mut self, cx: &mut Context<'_>) -> Poll<()> { + let sleep = self + .sleep + .get_or_insert_with(|| Box::pin(sleep_until(self.deadline))); + sleep.as_mut().poll(cx) + } + + /// Return the absolute deadline. + pub(crate) fn deadline(&self) -> Instant { + self.deadline + } +} From 044e824dbff776dffb47e26af28e804fe0d7f6cc Mon Sep 17 00:00:00 2001 From: Daniel McCarney Date: Fri, 19 Jun 2026 14:17:25 -0400 Subject: [PATCH 5/5] server: extend handshake timeout to LazyConfigAcceptor + into_stream --- src/server.rs | 78 +++++++++++++++++++++++++++++++++++++++++++-------- tests/test.rs | 52 ++++++++++++++++++++++++++++++++++ 2 files changed, 119 insertions(+), 11 deletions(-) diff --git a/src/server.rs b/src/server.rs index 1f530758..3f966bdf 100644 --- a/src/server.rs +++ b/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) -> Self { self.handshake_timeout = timeout; self @@ -97,6 +101,7 @@ pub struct LazyConfigAcceptor { acceptor: rustls::server::Acceptor, io: Option, alert: Option<(rustls::Error, AcceptedAlert)>, + timeout: Option, } impl LazyConfigAcceptor @@ -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) -> 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( + timeout: &mut Option, + cx: &mut Context<'_>, +) -> Poll> { + 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 { 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, } impl StartHandshake @@ -237,6 +292,7 @@ where Self { accepted, io: transport, + deadline: None, } } @@ -255,7 +311,7 @@ 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, @@ -263,20 +319,20 @@ where // 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, )) } } diff --git a/tests/test.rs b/tests/test.rs index 7c577d4f..40e550af 100644 --- a/tests/test.rs +++ b/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]