TLS handshake timeout (Fix #37)

Cargo.toml

@@ -13,7 +13,7 @@ version = "0.3.3"
 
 [features]
 default = []
-tls-rustls = ["arc-swap", "pin-project-lite", "rustls", "rustls-pemfile", "tokio/fs", "tokio-rustls"]
+tls-rustls = ["arc-swap", "pin-project-lite", "rustls", "rustls-pemfile", "tokio/fs", "tokio/time", "tokio-rustls"]
 
 [dependencies]
 bytes = "1"

examples/configure_addr_incoming.rs

@@ -0,0 +1,27 @@
+//! Run with `cargo run --example configure_http` command.
+//!
+//! To connect through browser, navigate to "http://localhost:3000" url.
+
+use axum::{routing::get, Router};
+use axum_server::AddrIncomingConfig;
+use std::net::SocketAddr;
+use std::time::Duration;
+
+#[tokio::main]
+async fn main() {
+    let app = Router::new().route("/", get(|| async { "Hello, world!" }));
+
+    let config = AddrIncomingConfig::new()
+        .tcp_nodelay(true)
+        .tcp_sleep_on_accept_errors(true)
+        .tcp_keepalive(Some(Duration::from_secs(32)))
+        .build();
+
+    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
+    println!("listening on {}", addr);
+    axum_server::bind(addr)
+        .addr_incoming_config(config)
+        .serve(app.into_make_service())
+        .await
+        .unwrap();
+}

src/addr_incoming_config.rs

@@ -0,0 +1,55 @@
+use std::time::Duration;
+
+/// A configuration for [`AddrIncoming`].
+#[derive(Debug, Clone)]
+pub struct AddrIncomingConfig {
+    pub(crate) tcp_sleep_on_accept_errors: bool,
+    pub(crate) tcp_keepalive: Option<Duration>,
+    pub(crate) tcp_nodelay: bool,
+}
+
+impl Default for AddrIncomingConfig {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl AddrIncomingConfig {
+    /// Creates a default [`AddrIncoming`] config.
+    pub fn new() -> AddrIncomingConfig {
+        Self {
+            tcp_sleep_on_accept_errors: true,
+            tcp_keepalive: None,
+            tcp_nodelay: false,
+        }
+    }
+
+    /// Builds the config, creating an owned version of it.
+    pub fn build(&mut self) -> Self {
+        self.clone()
+    }
+
+    /// Set whether to sleep on accept errors, to avoid exhausting file descriptor limits.
+    ///
+    /// Default is `true`.
+    pub fn tcp_sleep_on_accept_errors(&mut self, val: bool) -> &mut Self {
+        self.tcp_sleep_on_accept_errors = val;
+        self
+    }
+
+    /// Set how often to send TCP keepalive probes.
+    ///
+    /// Default is `false`.
+    pub fn tcp_keepalive(&mut self, val: Option<Duration>) -> &mut Self {
+        self.tcp_keepalive = val;
+        self
+    }
+
+    /// Set the value of `TCP_NODELAY` option for accepted connections.
+    ///
+    /// Default is `false`.
+    pub fn tcp_nodelay(&mut self, val: bool) -> &mut Self {
+        self.tcp_nodelay = val;
+        self
+    }
+}

src/lib.rs

@@ -87,6 +87,7 @@
 )]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 
+mod addr_incoming_config;
 mod handle;
 mod http_config;
 mod notify_once;
