diff --git a/examples/client.rs b/examples/client.rs index 601fc668..4c2c0ca4 100644 --- a/examples/client.rs +++ b/examples/client.rs @@ -1,35 +1,32 @@ +//! An example using tokio-rustls to make an HTTP GET request, demonstrating +//! how to apply timeouts to every phase of the connection lifecycle using +//! only `tokio::time` and the public tokio-rustls API: +//! +//! * TCP connect +//! * TLS handshake +//! * individual reads/writes on the established stream +//! * graceful TLS shutdown (close notify) +//! +//! The handshake timeout also applies unchanged to `TlsConnector::connect_with` +//! and `TlsConnectorWithAlpn::connect`, they all return the same `Connect` +//! future. + use std::error::Error as StdError; +use std::future::Future; use std::io; use std::net::ToSocketAddrs; use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; use argh::FromArgs; -use rustls::pki_types::pem::PemObject; -use rustls::pki_types::{CertificateDer, ServerName}; -use tokio::io::{AsyncWriteExt, copy, split, stdin as tokio_stdin, stdout as tokio_stdout}; +use tokio::io::{AsyncReadExt, AsyncWriteExt, stdout as tokio_stdout}; use tokio::net::TcpStream; -use tokio_rustls::{TlsConnector, rustls}; - -/// Tokio Rustls client example -#[derive(FromArgs)] -struct Options { - /// host - #[argh(positional)] - host: String, - - /// port - #[argh(option, short = 'p', default = "443")] - port: u16, - - /// domain - #[argh(option, short = 'd')] - domain: Option, - - /// cafile - #[argh(option, short = 'c')] - cafile: Option, -} +use tokio::time::timeout; +use tokio_rustls::TlsConnector; +use tokio_rustls::rustls::pki_types::pem::PemObject; +use tokio_rustls::rustls::pki_types::{CertificateDer, ServerName}; +use tokio_rustls::rustls::{ClientConfig, RootCertStore}; #[tokio::main] async fn main() -> Result<(), Box> { @@ -40,9 +37,9 @@ async fn main() -> Result<(), Box> { .next() .ok_or_else(|| io::Error::from(io::ErrorKind::NotFound))?; let domain = options.domain.unwrap_or(options.host); - let content = format!("GET / HTTP/1.0\r\nHost: {}\r\n\r\n", domain); + let content = format!("GET / HTTP/1.0\r\nHost: {domain}\r\n\r\n",); - let mut root_cert_store = rustls::RootCertStore::empty(); + let mut root_cert_store = RootCertStore::empty(); if let Some(cafile) = &options.cafile { for cert in CertificateDer::pem_file_iter(cafile)? { root_cert_store.add(cert?)?; @@ -51,30 +48,118 @@ async fn main() -> Result<(), Box> { root_cert_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); } - let config = rustls::ClientConfig::builder() - .with_root_certificates(root_cert_store) - .with_no_client_auth(); // i guess this was previously the default? - let connector = TlsConnector::from(Arc::new(config)); + let connector = TlsConnector::from(Arc::new( + ClientConfig::builder() + .with_root_certificates(root_cert_store) + .with_no_client_auth(), + )); - let stream = TcpStream::connect(&addr).await?; - - let (mut stdin, mut stdout) = (tokio_stdin(), tokio_stdout()); + // TCP connect, bounded by the connect timeout. + let stream = with_timeout( + TcpStream::connect(&addr), + "TCP connect", + Duration::from_secs(options.connect_timeout), + ) + .await?; + // TLS handshake, bounded by the handshake timeout. + // + // On timeout the `Connect` future, and the socket it owns, is dropped. let domain = ServerName::try_from(domain.as_str())?.to_owned(); - let mut stream = connector.connect(domain, stream).await?; - stream.write_all(content.as_bytes()).await?; - - let (mut reader, mut writer) = split(stream); - - tokio::select! { - ret = copy(&mut reader, &mut stdout) => { - ret?; - }, - ret = copy(&mut stdin, &mut writer) => { - ret?; - writer.shutdown().await? + let mut stream = with_timeout( + connector.connect(domain, stream), + "TLS handshake", + Duration::from_secs(options.handshake_timeout), + ) + .await?; + + let io_timeout = Duration::from_secs(options.io_timeout); + + // request/response, with a timeout on each individual write and + // read. + // + // Note a single `timeout()` around `read_to_end()` would bound the + // *total* response time instead. Wrapping each `read()` bounds idle time. + with_timeout( + stream.write_all(content.as_bytes()), + "request write", + io_timeout, + ) + .await?; + + let mut stdout = tokio_stdout(); + let mut buf = vec![0u8; 8192]; + loop { + let n = with_timeout(stream.read(&mut buf), "response read", io_timeout).await?; + if n == 0 { + break; } + stdout.write_all(&buf[..n]).await?; } + stdout.flush().await?; + + // graceful TLS shutdown. + // + // `shutdown()` writes a close_notify alert, so it too can block on a stalled peer. + with_timeout( + stream.shutdown(), + "TLS shutdown", + Duration::from_secs(options.shutdown_timeout), + ) + .await?; Ok(()) } + +/// Tokio Rustls client example with timeouts +#[derive(FromArgs)] +struct Options { + /// host + #[argh(positional)] + host: String, + + /// port + #[argh(option, short = 'p', default = "443")] + port: u16, + + /// domain + #[argh(option, short = 'd')] + domain: Option, + + /// cafile + #[argh(option, short = 'c')] + cafile: Option, + + /// TCP connect timeout (seconds) + #[argh(option, default = "10")] + connect_timeout: u64, + + /// TLS handshake timeout (seconds) + #[argh(option, default = "10")] + handshake_timeout: u64, + + /// per-read/write timeout on the established stream (seconds) + #[argh(option, default = "30")] + io_timeout: u64, + + /// graceful TLS shutdown (close_notify) timeout (seconds) + #[argh(option, default = "5")] + shutdown_timeout: u64, +} + +/// Await `fut` for at most `duration`, converting an elapsed timeout into an `io::Error`. +/// +/// On timeout, the returned error describes the `phase` that timed out. +async fn with_timeout( + fut: impl Future>, + phase: &str, + duration: Duration, +) -> io::Result { + match timeout(duration, fut).await { + Ok(result) => result, + Err(_) => Err(io::Error::new( + io::ErrorKind::TimedOut, + format!("{phase} timed out after {duration:?}"), + )), + } +} diff --git a/examples/server.rs b/examples/server.rs index e4b1cb35..5997715c 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -1,35 +1,31 @@ +//! An example using tokio-rustls to build an echo server, demonstrating how +//! to apply timeouts to every phase of each accepted connection. +//! +//! * TLS handshake via `TlsAcceptor::accept` (the default path) +//! * TLS handshake via `LazyConfigAcceptor` (`--lazy`), with one deadline +//! spanning both the wait for the ClientHello and the rest of the +//! handshake, and `take_io()` used to answer a stalled client in plaintext +//! * individual reads/writes on the established stream +//! * graceful TLS shutdown (close notify) + use std::error::Error as StdError; +use std::future::Future; use std::io; use std::net::ToSocketAddrs; use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; use argh::FromArgs; -use rustls::pki_types::pem::PemObject; -use rustls::pki_types::{CertificateDer, PrivateKeyDer}; -use tokio::io::{AsyncWriteExt, copy, sink, split}; -use tokio::net::TcpListener; -use tokio_rustls::{TlsAcceptor, rustls}; - -/// Tokio Rustls server example -#[derive(FromArgs)] -struct Options { - /// bind addr - #[argh(positional)] - addr: String, - - /// cert file - #[argh(option, short = 'c')] - cert: PathBuf, - - /// key file - #[argh(option, short = 'k')] - key: PathBuf, - - /// echo mode - #[argh(switch, short = 'e')] - echo_mode: bool, -} +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::time::{Instant, timeout, timeout_at}; +use tokio_rustls::rustls::ServerConfig; +use tokio_rustls::rustls::pki_types::pem::PemObject; +use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer}; +use tokio_rustls::rustls::server::Acceptor; +use tokio_rustls::server::TlsStream; +use tokio_rustls::{LazyConfigAcceptor, TlsAcceptor}; #[tokio::main] async fn main() -> Result<(), Box> { @@ -40,52 +36,188 @@ async fn main() -> Result<(), Box> { .to_socket_addrs()? .next() .ok_or_else(|| io::Error::from(io::ErrorKind::AddrNotAvailable))?; - let certs = CertificateDer::pem_file_iter(&options.cert)?.collect::, _>>()?; - let key = PrivateKeyDer::from_pem_file(&options.key)?; - let flag_echo = options.echo_mode; - let config = rustls::ServerConfig::builder() - .with_no_client_auth() - .with_single_cert(certs, key)?; - let acceptor = TlsAcceptor::from(Arc::new(config)); + let handshake_timeout = Duration::from_secs(options.handshake_timeout); + let io_timeout = Duration::from_secs(options.io_timeout); + let shutdown_timeout = Duration::from_secs(options.shutdown_timeout); + let config = Arc::new( + ServerConfig::builder() + .with_no_client_auth() + .with_single_cert( + CertificateDer::pem_file_iter(&options.cert)?.collect::>()?, + PrivateKeyDer::from_pem_file(&options.key)?, + )?, + ); + let acceptor = TlsAcceptor::from(config.clone()); let listener = TcpListener::bind(&addr).await?; loop { - let (stream, peer_addr) = listener.accept().await?; let acceptor = acceptor.clone(); + let config = config.clone(); + let (stream, peer_addr) = listener.accept().await?; let fut = async move { - let mut stream = acceptor.accept(stream).await?; - - if flag_echo { - let (mut reader, mut writer) = split(stream); - let n = copy(&mut reader, &mut writer).await?; - writer.flush().await?; - println!("Echo: {} - {}", peer_addr, n); - } else { - let mut output = sink(); - stream - .write_all( - &b"HTTP/1.0 200 ok\r\n\ - Connection: close\r\n\ - Content-length: 12\r\n\ - \r\n\ - Hello world!"[..], - ) - .await?; - stream.shutdown().await?; - copy(&mut stream, &mut output).await?; - println!("Hello: {}", peer_addr); - } + let stream = match options.lazy { + // Accept the connection lazily, applying the handshake_timeout. + true => accept_lazy(config, stream, handshake_timeout).await?, + // Or, directly accept the stream, applying the handshake_timeout. + false => { + with_timeout(acceptor.accept(stream), "TLS handshake", handshake_timeout) + .await? + } + }; + + let n = serve(stream, io_timeout, shutdown_timeout).await?; + println!("Echo: {} - {}", peer_addr, n); Ok(()) as io::Result<()> }; tokio::spawn(async move { if let Err(err) = fut.await { - eprintln!("{:?}", err); + eprintln!("{peer_addr}: {err:?}"); } }); } } + +/// Tokio Rustls server example with timeouts +#[derive(FromArgs)] +struct Options { + /// bind addr + #[argh(positional)] + addr: String, + + /// cert file + #[argh(option, short = 'c')] + cert: PathBuf, + + /// key file + #[argh(option, short = 'k')] + key: PathBuf, + + /// use LazyConfigAcceptor instead of TlsAcceptor + #[argh(switch, short = 'l')] + lazy: bool, + + /// TLS handshake timeout (seconds) + #[argh(option, default = "5")] + handshake_timeout: u64, + + /// per-read/write timeout on the established stream (seconds) + #[argh(option, default = "30")] + io_timeout: u64, + + /// graceful TLS shutdown (close_notify) timeout (seconds) + #[argh(option, default = "5")] + shutdown_timeout: u64, +} + +/// Accept a TLS connection with `LazyConfigAcceptor` under a single deadline. +/// +/// The handshake has two await points: +/// +/// 1. reading the ClientHello +/// 2. driving the rest of the handshake with the chosen config +/// +/// A shared `Instant` deadline (via `timeout_at`) keeps the total bounded +/// rather than granting each phase its own budget. +/// +/// If the client stalls before even sending a ClientHello we can take the +/// socket back with `take_io()` and respond in plaintext before hanging up. +async fn accept_lazy( + config: Arc, + stream: TcpStream, + handshake_timeout: Duration, +) -> io::Result> { + let deadline = Instant::now() + handshake_timeout; + + let acceptor = LazyConfigAcceptor::new(Acceptor::default(), stream); + tokio::pin!(acceptor); + + // Read the ClientHello, respecting the deadline. + let start = match timeout_at(deadline, acceptor.as_mut()).await { + Ok(result) => result?, + Err(_) => { + // No ClientHello arrived in time. No TLS handshake has happened yet, + // so we can still answer in plaintext. + if let Some(mut stream) = acceptor.take_io() { + let _ = stream + .write_all(b"HTTP/1.0 408 Request Timeout\r\n\r\n") + .await; + } + return Err(io::Error::new( + io::ErrorKind::TimedOut, + format!("timed out awaiting ClientHello after {handshake_timeout:?}"), + )); + } + }; + + // The ClientHello is now available, e.g. for SNI-based config selection. + if let Some(sni) = start.client_hello().server_name() { + println!("ClientHello with SNI: {sni}"); + } + + // Complete the handshake, respecting the deadline. + match timeout_at(deadline, start.into_stream(config)).await { + Ok(result) => result, + Err(_elapsed) => Err(io::Error::new( + io::ErrorKind::TimedOut, + format!("TLS handshake timed out after {handshake_timeout:?}"), + )), + } +} + +/// Serve a single established TLS `stream`, respecting timeouts. +/// +/// Echoes one request worth of bytes, bounding each read/write with the `io_timeout`, and the +/// final graceful shutdown with the `shutdown_timeout.` +async fn serve( + mut stream: TlsStream, + io_timeout: Duration, + shutdown_timeout: Duration, +) -> io::Result { + let mut echoed = 0; + let mut buf = vec![0u8; 8192]; + loop { + // Bounding each read() limits how long the peer may sit idle. A + // single timeout() around the whole loop would bound total + // connection time instead. + // For an idle peer, stop echoing and fall through to the graceful + // shutdown below so the peer sees close_notify rather than an abrupt + // close. + let n = match timeout(io_timeout, stream.read(&mut buf)).await { + Ok(result) => result?, + Err(_elapsed) => break, + }; + if n == 0 { + break; + } + with_timeout(stream.write_all(&buf[..n]), "write", io_timeout).await?; + echoed += n as u64; + } + with_timeout(stream.flush(), "flush", io_timeout).await?; + + // shutdown() writes a close_notify alert, so it too can block on a + // stalled peer. + with_timeout(stream.shutdown(), "TLS shutdown", shutdown_timeout).await?; + Ok(echoed) +} + +/// Await `fut` for at most `duration`, converting an elapsed timeout into an `io::Error`. +/// +/// On timeout, the returned error describes the `phase` that timed out. +async fn with_timeout( + fut: impl Future>, + phase: &str, + duration: Duration, +) -> io::Result { + match timeout(duration, fut).await { + Ok(result) => result, + Err(_) => Err(io::Error::new( + io::ErrorKind::TimedOut, + format!("{phase} timed out after {duration:?}"), + )), + } +} diff --git a/src/client.rs b/src/client.rs index a94b64a8..bb01cf94 100644 --- a/src/client.rs +++ b/src/client.rs @@ -35,6 +35,12 @@ impl TlsConnector { self } + /// Returns a future that performs a TLS handshake with `domain` using the `stream`. + /// + /// You likely want to wrap this in a timeout (for example with [`tokio::time::timeout`][]) + /// to bound the handshake time. + /// + /// [`tokio::time::timeout`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout.html #[inline] pub fn connect(&self, domain: ServerName<'static>, stream: IO) -> Connect where @@ -43,6 +49,15 @@ impl TlsConnector { self.connect_impl(domain, stream, None, |_| ()) } + /// Similar to [`Self::connect()`], but calls `f` before performing the handshake. + /// + /// As with [`Self::connect()`] you likely want to wrap this in a timeout to + /// bound the handshake time. + /// + /// The `f` handler is given a mutable reference to a [`ClientConnection`][] that can + /// be used for tasks like writing early data. + /// + /// [`ClientConnection`]: https://docs.rs/rustls/latest/rustls/client/struct.ClientConnection.html #[inline] pub fn connect_with(&self, domain: ServerName<'static>, stream: IO, f: F) -> Connect where @@ -128,6 +143,12 @@ pub struct TlsConnectorWithAlpn<'c> { } impl TlsConnectorWithAlpn<'_> { + /// Returns a future that performs a TLS handshake with `domain` using the `stream`. + /// + /// You likely want to wrap this in a timeout (for example with [`tokio::time::timeout`][]) + /// to bound the handshake time. + /// + /// [`tokio::time::timeout`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout.html #[inline] pub fn connect(self, domain: ServerName<'static>, stream: IO) -> Connect where @@ -137,6 +158,15 @@ impl TlsConnectorWithAlpn<'_> { .connect_impl(domain, stream, Some(self.alpn_protocols), |_| ()) } + /// Similar to [`Self::connect()`], but calls `f` before performing the handshake. + /// + /// As with [`Self::connect()`] you likely want to wrap this in a timeout to + /// bound the handshake time. + /// + /// The `f` handler is given a mutable reference to a [`ClientConnection`][] that can + /// be used for tasks like writing early data. + /// + /// [`ClientConnection`]: https://docs.rs/rustls/latest/rustls/client/struct.ClientConnection.html #[inline] pub fn connect_with(self, domain: ServerName<'static>, stream: IO, f: F) -> Connect where diff --git a/src/server.rs b/src/server.rs index a43dd8fc..2b068984 100644 --- a/src/server.rs +++ b/src/server.rs @@ -27,6 +27,12 @@ impl From> for TlsAcceptor { } impl TlsAcceptor { + /// Returns a future for completing a TLS handshake for a client using `stream`. + /// + /// You likely want to wrap this in a timeout (for example with [`tokio::time::timeout`][]) + /// to bound the handshake time. + /// + /// [`tokio::time::timeout`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout.html #[inline] pub fn accept(&self, stream: IO) -> Accept where @@ -35,6 +41,18 @@ impl TlsAcceptor { self.accept_with(stream, |_| ()) } + /// Similar to [`Self::accept()`], but calls `f` before performing the handshake. + /// + /// As with [`Self::accept()`] you likely want to wrap this in a timeout to + /// bound the handshake time. + /// + /// The `f` handler is given a mutable reference to a [`ServerConnection`][] that can be used + /// to configure the connection before the handshake, for example, adjusting the buffer limit. + /// + /// Because no data has been read from `stream` yet when `f` is called ClientHello + /// dependent state (like early data) is not yet available. + /// + /// [`ServerConnection`]: https://docs.rs/rustls/latest/rustls/server/struct.ServerConnection.html pub fn accept_with(&self, stream: IO, f: F) -> Accept where IO: AsyncRead + AsyncWrite + Unpin, @@ -67,6 +85,14 @@ impl TlsAcceptor { } } +/// A future for reading a `ClientHello` from `io` without committing to a [`ServerConfig`][]. +/// +/// Awaiting it yields a [`StartHandshake`], which exposes the +/// [`ClientHello`][] (for example, to choose a config based on SNI) and performs +/// the rest of the handshake via [`StartHandshake::into_stream()`]. +/// +/// [`ServerConfig`]: https://docs.rs/rustls/latest/rustls/server/struct.ServerConfig.html +/// [`ClientHello`]: https://docs.rs/rustls/latest/rustls/server/struct.ClientHello.html pub struct LazyConfigAcceptor { acceptor: rustls::server::Acceptor, io: Option, @@ -77,6 +103,22 @@ impl LazyConfigAcceptor where IO: AsyncRead + AsyncWrite + Unpin, { + /// Returns a new `LazyConfigAcceptor` that reads a `ClientHello` from `io`. + /// + /// You likely want to wrap awaiting the acceptor in a timeout to bound how long the + /// peer may take to send the `ClientHello`. + /// + /// Note that awaiting the acceptor is only the first half of the handshake and + /// [`StartHandshake::into_stream()`] performs the rest. + /// + /// To bound the time for the complete handshake, share one deadline across + /// both awaits (for example with [`tokio::time::timeout_at`][]) rather than giving each + /// its own timeout. + /// + /// If a timeout elapses before the `ClientHello` arrives, [`Self::take_io()`] can + /// recover the `io`, for example to answer the peer in plaintext before closing. + /// + /// [`tokio::time::timeout_at`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout_at.html #[inline] pub fn new(acceptor: rustls::server::Acceptor, io: IO) -> Self { Self { @@ -218,10 +260,27 @@ where self.accepted.client_hello() } + /// Returns a future that performs the rest of the TLS handshake using `config`. + /// + /// You likely want to wrap this in a timeout to bound the handshake time. Ideally + /// with [`tokio::time::timeout_at`][], reusing the deadline that also bounded + /// awaiting the [`LazyConfigAcceptor`] so both halves of the handshake share + /// one budget. See [`LazyConfigAcceptor::new()`]. + /// + /// [`tokio::time::timeout_at`]: https://docs.rs/tokio/latest/tokio/time/fn.timeout_at.html pub fn into_stream(self, config: Arc) -> Accept { self.into_stream_with(config, |_| ()) } + /// Similar to [`Self::into_stream()`], but calls `f` before performing the handshake. + /// + /// As with [`Self::into_stream()`] you likely want to wrap this in a timeout to + /// bound the handshake time. + /// + /// The `f` handler is given a mutable reference to a [`ServerConnection`][] that can be + /// used to configure the connection before the handshake. + /// + /// [`ServerConnection`]: https://docs.rs/rustls/latest/rustls/server/struct.ServerConnection.html pub fn into_stream_with(self, config: Arc, f: F) -> Accept where F: FnOnce(&mut ServerConnection),