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

@@ -110,19 +110,40 @@ 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 timeout_client = ActiveStream::new(client, tracker.clone());
+    let mut timeout_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 timeout_client, &mut timeout_proxy);
+
+    let monitor_task = async {
+        loop {
+            tokio::time::sleep(Duration::from_secs(5)).await;
+            if tracker.is_idle(IDLE_TIMEOUT) {
+                break;
+            }
+        }
+    };
+
+    tokio::select! {
+        res = copy_task => {
+            if let Err(e) = res {
+                debug!("TCP relay error or closed: {}", e);
+            }
         }
-        Err(e) => {
-            debug!("TCP relay error: {}", e);
-            Ok(())
+        _ = monitor_task => {
+            debug!("TCP relay idle timeout for {}", target_addr);
         }
     }
+
+    stats::update_metrics(
+        runtime,
+        Protocol::Tcp,
+        proxy_name,
+        target_addr,
+        timeout_client.read_bytes,
+        timeout_proxy.read_bytes,
+    );
+    Ok(())
 }
 
 async fn find_session_target(
@@ -178,94 +199,62 @@ impl SharedIdleTracker {
     }
 }
 
-struct IdleTimeoutStream<T> {
+struct ActiveStream<T> {
     inner: T,
     tracker: Arc<SharedIdleTracker>,
-    timeout: Duration,
+    read_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,
+            read_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();
+                self.read_bytes += n as u64;
             }
-            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()))
+        if let Poll::Ready(Ok(n)) = &poll {
+            if *n > 0 {
+                self.update_activity();
             }
-            _ => poll,
         }
+
+        poll
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {