diff --git a/Cargo.toml b/Cargo.toml index 88def1b..9dace85 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ webpki-roots = { version = "0.26", optional = true } headers = "0.4" [dev-dependencies] +rustls = { version = "0.23", default-features = false, features = ["ring"] } tokio = { version = "1.35", features = ["full"] } hyper = { version = "1.0", features = ["client", "http1"] } hyper-util = { version = "0.1", features = ["client", "client-legacy", "http1", "tokio"] } @@ -46,7 +47,7 @@ native-tls-vendored = ["native-tls", "tokio-native-tls?/vendored"] rustls-tls-webpki-roots = ["dep:webpki-roots", "__rustls", "hyper-rustls/webpki-roots"] rustls-tls-native-roots = ["__rustls", "hyper-rustls/rustls-native-certs"] -__tls = [] +__tls = ["tokio/time"] # Enables common rustls code. # Equivalent to rustls-tls-manual-roots but shorter :) diff --git a/src/lib.rs b/src/lib.rs index 4a4cc8d..8bf6d72 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -67,6 +67,7 @@ use std::{ future::Future, pin::Pin, task::{Context, Poll}, + time::Duration, }; use futures_util::future::TryFutureExt; @@ -255,6 +256,7 @@ impl Proxy { pub struct ProxyConnector { proxies: Vec, connector: C, + tls_handshake_timeout: Option, #[cfg(all(not(feature = "__rustls"), feature = "native-tls"))] tls: Option, @@ -293,6 +295,7 @@ impl ProxyConnector { Ok(ProxyConnector { proxies: Vec::new(), connector: connector, + tls_handshake_timeout: None, tls: Some(tls), }) } @@ -314,6 +317,7 @@ impl ProxyConnector { Ok(ProxyConnector { proxies: Vec::new(), connector, + tls_handshake_timeout: None, tls: Some(tls), }) } @@ -323,6 +327,7 @@ impl ProxyConnector { ProxyConnector { proxies: Vec::new(), connector, + tls_handshake_timeout: None, tls: None, } } @@ -347,6 +352,7 @@ impl ProxyConnector { ProxyConnector { connector, proxies: self.proxies, + tls_handshake_timeout: self.tls_handshake_timeout, tls: self.tls, } } @@ -363,6 +369,14 @@ impl ProxyConnector { self.tls = tls; } + /// Set or unset the timeout bounding the TLS handshake performed when tunneling through + /// this proxy to an HTTPS destination. + /// + /// `None` (the default) means the handshake is not bounded by a timeout. + pub fn set_tls_handshake_timeout(&mut self, timeout: Option) { + self.tls_handshake_timeout = timeout; + } + /// Get the current proxies pub fn proxies(&self) -> &[Proxy] { &self.proxies @@ -443,6 +457,8 @@ where } else { None }; + #[cfg(feature = "__tls")] + let tls_handshake_timeout = self.tls_handshake_timeout; Box::pin(async move { // this hack will gone once `try_blocks` will eventually stabilized @@ -456,10 +472,14 @@ where Some(tls) => { use hyper_util::rt::TokioIo; let tls = TlsConnector::from(tls); - let secure_stream = mtry!(tls - .connect(&host, TokioIo::new(tunnel_stream)) + let secure_stream = mtry!(mtry!( + with_optional_timeout( + tls_handshake_timeout, + tls.connect(&host, TokioIo::new(tunnel_stream)) + ) .await - .map_err(io_err)); + ) + .map_err(io_err)); Ok(ProxyStream::Secured(Box::new(TokioIo::new(secure_stream)))) } @@ -469,10 +489,14 @@ where use hyper_util::rt::TokioIo; let server_name = mtry!(ServerName::try_from(host.to_string()).map_err(io_err)); - let secure_stream = mtry!(tls - .connect(server_name, TokioIo::new(tunnel_stream)) + let secure_stream = mtry!(mtry!( + with_optional_timeout( + tls_handshake_timeout, + tls.connect(server_name, TokioIo::new(tunnel_stream)) + ) .await - .map_err(io_err)); + ) + .map_err(io_err)); Ok(ProxyStream::Secured(Box::new(TokioIo::new(secure_stream)))) } @@ -506,6 +530,22 @@ where } } +/// Awaits `fut`, bounding it by `timeout` when set. +/// +/// A `None` timeout awaits `fut` directly, so it never yields the outer timeout error. +#[cfg(feature = "__tls")] +async fn with_optional_timeout( + timeout: Option, + fut: F, +) -> io::Result { + match timeout { + Some(timeout) => tokio::time::timeout(timeout, fut) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "TLS handshake timed out")), + None => Ok(fut.await), + } +} + fn proxy_dst(dst: &Uri, proxy: &Uri) -> io::Result { Uri::builder() .scheme( @@ -562,4 +602,58 @@ mod tests { assert_eq!(proxy.headers().get("authorization"), None); } + + #[cfg(feature = "__tls")] + #[tokio::test] + async fn tls_handshake_timeout_fires_against_a_stalled_proxy_tunnel() { + use std::time::Duration; + + use hyper_util::client::legacy::connect::HttpConnector; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + use tokio::time::timeout; + use tower_service::Service as _; + + use crate::{Proxy, ProxyConnector}; + + #[cfg(feature = "__rustls")] + let _ = tokio_rustls::rustls::crypto::ring::default_provider().install_default(); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + // Emulates a proxy that completes the `CONNECT` tunnel but then never speaks TLS, so the + // client's handshake deadline is the only thing that can end the connection attempt. + let proxy_task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf).await.unwrap(); + stream.write_all(b"HTTP/1.1 200 OK\r\n\r\n").await.unwrap(); + tokio::time::sleep(Duration::from_secs(30)).await; + drop(stream); + }); + + let proxy_uri: Uri = format!("http://{addr}").parse().unwrap(); + let proxy = Proxy::new(Intercept::All, proxy_uri); + + let mut proxy_connector = ProxyConnector::new(HttpConnector::new()).unwrap(); + proxy_connector.add_proxy(proxy); + proxy_connector.set_tls_handshake_timeout(Some(Duration::from_millis(200))); + + let dst: Uri = "https://example.invalid/".parse().unwrap(); + let result = timeout(Duration::from_secs(5), proxy_connector.call(dst)) + .await + .expect("request should not hit the outer test timeout"); + let error = match result { + Ok(_) => panic!("handshake should time out before completing"), + Err(error) => error, + }; + + assert!( + error.to_string().contains("TLS handshake timed out"), + "expected a TLS handshake timeout, got: {error}" + ); + + proxy_task.abort(); + } }