fix(dgw): proper timeout for HTTP listeners (#561)

Devolutions · Oct 16, 2023 · 90a0725 · 90a0725
1 parent 375ec71
commit 90a0725
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 20 deletions.
diff --git a/devolutions-gateway/Cargo.toml b/devolutions-gateway/Cargo.toml
@@ -66,7 +66,7 @@ tokio-rustls = { version = "0.24", features = ["dangerous_configuration", "tls12
 reqwest = { version = "0.11", features = ["json"] } # TODO: directly use hyper in subscriber module
 futures = "0.3"
 async-trait = "0.1"
-tower = { version = "0.4", features = ["timeout"] }
+tower = "0.4"
 ngrok = "0.13"
 
 # HTTP

diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs
@@ -44,26 +44,23 @@ pub struct DgwState {
 }
 
 pub fn make_http_service(state: DgwState) -> axum::Router<()> {
+    use tower::ServiceBuilder;
+
     trace!("make http service");
 
     axum::Router::new()
         .merge(api::make_router(state.clone()))
         .nest_service("/KdcProxy", api::kdc_proxy::make_router(state.clone()))
         .nest_service("/jet/KdcProxy", api::kdc_proxy::make_router(state.clone()))
-        .layer(axum::middleware::from_fn_with_state(
-            state,
-            middleware::auth::auth_middleware,
-        ))
-        .layer(middleware::cors::make_middleware())
-        .layer(axum::middleware::from_fn(middleware::log::log_middleware))
         .layer(
-            tower::ServiceBuilder::new()
-                .layer(axum::error_handling::HandleErrorLayer::new(
-                    |error: tower::BoxError| async move {
-                        debug!(%error, "timed out");
-                        axum::http::StatusCode::REQUEST_TIMEOUT
-                    },
-                ))
-                .timeout(std::time::Duration::from_secs(15)),
+            // NOTE: It is recommended to use `tower::ServiceBuilder` when applying multiple middlewares at once:
+            // https://docs.rs/axum/0.6.20/axum/middleware/index.html#applying-multiple-middleware
+            ServiceBuilder::new()
+                .layer(axum::middleware::from_fn(middleware::log::log_middleware))
+                .layer(middleware::cors::make_middleware())
+                .layer(axum::middleware::from_fn_with_state(
+                    state,
+                    middleware::auth::auth_middleware,
+                )),
         )
 }
diff --git a/devolutions-gateway/src/listener.rs b/devolutions-gateway/src/listener.rs
@@ -1,6 +1,7 @@
 use anyhow::Context;
 use async_trait::async_trait;
 use devolutions_gateway_task::{ChildTask, ShutdownSignal, Task};
+use futures::TryFutureExt as _;
 use std::net::SocketAddr;
 use tap::Pipe as _;
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -12,6 +13,8 @@ use crate::generic_client::GenericClient;
 use crate::utils::url_to_socket_addr;
 use crate::DgwState;
 
+const HTTP_REQUEST_TIMEOUT: tokio::time::Duration = tokio::time::Duration::from_secs(15);
+
 #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
 #[derive(Debug, Clone, Serialize)]
 pub struct ListenerUrls {
@@ -167,11 +170,12 @@ async fn run_http_listener(listener: TcpListener, state: DgwState) -> anyhow::Re
             Ok((stream, peer_addr)) => {
                 let state = state.clone();
 
-                let fut = async move {
+                let fut = tokio::time::timeout(HTTP_REQUEST_TIMEOUT, async move {
                     if let Err(e) = handle_http_peer(stream, state, peer_addr).await {
                         error!(error = format!("{e:#}"), "handle_http_peer failed");
                     }
-                }
+                })
+                .map_err(|error| warn!(%error, "request timed out"))
                 .instrument(info_span!("http", client = %peer_addr));
 
                 ChildTask::spawn(fut).detach();
@@ -194,11 +198,12 @@ async fn run_https_listener(listener: TcpListener, state: DgwState) -> anyhow::R
                 let tls_acceptor = tls_conf.acceptor.clone();
                 let state = state.clone();
 
-                let fut = async move {
+                let fut = tokio::time::timeout(HTTP_REQUEST_TIMEOUT, async move {
                     if let Err(e) = handle_https_peer(stream, tls_acceptor, state, peer_addr).await {
                         error!(error = format!("{e:#}"), "handle_https_peer failed");
                     }
-                }
+                })
+                .map_err(|error| warn!(%error, "request timed out"))
                 .instrument(info_span!("https", client = %peer_addr));
 
                 ChildTask::spawn(fut).detach();
@@ -237,7 +242,7 @@ where
         .serve_connection(io, app)
         .with_upgrades()
         .await
-        .context("HTTP server (TLS)")
+        .context("HTTP server")
 }
 
 pub trait ToInternalUrl {