From 9c93b1cc31a3d6a154fbfd589eca62c87f5bae39 Mon Sep 17 00:00:00 2001
From: Jun Kurihara <junkurihara@users.noreply.github.com>
Date: Tue, 12 Jul 2022 01:12:47 +0900
Subject: [PATCH] add some timeouts

---
 src/msg_handler/handler.rs | 19 ++++++++++++++-----
 src/proxy/proxy_h3.rs      | 11 +++++++++--
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/msg_handler/handler.rs b/src/msg_handler/handler.rs
index 59ed44c..b438d98 100644
--- a/src/msg_handler/handler.rs
+++ b/src/msg_handler/handler.rs
@@ -91,13 +91,22 @@ where
     req_forwarded.log(&client_addr, Some("(Request to Backend)"));
 
     // Forward request to
-    let mut res_backend = match self.forwarder.request(req_forwarded).await {
-      Ok(res) => res,
-      Err(e) => {
-        error!("Failed to get response from backend: {}", e);
-        return http_error(StatusCode::BAD_REQUEST);
+    let mut res_backend = {
+      match tokio::time::timeout(self.globals.timeout, self.forwarder.request(req_forwarded)).await
+      {
+        Err(_) => {
+          return http_error(StatusCode::GATEWAY_TIMEOUT);
+        }
+        Ok(x) => match x {
+          Ok(res) => res,
+          Err(e) => {
+            error!("Failed to get response from backend: {}", e);
+            return http_error(StatusCode::BAD_REQUEST);
+          }
+        },
       }
     };
+
     res_backend.log_debug(
       &backend.server_name,
       &client_addr,
diff --git a/src/proxy/proxy_h3.rs b/src/proxy/proxy_h3.rs
index b9e5c30..b6f286b 100644
--- a/src/proxy/proxy_h3.rs
+++ b/src/proxy/proxy_h3.rs
@@ -17,8 +17,10 @@ where
       return;
     }
     let fut = self.clone().handle_connection_h3(conn);
+    let timeout_sec = self.globals.timeout;
     self.globals.runtime_handle.spawn(async move {
-      if let Err(e) = fut.await {
+      if let Err(e) = tokio::time::timeout(timeout_sec + Duration::from_secs(1), fut).await {
+        // TODO: ここのtimeoutはどの値を使うべき？
         warn!("QUIC or HTTP/3 connection failed: {}", e)
       }
       clients_count.decrement();
@@ -79,7 +81,12 @@ where
 
           let self_inner = self.clone();
           self.globals.runtime_handle.spawn(async move {
-            if let Err(e) = self_inner.handle_stream_h3(req, stream, client_addr).await {
+            if let Err(e) = tokio::time::timeout(
+              self_inner.globals.timeout + Duration::from_secs(1), // timeout per stream
+              self_inner.handle_stream_h3(req, stream, client_addr),
+            )
+            .await
+            {
               error!("HTTP/3 request failed: {}", e);
             }
             // // TODO: Work around for timeout