diff --git a/Cargo.toml b/Cargo.toml index 60dde864..075b40c2 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"] diff --git a/src/client.rs b/src/client.rs index 128e840e..ee806aa0 100644 --- a/src/client.rs +++ b/src/client.rs @@ -7,6 +7,7 @@ use std::pin::Pin; #[cfg(feature = "early-data")] use std::task::Waker; use std::task::{Context, Poll}; +use std::time::Duration; use std::{ io::{self, BufRead as _}, sync::Arc, @@ -14,6 +15,7 @@ use std::{ use rustls::{pki_types::ServerName, ClientConfig, ClientConnection}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; +use tokio::time::{self, Sleep}; use crate::common::{IoSession, MidHandshake, Stream, TlsState}; @@ -21,11 +23,20 @@ use crate::common::{IoSession, MidHandshake, Stream, TlsState}; #[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, @@ -68,36 +79,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 { + inner: MidHandshake::Error { + io: stream, + // TODO(eliza): should this really return an `io::Error`? + // Probably not... + error: io::Error::new(io::ErrorKind::Other, error), + }, + timeout: self.handshake_timeout.map(HandshakeTimeout::new), + }; } }; f(&mut session); - Connect(MidHandshake::Handshaking(TlsStream { - io: stream, + Connect { + inner: 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, + }), + timeout: self.handshake_timeout.map(HandshakeTimeout::new), + } } pub fn with_alpn(&self, alpn_protocols: Vec>) -> TlsConnectorWithAlpn<'_> { @@ -117,6 +134,7 @@ impl From> for TlsConnector { fn from(inner: Arc) -> Self { Self { inner, + handshake_timeout: None, #[cfg(feature = "early-data")] early_data: false, } @@ -151,16 +169,43 @@ impl TlsConnectorWithAlpn<'_> { /// Future returned from `TlsConnector::connect` which will resolve /// once the connection handshake has finished. -pub struct Connect(MidHandshake>); +pub struct Connect { + inner: MidHandshake>, + timeout: Option, +} + +struct HandshakeTimeout { + duration: Duration, + sleep: Option>>, +} + +impl HandshakeTimeout { + fn new(duration: Duration) -> Self { + Self { + duration, + sleep: None, + } + } + + fn poll(&mut self, cx: &mut Context<'_>) -> Poll<()> { + let sleep = self + .sleep + .get_or_insert_with(|| Box::pin(time::sleep(self.duration))); + sleep.as_mut().poll(cx) + } +} impl Connect { #[inline] pub fn into_fallible(self) -> FallibleConnect { - FallibleConnect(self.0) + FallibleConnect { + inner: self.inner, + timeout: self.timeout, + } } pub fn get_ref(&self) -> Option<&IO> { - match &self.0 { + match &self.inner { MidHandshake::Handshaking(sess) => Some(sess.get_ref().0), MidHandshake::SendAlert { io, .. } => Some(io), MidHandshake::Error { io, .. } => Some(io), @@ -169,7 +214,7 @@ impl Connect { } pub fn get_mut(&mut self) -> Option<&mut IO> { - match &mut self.0 { + match &mut self.inner { MidHandshake::Handshaking(sess) => Some(sess.get_mut().0), MidHandshake::SendAlert { io, .. } => Some(io), MidHandshake::Error { io, .. } => Some(io), @@ -183,7 +228,8 @@ impl Future for Connect { #[inline] fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err) + let this = self.as_mut().get_mut(); + poll_fallible_connect(&mut this.inner, &mut this.timeout, cx).map_err(|(err, _)| err) } } @@ -192,12 +238,44 @@ impl Future for FallibleConnect { #[inline] fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.0).poll(cx) + let this = self.as_mut().get_mut(); + poll_fallible_connect(&mut this.inner, &mut this.timeout, cx) } } /// Like [Connect], but returns `IO` on failure. -pub struct FallibleConnect(MidHandshake>); +pub struct FallibleConnect { + inner: MidHandshake>, + timeout: Option, +} + +fn handshake_timeout_error() -> io::Error { + io::Error::new(io::ErrorKind::TimedOut, "TLS handshake timed out") +} + +fn poll_fallible_connect( + inner: &mut MidHandshake>, + timeout: &mut Option, + cx: &mut Context<'_>, +) -> Poll, (io::Error, IO)>> +where + IO: AsyncRead + AsyncWrite + Unpin, +{ + match Pin::new(&mut *inner).poll(cx) { + Poll::Ready(result) => Poll::Ready(result), + Poll::Pending => match timeout { + Some(timeout) => { + if timeout.poll(cx).is_pending() { + return Poll::Pending; + } + + let io = inner.take_io().expect("handshake missing IO"); + Poll::Ready(Err((handshake_timeout_error(), io))) + } + _ => Poll::Pending, + }, + } +} /// 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/tests/test.rs b/tests/test.rs index 8147e92d..28850d39 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -143,6 +143,20 @@ async fn fail() -> io::Result<()> { Ok(()) } +#[tokio::test] +async fn handshake_timeout() { + let (_, config) = utils::make_configs(); + let config = TlsConnector::from(Arc::new(config)) + .with_handshake_timeout(Some(Duration::from_millis(10))); + let domain = ServerName::try_from(utils::TEST_SERVER_DOMAIN) + .unwrap() + .to_owned(); + let (client, _server) = tokio::io::duplex(4096); + + let err = config.connect(domain, client).await.unwrap_err(); + assert_eq!(err.kind(), ErrorKind::TimedOut); +} + #[tokio::test] async fn test_lazy_config_acceptor() -> io::Result<()> { let (sconfig, cconfig) = utils::make_configs();