client/server timeout examples, doc pass hinting towards timeouts.

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<String>,
-
-    /// cafile
-    #[argh(option, short = 'c')]
-    cafile: Option<PathBuf>,
-}
+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<dyn StdError + Send + Sync + 'static>> {
@@ -40,9 +37,9 @@ async fn main() -> Result<(), Box<dyn StdError + Send + Sync + 'static>> {
         .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<dyn StdError + Send + Sync + 'static>> {
         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<String>,
+
+    /// cafile
+    #[argh(option, short = 'c')]
+    cafile: Option<PathBuf>,
+
+    /// 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<T>(
+    fut: impl Future<Output = io::Result<T>>,
+    phase: &str,
+    duration: Duration,
+) -> io::Result<T> {
+    match timeout(duration, fut).await {
+        Ok(result) => result,
+        Err(_) => Err(io::Error::new(
+            io::ErrorKind::TimedOut,
+            format!("{phase} timed out after {duration:?}"),
+        )),
+    }
+}