fix(gateway): resolve TCP Connection Reset issue on Keep-Alive clients

Refactored idle timeout mechanism in TcpRelay to avoid prematurely closing Keep-Alive connections (e.g., in Chrome). Swapped out the custom active polling with tokio::select! to correctly close idle sockets without triggering TCP RST packets.

Author: yinheli

src/gateway/relay_tcp.rs

@@ -7,7 +7,7 @@ use std::{
         atomic::{AtomicU64, Ordering},
     },
     task::{Context, Poll},
-    time::{Duration, SystemTime, UNIX_EPOCH},
+    time::{Duration, Instant},
 };
 
 use anyhow::{Context as _, Result};
@@ -110,19 +110,39 @@ async fn copy_with_idle_timeout(
 ) -> Result<()> {
     let tracker = Arc::new(SharedIdleTracker::new());
 
-    let mut timeout_client = IdleTimeoutStream::new(client, tracker.clone(), IDLE_TIMEOUT);
-    let mut timeout_proxy = IdleTimeoutStream::new(proxy, tracker, IDLE_TIMEOUT);
+    let mut active_client = ActiveStream::new(client, tracker.clone());
+    let mut active_proxy = ActiveStream::new(proxy, tracker.clone());
 
-    match copy_bidirectional(&mut timeout_client, &mut timeout_proxy).await {
-        Ok((up, down)) => {
-            stats::update_metrics(runtime, Protocol::Tcp, proxy_name, target_addr, up, down);
-            Ok(())
+    let copy_task = copy_bidirectional(&mut active_client, &mut active_proxy);
+
+    let monitor_task = async {
+        loop {
+            let elapsed = tracker.elapsed();
+            if elapsed >= IDLE_TIMEOUT {
+                break;
+            }
+            tokio::time::sleep(IDLE_TIMEOUT - elapsed).await;
+        }
+    };
+
+    let (up, down) = tokio::select! {
+        res = copy_task => {
+            match res {
+                Ok((up, down)) => (up, down),
+                Err(e) => {
+                    debug!("TCP relay error or closed: {}", e);
+                    (active_proxy.written_bytes, active_client.written_bytes)
+                }
+            }
         }
-        Err(e) => {
-            debug!("TCP relay error: {}", e);
-            Ok(())
+        _ = monitor_task => {
+            debug!("TCP relay idle timeout for {}", target_addr);
+            (active_proxy.written_bytes, active_client.written_bytes)
         }
-    }
+    };
+
+    stats::update_metrics(runtime, Protocol::Tcp, proxy_name, target_addr, up, down);
+    Ok(())
 }
 
 async fn find_session_target(
@@ -141,131 +161,86 @@ async fn find_session_target(
 }
 
 struct SharedIdleTracker {
+    base_instant: Instant,
     last_activity: Arc<AtomicU64>,
 }
 
 impl SharedIdleTracker {
     fn new() -> Self {
-        let now_millis = SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .unwrap()
-            .as_millis() as u64;
-
         Self {
-            last_activity: Arc::new(AtomicU64::new(now_millis)),
+            base_instant: Instant::now(),
+            last_activity: Arc::new(AtomicU64::new(0)),
         }
     }
 
     fn update_activity(&self) {
-        let now_millis = SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .unwrap()
-            .as_millis() as u64;
-        self.last_activity.store(now_millis, Ordering::Relaxed);
+        let elapsed = self.base_instant.elapsed().as_millis() as u64;
+        self.last_activity.store(elapsed, Ordering::Relaxed);
     }
 
     fn elapsed(&self) -> Duration {
         let last_millis = self.last_activity.load(Ordering::Relaxed);
-        let now_millis = SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .unwrap()
-            .as_millis() as u64;
+        let now_millis = self.base_instant.elapsed().as_millis() as u64;
         Duration::from_millis(now_millis.saturating_sub(last_millis))
     }
-
-    fn is_idle(&self, timeout: Duration) -> bool {
-        self.elapsed() > timeout
-    }
 }
 
-struct IdleTimeoutStream<T> {
+struct ActiveStream<T> {
     inner: T,
     tracker: Arc<SharedIdleTracker>,
-    timeout: Duration,
+    written_bytes: u64,
 }
 
-impl<T> IdleTimeoutStream<T> {
-    fn new(inner: T, tracker: Arc<SharedIdleTracker>, timeout: Duration) -> Self {
+impl<T> ActiveStream<T> {
+    fn new(inner: T, tracker: Arc<SharedIdleTracker>) -> Self {
         Self {
             inner,
             tracker,
-            timeout,
+            written_bytes: 0,
         }
     }
 
     fn update_activity(&self) {
         self.tracker.update_activity();
     }
-
-    fn check_idle(&self) -> tokio::io::Result<()> {
-        if self.tracker.is_idle(self.timeout) {
-            return Err(tokio::io::Error::new(
-                tokio::io::ErrorKind::TimedOut,
-                "idle timeout - no activity on either side",
-            ));
-        }
-        Ok(())
-    }
-
-    fn is_normal_close(e: &std::io::Error) -> bool {
-        matches!(
-            e.kind(),
-            std::io::ErrorKind::BrokenPipe
-                | std::io::ErrorKind::ConnectionReset
-                | std::io::ErrorKind::ConnectionAborted
-                | std::io::ErrorKind::UnexpectedEof
-        )
-    }
 }
 
-impl<T: AsyncRead + Unpin> AsyncRead for IdleTimeoutStream<T> {
+impl<T: AsyncRead + Unpin> AsyncRead for ActiveStream<T> {
     fn poll_read(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         buf: &mut tokio::io::ReadBuf<'_>,
     ) -> Poll<std::io::Result<()>> {
-        self.check_idle()?;
-
         let initial_filled = buf.filled().len();
         let poll = Pin::new(&mut self.inner).poll_read(cx, buf);
 
-        match &poll {
-            Poll::Ready(Ok(())) if buf.filled().len() > initial_filled => {
+        if let Poll::Ready(Ok(())) = &poll {
+            let n = buf.filled().len() - initial_filled;
+            if n > 0 {
                 self.update_activity();
             }
-            Poll::Ready(Err(e)) if Self::is_normal_close(e) => {
-                return Poll::Ready(Ok(()));
-            }
-            _ => {}
         }
 
         poll
     }
 }
 
-impl<T: AsyncWrite + Unpin> AsyncWrite for IdleTimeoutStream<T> {
+impl<T: AsyncWrite + Unpin> AsyncWrite for ActiveStream<T> {
     fn poll_write(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<Result<usize, Error>> {
-        self.check_idle()?;
-
         let poll = Pin::new(&mut self.inner).poll_write(cx, buf);
 
-        match poll {
-            Poll::Ready(Ok(n)) => {
-                if n > 0 {
-                    self.update_activity();
-                }
-                Poll::Ready(Ok(n))
-            }
-            Poll::Ready(Err(e)) if Self::is_normal_close(&e) => {
-                // Treat normal close as successful write of all bytes
-                Poll::Ready(Ok(buf.len()))
-            }
-            _ => poll,
+        if let Poll::Ready(Ok(n)) = &poll
+            && *n > 0
+        {
+            self.update_activity();
+            self.written_bytes += *n as u64;
         }
+
+        poll
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {

src/main.rs

@@ -1,4 +1,4 @@
-#![feature(test)]
+#![cfg_attr(test, feature(test))]
 
 #[macro_use]
 extern crate lazy_static;