Client and server timeout support

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]

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<ClientConfig>,
+    handshake_timeout: Option<Duration>,
     #[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<Duration>) -> 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<Vec<u8>>) -> TlsConnectorWithAlpn<'_> {
@@ -116,6 +132,7 @@ impl From<Arc<ClientConfig>> for TlsConnector {
     fn from(inner: Arc<ClientConfig>) -> 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<IO>(MidHandshake<TlsStream<IO>>);
+pub struct Connect<IO>(HandshakeFuture<TlsStream<IO>>);
 
 impl<IO> Connect<IO> {
     #[inline]
@@ -159,7 +176,7 @@ impl<IO> Connect<IO> {
     }
 
     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<IO> Connect<IO> {
     }
 
     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,22 +198,22 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Connect<IO> {
     type Output = io::Result<TlsStream<IO>>;
 
     #[inline]
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err)
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.get_mut().0.poll(cx).map_err(|(err, _)| err)
     }
 }
 
 impl<IO: AsyncRead + AsyncWrite + Unpin> Future for FallibleConnect<IO> {
     type Output = Result<TlsStream<IO>, (io::Error, IO)>;
 
     #[inline]
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        Pin::new(&mut self.0).poll(cx)
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.get_mut().0.poll(cx)
     }
 }
 
 /// Like [Connect], but returns `IO` on failure.
-pub struct FallibleConnect<IO>(MidHandshake<TlsStream<IO>>);
+pub struct FallibleConnect<IO>(HandshakeFuture<TlsStream<IO>>);
 
 /// A wrapper around an underlying raw stream which implements the TLS or SSL
 /// protocol.

src/common/handshake.rs

@@ -33,6 +33,16 @@ pub(crate) enum MidHandshake<IS: IoSession> {
     },
 }
 
+impl<IS: IoSession> MidHandshake<IS> {
+    pub(crate) fn take_io(&mut self) -> Option<IS::Io> {
+        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<IS, SD> Future for MidHandshake<IS>
 where
     IS: IoSession + Unpin,

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, Timeout};
 
 #[derive(Debug)]
 pub(crate) enum TlsState {

src/common/timeout.rs

@@ -0,0 +1,123 @@
+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::{Instant, Sleep, sleep_until};
+
+use crate::common::{IoSession, MidHandshake};
+
+/// A `MidHandshake` future bundled with an optional deadline.
+pub(crate) struct HandshakeFuture<IS: IoSession> {
+    inner: MidHandshake<IS>,
+    timeout: Option<Timeout>,
+}
+
+impl<IS: IoSession> HandshakeFuture<IS> {
+    /// 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<IS>, timeout: Option<Duration>) -> Self {
+        Self {
+            inner,
+            timeout: timeout.map(Timeout::from_duration),
+        }
+    }
+
+    /// 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<IS>, deadline: Option<Instant>) -> Self {
+        Self {
+            inner,
+            timeout: deadline.map(Timeout::from_deadline),
+        }
+    }
+
+    pub(crate) fn handshake(&self) -> &MidHandshake<IS> {
+        &self.inner
+    }
+
+    pub(crate) fn handshake_mut(&mut self) -> &mut MidHandshake<IS> {
+        &mut self.inner
+    }
+}
+
+impl<IS, SD> HandshakeFuture<IS>
+where
+    IS: IoSession + Unpin,
+    IS::Io: AsyncRead + AsyncWrite + Unpin,
+    IS::Session: DerefMut + Deref<Target = ConnectionCommon<SD>> + Unpin,
+    SD: SideData,
+{
+    pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<IS, (io::Error, IS::Io)>> {
+        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,
+        };
+        if timeout.poll_deadline(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"),
+        }
+    }
+}
+
+/// An absolute deadline paired with a lazily-initialized `Sleep` future.
+pub(crate) struct Timeout {
+    deadline: Instant,
+    sleep: Option<Pin<Box<Sleep>>>,
+}
+
+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
+    }
+}