From 0ceb58586b94cd67979e1b7b68807d6825dfbbe3 Mon Sep 17 00:00:00 2001
From: "Sijie.Sun" <sunsijie@buaa.edu.cn>
Date: Tue, 26 Aug 2025 23:30:30 +0800
Subject: [PATCH] fix keepalive on accepted tcp proxy connection (#1302)

---
 easytier/src/gateway/tcp_proxy.rs | 39 ++++++++++++++++++-------------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/easytier/src/gateway/tcp_proxy.rs b/easytier/src/gateway/tcp_proxy.rs
index 17ec194..1ead0cf 100644
--- a/easytier/src/gateway/tcp_proxy.rs
+++ b/easytier/src/gateway/tcp_proxy.rs
@@ -71,27 +71,12 @@ impl NatDstConnector for NatDstTcpConnector {
                 return Err(e.into());
             }
         };
-        if let Err(e) = socket.set_nodelay(true) {
-            tracing::warn!("set_nodelay failed, ignore it: {:?}", e);
-        }
-
-        const TCP_KEEPALIVE_TIME: Duration = Duration::from_secs(5);
-        const TCP_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(2);
-        const TCP_KEEPALIVE_RETRIES: u32 = 2;
 
         let stream = timeout(Duration::from_secs(10), socket.connect(nat_dst))
             .await?
             .with_context(|| format!("connect to nat dst failed: {:?}", nat_dst))?;
 
-        let ka = TcpKeepalive::new()
-            .with_time(TCP_KEEPALIVE_TIME)
-            .with_interval(TCP_KEEPALIVE_INTERVAL);
-
-        #[cfg(not(target_os = "windows"))]
-        let ka = ka.with_retries(TCP_KEEPALIVE_RETRIES);
-
-        let sf = SockRef::from(&stream);
-        sf.set_tcp_keepalive(&ka)?;
+        prepare_kernel_tcp_socket(&stream)?;
 
         Ok(stream)
     }
@@ -280,11 +265,33 @@ enum ProxyTcpListener {
     SmolTcpListener(SmolTcpListener),
 }
 
+fn prepare_kernel_tcp_socket(stream: &TcpStream) -> Result<()> {
+    const TCP_KEEPALIVE_TIME: Duration = Duration::from_secs(5);
+    const TCP_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(2);
+    const TCP_KEEPALIVE_RETRIES: u32 = 2;
+
+    let ka = TcpKeepalive::new()
+        .with_time(TCP_KEEPALIVE_TIME)
+        .with_interval(TCP_KEEPALIVE_INTERVAL);
+
+    #[cfg(not(target_os = "windows"))]
+    let ka = ka.with_retries(TCP_KEEPALIVE_RETRIES);
+
+    let sf = SockRef::from(&stream);
+    sf.set_tcp_keepalive(&ka)?;
+    if let Err(e) = sf.set_nodelay(true) {
+        tracing::warn!("set_nodelay failed, ignore it: {:?}", e);
+    }
+
+    Ok(())
+}
+
 impl ProxyTcpListener {
     pub async fn accept(&mut self) -> Result<(ProxyTcpStream, SocketAddr)> {
         match self {
             Self::KernelTcpListener(listener) => {
                 let (stream, addr) = listener.accept().await?;
+                prepare_kernel_tcp_socket(&stream)?;
                 Ok((ProxyTcpStream::KernelTcpStream(stream), addr))
             }
             #[cfg(feature = "smoltcp")]