From 8081eb9f9281faed06ccb3c773b013b514a37245 Mon Sep 17 00:00:00 2001
From: Jun Kurihara <junkurihara@users.noreply.github.com>
Date: Tue, 23 Jan 2024 02:29:05 +0900
Subject: [PATCH] refactor: reconsider http read header timeouts

---
 rpxy-lib/src/constants.rs        |  2 +-
 rpxy-lib/src/forwarder/client.rs | 14 ++++++++++++--
 rpxy-lib/src/proxy/mod.rs        | 10 ++++++++--
 rpxy-lib/src/proxy/proxy_main.rs |  3 +--
 4 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/rpxy-lib/src/constants.rs b/rpxy-lib/src/constants.rs
index 064f5fde..b0231da8 100644
--- a/rpxy-lib/src/constants.rs
+++ b/rpxy-lib/src/constants.rs
@@ -12,7 +12,7 @@ pub const MAX_CONCURRENT_STREAMS: u32 = 64;
 pub const CERTS_WATCH_DELAY_SECS: u32 = 60;
 pub const LOAD_CERTS_ONLY_WHEN_UPDATED: bool = true;
 
-pub const CONNECTION_TIMEOUT_SEC: u64 = 30; // timeout to serve a connection. this might limits the max length of response.
+pub const CONNECTION_TIMEOUT_SEC: u64 = 60; // timeout to serve a connection, total time of receive request, serve, and send response. this might limits the max length of response.
 
 // #[cfg(feature = "http3")]
 // pub const H3_RESPONSE_BUF_SIZE: usize = 65_536; // 64KB
diff --git a/rpxy-lib/src/forwarder/client.rs b/rpxy-lib/src/forwarder/client.rs
index 2e699d05..9be7b005 100644
--- a/rpxy-lib/src/forwarder/client.rs
+++ b/rpxy-lib/src/forwarder/client.rs
@@ -133,7 +133,9 @@ Please enable native-tls-backend or rustls-backend feature to enable TLS support
     );
     let executor = LocalExecutor::new(_globals.runtime_handle.clone());
     let mut http = HttpConnector::new();
+    http.enforce_http(true);
     http.set_reuse_address(true);
+    http.set_keepalive(Some(_globals.proxy_config.upstream_idle_timeout));
     let inner = Client::builder(executor).build::<_, B>(http);
     let inner_h2 = inner.clone();
 
@@ -216,8 +218,16 @@ where
     #[cfg(not(feature = "rustls-backend-webpki"))]
     info!("Native cert store with rustls is used for the connection to backend applications");
 
-    let connector = builder.https_or_http().enable_all_versions().build();
-    let connector_h2 = builder_h2.https_or_http().enable_http2().build();
+    let mut http = HttpConnector::new();
+    http.enforce_http(false);
+    http.set_reuse_address(true);
+    http.set_keepalive(Some(_globals.proxy_config.upstream_idle_timeout));
+
+    let connector = builder
+      .https_or_http()
+      .enable_all_versions()
+      .wrap_connector(http.clone());
+    let connector_h2 = builder_h2.https_or_http().enable_http2().wrap_connector(http);
     let inner = Client::builder(LocalExecutor::new(_globals.runtime_handle.clone())).build::<_, B1>(connector);
     let inner_h2 = Client::builder(LocalExecutor::new(_globals.runtime_handle.clone())).build::<_, B1>(connector_h2);
 
diff --git a/rpxy-lib/src/proxy/mod.rs b/rpxy-lib/src/proxy/mod.rs
index 3f19059d..2cc9b755 100644
--- a/rpxy-lib/src/proxy/mod.rs
+++ b/rpxy-lib/src/proxy/mod.rs
@@ -29,8 +29,14 @@ pub(crate) fn connection_builder(globals: &Arc<Globals>) -> Arc<ConnectionBuilde
     .pipeline_flush(true);
   http_server
     .http2()
-    .keep_alive_interval(Some(globals.proxy_config.proxy_idle_timeout))
-    .timer(TokioTimer)
     .max_concurrent_streams(globals.proxy_config.max_concurrent_streams);
+
+  if globals.proxy_config.keepalive {
+    http_server
+      .http2()
+      .keep_alive_interval(Some(globals.proxy_config.proxy_idle_timeout))
+      .keep_alive_timeout(globals.proxy_config.proxy_idle_timeout + std::time::Duration::from_secs(1))
+      .timer(TokioTimer);
+  }
   Arc::new(http_server)
 }
diff --git a/rpxy-lib/src/proxy/proxy_main.rs b/rpxy-lib/src/proxy/proxy_main.rs
index 61176f3c..35768c1d 100644
--- a/rpxy-lib/src/proxy/proxy_main.rs
+++ b/rpxy-lib/src/proxy/proxy_main.rs
@@ -88,10 +88,9 @@ where
     let message_handler_clone = self.message_handler.clone();
     let tls_enabled = self.tls_enabled;
     let listening_on = self.listening_on;
-    let timeout_sec = Duration::from_secs(CONNECTION_TIMEOUT_SEC + 1); // just in case...
     self.globals.runtime_handle.clone().spawn(async move {
       timeout(
-        timeout_sec + Duration::from_secs(1), // just in case...
+        Duration::from_secs(CONNECTION_TIMEOUT_SEC) + Duration::from_secs(1), // just in case...
         server_clone.serve_connection_with_upgrades(
           stream,
           service_fn(move |req: Request<Incoming>| {