@@ -96,6 +97,7 @@ pub mod accept;
 pub mod service;
 
 pub use self::{
+    addr_incoming_config::AddrIncomingConfig,
     handle::Handle,
     http_config::HttpConfig,
     server::{bind, Server},

src/server.rs

@@ -1,3 +1,4 @@
+use crate::addr_incoming_config::AddrIncomingConfig;
 use crate::{
     accept::{Accept, DefaultAcceptor},
     handle::Handle,
@@ -25,6 +26,7 @@ use tokio::{
 pub struct Server<A = DefaultAcceptor> {
     acceptor: A,
     addr: SocketAddr,
+    addr_incoming_conf: AddrIncomingConfig,
     handle: Handle,
     http_conf: HttpConfig,
 }
@@ -43,6 +45,7 @@ impl Server {
         Self {
             acceptor,
             addr,
+            addr_incoming_conf: AddrIncomingConfig::default(),
             handle,
             http_conf: HttpConfig::default(),
         }
@@ -55,6 +58,7 @@ impl<A> Server<A> {
         Server {
             acceptor,
             addr: self.addr,
+            addr_incoming_conf: self.addr_incoming_conf,
             handle: self.handle,
             http_conf: self.http_conf,
         }
@@ -72,6 +76,12 @@ impl<A> Server<A> {
         self
     }
 
+    /// Overwrite addr incoming configuration.
+    pub fn addr_incoming_config(mut self, config: AddrIncomingConfig) -> Self {
+        self.addr_incoming_conf = config;
+        self
+    }
+
     /// Serve provided [`MakeService`].
     ///
     /// # Errors
@@ -93,11 +103,15 @@ impl<A> Server<A> {
         A::Future: Send,
     {
         let acceptor = self.acceptor;
+        let addr_incoming_conf = self.addr_incoming_conf;
         let handle = self.handle;
         let http_conf = self.http_conf;
 
         let listener = TcpListener::bind(self.addr).await?;
         let mut incoming = AddrIncoming::from_listener(listener).map_err(io_other)?;
+        incoming.set_sleep_on_errors(addr_incoming_conf.tcp_sleep_on_accept_errors);
+        incoming.set_keepalive(addr_incoming_conf.tcp_keepalive);
+        incoming.set_nodelay(addr_incoming_conf.tcp_nodelay);
 
         handle.notify_listening(incoming.local_addr());
 

src/tls_rustls/future.rs

@@ -2,6 +2,8 @@
 
 use crate::tls_rustls::RustlsConfig;
 use pin_project_lite::pin_project;
+use std::io::{Error, ErrorKind};
+use std::time::Duration;
 use std::{
     fmt,
     future::Future,
@@ -10,6 +12,7 @@ use std::{
     task::{Context, Poll},
 };
 use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::time::{timeout, Timeout};
 use tokio_rustls::{server::TlsStream, Accept, TlsAcceptor};
 
 pin_project! {
@@ -22,8 +25,11 @@ pin_project! {
 }
 
 impl<F, I, S> RustlsAcceptorFuture<F, I, S> {
-    pub(crate) fn new(future: F, config: RustlsConfig) -> Self {
-        let inner = AcceptFuture::Inner { future };
+    pub(crate) fn new(future: F, config: RustlsConfig, handshake_timeout: Duration) -> Self {
+        let inner = AcceptFuture::Inner {
+            future,
+            handshake_timeout,
+        };
         let config = Some(config);
 
         Self { inner, config }
@@ -42,10 +48,11 @@ pin_project! {
         Inner {
             #[pin]
             future: F,
+            handshake_timeout: Duration,
         },
         Accept {
             #[pin]
-            future: Accept<I>,
+            future: Timeout<Accept<I>>,
             service: Option<S>,
         },
     }
@@ -63,7 +70,10 @@ where
 
         loop {
             match this.inner.as_mut().project() {
-                AcceptFutureProj::Inner { future } => {
+                AcceptFutureProj::Inner {
+                    future,
+                    handshake_timeout,
+                } => {
                     match future.poll(cx) {
                         Poll::Ready(Ok((stream, service))) => {
                             let server_config = this.config
@@ -75,20 +85,27 @@ where
                             let future = acceptor.accept(stream);
 
                             let service = Some(service);
+                            let handshake_timeout = *handshake_timeout;
 
-                            this.inner.set(AcceptFuture::Accept { future, service });
+                            this.inner.set(AcceptFuture::Accept {
+                                future: timeout(handshake_timeout, future),
+                                service,
+                            });
                         }
                         Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                         Poll::Pending => return Poll::Pending,
                     }
                 }
                 AcceptFutureProj::Accept { future, service } => match future.poll(cx) {
-                    Poll::Ready(Ok(stream)) => {
+                    Poll::Ready(Ok(Ok(stream))) => {
                         let service = service.take().expect("future polled after ready");
 
                         return Poll::Ready(Ok((stream, service)));
                     }
-                    Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
+                    Poll::Ready(Ok(Err(e))) => return Poll::Ready(Err(e)),
+                    Poll::Ready(Err(timeout)) => {
+                        return Poll::Ready(Err(Error::new(ErrorKind::TimedOut, timeout)))
+                    }
                     Poll::Pending => return Poll::Pending,
                 },
             }

src/tls_rustls/mod.rs

@@ -34,6 +34,7 @@ use crate::{
 };
 use arc_swap::ArcSwap;
 use rustls::{Certificate, PrivateKey, ServerConfig};
+use std::time::Duration;
 use std::{fmt, io, net::SocketAddr, path::Path, sync::Arc};
 use tokio::{
     io::{AsyncRead, AsyncWrite},
@@ -65,14 +66,32 @@ pub fn bind_rustls(addr: SocketAddr, config: RustlsConfig) -> Server<RustlsAccep
 pub struct RustlsAcceptor<A = DefaultAcceptor> {
     inner: A,
     config: RustlsConfig,
+    handshake_timeout: Duration,
 }
 
 impl RustlsAcceptor {
     /// Create a new rustls acceptor.
     pub fn new(config: RustlsConfig) -> Self {
         let inner = DefaultAcceptor::new();
 
-        Self { inner, config }
+        #[cfg(not(test))]
+        let handshake_timeout = Duration::from_secs(10);
+
+        // Don't force tests to wait too long.
+        #[cfg(test)]
+        let handshake_timeout = Duration::from_secs(1);
+
+        Self {
+            inner,
+            config,
+            handshake_timeout,
+        }
+    }
+
+    /// Override the default TLS handshake timeout of 10 seconds, except during testing.
+    pub fn handshake_timeout(mut self, val: Duration) -> Self {
+        self.handshake_timeout = val;
+        self
     }
 }
 
@@ -82,6 +101,7 @@ impl<A> RustlsAcceptor<A> {
         RustlsAcceptor {
             inner: acceptor,
             config: self.config,
+            handshake_timeout: self.handshake_timeout,
         }
     }
 }
@@ -99,7 +119,7 @@ where
         let inner_future = self.inner.accept(stream, service);
         let config = self.config.clone();
 
-        RustlsAcceptorFuture::new(inner_future, config)
+        RustlsAcceptorFuture::new(inner_future, config, self.handshake_timeout)
     }
 }
 
@@ -283,6 +303,7 @@ mod tests {
         sync::Arc,
         time::{Duration, SystemTime},
     };
+    use tokio::time::sleep;
     use tokio::{net::TcpStream, task::JoinHandle, time::timeout};
     use tokio_rustls::TlsConnector;
     use tower::{Service, ServiceExt};
@@ -298,6 +319,22 @@ mod tests {
         assert_eq!(body.as_ref(), b"Hello, world!");
     }
 
+    #[tokio::test]
+    async fn tls_timeout() {
+        let (handle, _server_task, addr) = start_server().await;
+        assert_eq!(handle.connection_count(), 0);
+
+        // We intentionally avoid driving a TLS handshake to completion.
+        let _stream = TcpStream::connect(addr).await.unwrap();
+
+        sleep(Duration::from_millis(500)).await;
+        assert_eq!(handle.connection_count(), 1);
+
+        tokio::time::sleep(Duration::from_millis(1000)).await;
+        // Timeout defaults to 1s during testing, and we have waited 1.5 seconds.
+        assert_eq!(handle.connection_count(), 0);
+    }
+
     #[tokio::test]
     async fn test_reload() {
         let handle = Handle::new();