enforce timeout on request layers

Signed-off-by: Jason Volk <jason@zemos.net>

Cargo.toml

@@ -127,12 +127,13 @@ version = "0.6.2"
 default-features = false
 features = [
     "add-extension",
+    "catch-panic",
     "cors",
     "sensitive-headers",
     "set-header",
+    "timeout",
     "trace",
     "util",
-    "catch-panic",
 ]
 
 [workspace.dependencies.rustls]

conduwuit-example.toml

@@ -377,6 +377,18 @@
 #
 #pusher_idle_timeout = 15
 
+# Maximum time to receive a request from a client (seconds).
+#
+#client_receive_timeout = 75
+
+# Maximum time to process a request received from a client (seconds).
+#
+#client_request_timeout = 180
+
+# Maximum time to transmit a response to a client (seconds)
+#
+#client_response_timeout = 120
+
 # Enables registration. If set to false, no users can register on this
 # server.
 #

src/core/config/mod.rs

@@ -480,6 +480,24 @@ pub struct Config {
 	#[serde(default = "default_pusher_idle_timeout")]
 	pub pusher_idle_timeout: u64,
 
+	/// Maximum time to receive a request from a client (seconds).
+	///
+	/// default: 75
+	#[serde(default = "default_client_receive_timeout")]
+	pub client_receive_timeout: u64,
+
+	/// Maximum time to process a request received from a client (seconds).
+	///
+	/// default: 180
+	#[serde(default = "default_client_request_timeout")]
+	pub client_request_timeout: u64,
+
+	/// Maximum time to transmit a response to a client (seconds)
+	///
+	/// default: 120
+	#[serde(default = "default_client_response_timeout")]
+	pub client_response_timeout: u64,
+
 	/// Enables registration. If set to false, no users can register on this
 	/// server.
 	///
@@ -2170,3 +2188,9 @@ fn default_stream_width_default() -> usize { 32 }
 fn default_stream_width_scale() -> f32 { 1.0 }
 
 fn default_stream_amplification() -> usize { 1024 }
+
+fn default_client_receive_timeout() -> u64 { 75 }
+
+fn default_client_request_timeout() -> u64 { 180 }
+
+fn default_client_response_timeout() -> u64 { 120 }

src/router/layers.rs

@@ -18,6 +18,7 @@ use tower_http::{
 	cors::{self, CorsLayer},
 	sensitive_headers::SetSensitiveHeadersLayer,
 	set_header::SetResponseHeaderLayer,
+	timeout::{RequestBodyTimeoutLayer, ResponseBodyTimeoutLayer, TimeoutLayer},
 	trace::{DefaultOnFailure, DefaultOnRequest, DefaultOnResponse, TraceLayer},
 };
 use tracing::Level;
@@ -59,6 +60,9 @@ pub(crate) fn build(services: &Arc<Services>) -> Result<(Router, Guard)> {
 		)
 		.layer(axum::middleware::from_fn_with_state(Arc::clone(services), request::handle))
 		.layer(SecureClientIpSource::ConnectInfo.into_extension())
+		.layer(ResponseBodyTimeoutLayer::new(Duration::from_secs(server.config.client_response_timeout)))
+		.layer(RequestBodyTimeoutLayer::new(Duration::from_secs(server.config.client_receive_timeout)))
+		.layer(TimeoutLayer::new(Duration::from_secs(server.config.client_request_timeout)))
 		.layer(SetResponseHeaderLayer::if_not_present(
 			HeaderName::from_static("origin-agent-cluster"), // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Origin-Agent-Cluster
 			HeaderValue::from_static("?1"),

src/router/request.rs

@@ -10,8 +10,10 @@ use axum::{
 use conduwuit::{debug, debug_error, debug_warn, err, error, trace, Result};
 use conduwuit_service::Services;
 use http::{Method, StatusCode, Uri};
+use tracing::Span;
 
 #[tracing::instrument(
+	name = "request",
 	level = "debug",
 	skip_all,
 	fields(
@@ -57,23 +59,34 @@ pub(crate) async fn handle(
 	let uri = req.uri().clone();
 	let method = req.method().clone();
 	let services_ = services.clone();
-	let task = services
-		.server
-		.runtime()
-		.spawn(async move { execute(services_, req, next).await });
+	let parent = Span::current();
+	let task = services.server.runtime().spawn(async move {
+		tokio::select! {
+			response = execute(&services_, req, next, parent) => response,
+			() = services_.server.until_shutdown() =>
+				StatusCode::SERVICE_UNAVAILABLE.into_response(),
+		}
+	});
 
 	task.await
 		.map_err(unhandled)
 		.and_then(move |result| handle_result(&method, &uri, result))
 }
 
+#[tracing::instrument(
+	name = "handle",
+	level = "debug",
+	parent = parent,
+	skip_all,
+)]
 async fn execute(
 	// we made a safety contract that Services will not go out of scope
 	// during the request; this ensures a reference is accounted for at
 	// the base frame of the task regardless of its detachment.
-	_services: Arc<Services>,
+	_services: &Arc<Services>,
 	req: http::Request<axum::body::Body>,
 	next: axum::middleware::Next,
+	parent: Span,
 ) -> Response {
 	next.run(req).await
 }