From 075086d28c3dc1bb21598dd7d6d090d1d40f52d1 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Sun, 15 Dec 2024 04:28:34 +0000
Subject: [PATCH] tproxy: Config for overall timeout for a connection

---
 tproxy/src/config.rs               |  2 ++
 tproxy/src/proxy.rs                | 28 +++++++++++++++++++---------
 tproxy/src/proxy/tls_passthough.rs |  2 +-
 tproxy/src/proxy/tls_terminate.rs  |  2 +-
 tproxy/tproxy.toml                 |  2 ++
 5 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/tproxy/src/config.rs b/tproxy/src/config.rs
index 2b2fbde2..d1913f2a 100644
--- a/tproxy/src/config.rs
+++ b/tproxy/src/config.rs
@@ -47,6 +47,8 @@ pub struct Timeouts {
     pub write: Duration,
     #[serde(with = "serde_duration")]
     pub shutdown: Duration,
+    #[serde(with = "serde_duration")]
+    pub total: Duration,
 }
 
 #[derive(Debug, Clone, Deserialize)]
diff --git a/tproxy/src/proxy.rs b/tproxy/src/proxy.rs
index c48b158d..b5fd2e1e 100644
--- a/tproxy/src/proxy.rs
+++ b/tproxy/src/proxy.rs
@@ -153,20 +153,30 @@ pub async fn run(config: &ProxyConfig, app_state: AppState) -> Result<()> {
     loop {
         match listener.accept().await {
             Ok((inbound, addr)) => {
-                info!("new connection from {addr}");
+                info!(%addr, "new connection received");
                 let app_state = app_state.clone();
                 let dotted_base_domain = dotted_base_domain.clone();
                 let tls_terminate_proxy = tls_terminate_proxy.clone();
                 tokio::spawn(async move {
-                    if let Err(e) = handle_connection(
-                        inbound,
-                        app_state,
-                        &dotted_base_domain,
-                        tls_terminate_proxy,
+                    let timeouts = &app_state.config.proxy.timeouts;
+                    let result = timeout(
+                        timeouts.total,
+                        handle_connection(
+                            inbound,
+                            app_state,
+                            &dotted_base_domain,
+                            tls_terminate_proxy,
+                        ),
                     )
-                    .await
-                    {
-                        error!("connection error: {e:?}");
+                    .await;
+                    match result {
+                        Ok(Ok(_)) => {}
+                        Ok(Err(e)) => {
+                            error!("connection error: {e:?}");
+                        }
+                        Err(_) => {
+                            info!(%addr, "connection kept too long");
+                        }
                     }
                 });
             }
diff --git a/tproxy/src/proxy/tls_passthough.rs b/tproxy/src/proxy/tls_passthough.rs
index f2856739..dcd3255a 100644
--- a/tproxy/src/proxy/tls_passthough.rs
+++ b/tproxy/src/proxy/tls_passthough.rs
@@ -72,7 +72,7 @@ pub(crate) async fn proxy_to_app(
         TcpStream::connect((target_ip, port)),
     )
     .await
-    .context("connection timeout")?
+    .context("connecting timeout")?
     .context("failed to connect to tapp")?;
     outbound
         .write_all(&buffer)
diff --git a/tproxy/src/proxy/tls_terminate.rs b/tproxy/src/proxy/tls_terminate.rs
index 8745af9b..5b45722a 100644
--- a/tproxy/src/proxy/tls_terminate.rs
+++ b/tproxy/src/proxy/tls_terminate.rs
@@ -145,7 +145,7 @@ impl TlsTerminateProxy {
             TcpStream::connect((host.ip, port)),
         )
         .await
-        .map_err(|_| anyhow::anyhow!("connection timeout"))?
+        .map_err(|_| anyhow::anyhow!("connecting timeout"))?
         .context("failed to connect to app")?;
         bridge(
             IgnoreUnexpectedEofStream::new(tls_stream),
diff --git a/tproxy/tproxy.toml b/tproxy/tproxy.toml
index 9a2caf22..40783e8f 100644
--- a/tproxy/tproxy.toml
+++ b/tproxy/tproxy.toml
@@ -49,6 +49,8 @@ idle = "10m"
 write = "5s"
 # Timeout for shutting down a connection.
 shutdown = "5s"
+# Timeout for total connection duration.
+total = "5h"
 
 [core.recycle]
 enabled = true