allow customizing keep-alive seconds on TCP sockets

- ref #490

Author: zonyitoo

README.md

@@ -458,6 +458,9 @@ Example configuration:
     // TCP_NODELAY
     "no_delay": false,
 
+    // Enables `SO_KEEPALIVE` and set `TCP_KEEPIDLE`, `TCP_KEEPINTVL` to the specified seconds
+    "keep_alive": 15,
+
     // Soft and Hard limit of file descriptors on *NIX systems
     "nofile": 10240,
 

crates/shadowsocks-service/src/config.rs

@@ -123,6 +123,8 @@ struct SSConfig {
     mode: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     no_delay: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    keep_alive: Option<u64>,
     #[cfg(all(unix, not(target_os = "android")))]
     #[serde(skip_serializing_if = "Option::is_none")]
     nofile: Option<u64>,
@@ -769,6 +771,9 @@ pub struct Config {
     pub no_delay: bool,
     /// Set `TCP_FASTOPEN` socket option
     pub fast_open: bool,
+    /// Set TCP Keep-Alive duration, will set both `TCP_KEEPIDLE` and `TCP_KEEPINTVL`
+    pub keep_alive: Option<Duration>,
+
     /// `RLIMIT_NOFILE` option for *nix systems
     #[cfg(all(unix, not(target_os = "android")))]
     pub nofile: Option<u64>,
@@ -883,6 +888,8 @@ impl Config {
 
             no_delay: false,
             fast_open: false,
+            keep_alive: None,
+
             #[cfg(all(unix, not(target_os = "android")))]
             nofile: None,
 
@@ -1374,6 +1381,11 @@ impl Config {
             nconfig.fast_open = b;
         }
 
+        // TCP Keep-Alive
+        if let Some(d) = config.keep_alive {
+            nconfig.keep_alive = Some(Duration::from_secs(d));
+        }
+
         // UDP
         nconfig.udp_timeout = config.udp_timeout.map(Duration::from_secs);
 
@@ -1864,6 +1876,10 @@ impl fmt::Display for Config {
             jconf.fast_open = Some(self.fast_open);
         }
 
+        if let Some(keepalive) = self.keep_alive {
+            jconf.keep_alive = Some(keepalive.as_secs());
+        }
+
         match self.dns {
             DnsConfig::System => {}
             #[cfg(feature = "trust-dns")]

crates/shadowsocks-service/src/local/mod.rs

@@ -82,14 +82,17 @@ pub async fn run(mut config: Config) -> io::Result<()> {
     };
     connect_opts.tcp.send_buffer_size = config.outbound_send_buffer_size;
     connect_opts.tcp.recv_buffer_size = config.outbound_recv_buffer_size;
+    connect_opts.tcp.nodelay = config.no_delay;
     connect_opts.tcp.fastopen = config.fast_open;
+    connect_opts.tcp.keepalive = config.keep_alive;
     context.set_connect_opts(connect_opts);
 
     let mut accept_opts = AcceptOpts::default();
     accept_opts.tcp.send_buffer_size = config.inbound_send_buffer_size;
     accept_opts.tcp.recv_buffer_size = config.inbound_recv_buffer_size;
     accept_opts.tcp.nodelay = config.no_delay;
     accept_opts.tcp.fastopen = config.fast_open;
+    accept_opts.tcp.keepalive = config.keep_alive;
 
     if let Some(resolver) = build_dns_resolver(config.dns, config.ipv6_first, context.connect_opts_ref()).await {
         context.set_dns_resolver(Arc::new(resolver));
@@ -215,9 +218,6 @@ pub async fn run(mut config: Config) -> io::Result<()> {
                 if let Some(b) = local_config.udp_addr {
                     server.set_udp_bind_addr(b.clone());
                 }
-                if config.no_delay {
-                    server.set_nodelay(true);
-                }
 
                 vfut.push(async move { server.run(&client_addr, balancer).await }.boxed());
             }
@@ -236,9 +236,6 @@ pub async fn run(mut config: Config) -> io::Result<()> {
                     server.set_udp_expiry_duration(d);
                 }
                 server.set_mode(local_config.mode);
-                if config.no_delay {
-                    server.set_nodelay(true);
-                }
 
                 let udp_addr = local_config.udp_addr.unwrap_or_else(|| client_addr.clone());
                 vfut.push(async move { server.run(&client_addr, &udp_addr, balancer).await }.boxed());
@@ -262,9 +259,6 @@ pub async fn run(mut config: Config) -> io::Result<()> {
                     server.set_udp_expiry_duration(d);
                 }
                 server.set_mode(local_config.mode);
-                if config.no_delay {
-                    server.set_nodelay(true);
-                }
                 server.set_tcp_redir(local_config.tcp_redir);
                 server.set_udp_redir(local_config.udp_redir);
 

crates/shadowsocks-service/src/local/redir/server.rs

@@ -18,7 +18,6 @@ pub struct Redir {
     mode: Mode,
     udp_expiry_duration: Option<Duration>,
     udp_capacity: Option<usize>,
-    nodelay: bool,
     tcp_redir: RedirType,
     udp_redir: RedirType,
 }
@@ -43,7 +42,6 @@ impl Redir {
             mode: Mode::TcpOnly,
             udp_expiry_duration: None,
             udp_capacity: None,
-            nodelay: false,
             tcp_redir: RedirType::tcp_default(),
             udp_redir: RedirType::udp_default(),
         }
@@ -64,11 +62,6 @@ impl Redir {
         self.mode = mode;
     }
 
-    /// Set `TCP_NODELAY`
-    pub fn set_nodelay(&mut self, nodelay: bool) {
-        self.nodelay = nodelay;
-    }
-
     /// Set transparent proxy type of TCP relay, which is platform dependent
     pub fn set_tcp_redir(&mut self, ty: RedirType) {
         self.tcp_redir = ty;
@@ -96,14 +89,7 @@ impl Redir {
     }
 
     async fn run_tcp_tunnel(&self, client_config: &ServerAddr, balancer: PingBalancer) -> io::Result<()> {
-        run_tcp_redir(
-            self.context.clone(),
-            client_config,
-            balancer,
-            self.tcp_redir,
-            self.nodelay,
-        )
-        .await
+        run_tcp_redir(self.context.clone(), client_config, balancer, self.tcp_redir).await
     }
 
     async fn run_udp_tunnel(&self, client_config: &ServerAddr, balancer: PingBalancer) -> io::Result<()> {

crates/shadowsocks-service/src/local/redir/tcprelay/mod.rs

@@ -39,17 +39,12 @@ async fn establish_client_tcp_redir<'a>(
     mut stream: TcpStream,
     peer_addr: SocketAddr,
     addr: &Address,
-    nodelay: bool,
 ) -> io::Result<()> {
     let server = balancer.best_tcp_server();
     let svr_cfg = server.server_config();
 
     let mut remote = AutoProxyClientStream::connect(context, &server, addr).await?;
 
-    if nodelay {
-        remote.set_nodelay(true)?;
-    }
-
     establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, addr).await
 }
 
@@ -59,20 +54,7 @@ async fn handle_redir_client(
     s: TcpStream,
     peer_addr: SocketAddr,
     mut daddr: SocketAddr,
-    nodelay: bool,
 ) -> io::Result<()> {
-    // let svr_cfg = server.server_config();
-    //
-    // if let Err(err) = s.set_keepalive(svr_cfg.timeout()) {
-    //     error!("failed to set keep alive: {:?}", err);
-    // }
-
-    if nodelay {
-        if let Err(err) = s.set_nodelay(true) {
-            error!("failed to set TCP_NODELAY on accepted socket, error: {:?}", err);
-        }
-    }
-
     // Get forward address from socket
     //
     // Try to convert IPv4 mapped IPv6 address for dual-stack mode.
@@ -82,15 +64,14 @@ async fn handle_redir_client(
         }
     }
     let target_addr = Address::from(daddr);
-    establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr, nodelay).await
+    establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr).await
 }
 
 pub async fn run_tcp_redir(
     context: Arc<ServiceContext>,
     client_config: &ServerAddr,
     balancer: PingBalancer,
     redir_ty: RedirType,
-    nodelay: bool,
 ) -> io::Result<()> {
     let listener = match *client_config {
         ServerAddr::SocketAddr(ref saddr) => TcpListener::bind_redir(redir_ty, *saddr).await?,
@@ -137,7 +118,7 @@ pub async fn run_tcp_redir(
                 }
             };
 
-            if let Err(err) = handle_redir_client(context, balancer, socket, peer_addr, dst_addr, nodelay).await {
+            if let Err(err) = handle_redir_client(context, balancer, socket, peer_addr, dst_addr).await {
                 debug!("TCP redirect client, error: {:?}", err);
             }
         });

crates/shadowsocks-service/src/local/socks/server/mod.rs

@@ -24,7 +24,6 @@ pub struct Socks {
     udp_expiry_duration: Option<Duration>,
     udp_capacity: Option<usize>,
     udp_bind_addr: Option<ServerAddr>,
-    nodelay: bool,
 }
 
 impl Default for Socks {
@@ -48,7 +47,6 @@ impl Socks {
             udp_expiry_duration: None,
             udp_capacity: None,
             udp_bind_addr: None,
-            nodelay: false,
         }
     }
 
@@ -75,11 +73,6 @@ impl Socks {
         self.udp_bind_addr = Some(a);
     }
 
-    /// Set `TCP_NODELAY`
-    pub fn set_nodelay(&mut self, nodelay: bool) {
-        self.nodelay = nodelay;
-    }
-
     /// Start serving
     pub async fn run(self, client_config: &ServerAddr, balancer: PingBalancer) -> io::Result<()> {
         let mut vfut = Vec::new();
@@ -133,13 +126,8 @@ impl Socks {
                 }
             };
 
-            if self.nodelay {
-                let _ = stream.set_nodelay(true);
-            }
-
             let balancer = balancer.clone();
             let context = self.context.clone();
-            let nodelay = self.nodelay;
             let udp_bind_addr = udp_bind_addr.clone();
             let mode = self.mode;
 
@@ -150,7 +138,6 @@ impl Socks {
                 balancer,
                 peer_addr,
                 mode,
-                nodelay,
             ));
         }
     }
@@ -163,7 +150,6 @@ impl Socks {
         balancer: PingBalancer,
         peer_addr: SocketAddr,
         mode: Mode,
-        nodelay: bool,
     ) -> io::Result<()> {
         use std::io::ErrorKind;
 
@@ -175,12 +161,12 @@ impl Socks {
 
         match version_buffer[0] {
             0x04 => {
-                let handler = Socks4TcpHandler::new(context, nodelay, balancer, mode);
+                let handler = Socks4TcpHandler::new(context, balancer, mode);
                 handler.handle_socks4_client(stream, peer_addr).await
             }
 
             0x05 => {
-                let handler = Socks5TcpHandler::new(context, udp_bind_addr, nodelay, balancer, mode);
+                let handler = Socks5TcpHandler::new(context, udp_bind_addr, balancer, mode);
                 handler.handle_socks5_client(stream, peer_addr).await
             }
 
@@ -200,9 +186,8 @@ impl Socks {
         balancer: PingBalancer,
         peer_addr: SocketAddr,
         mode: Mode,
-        nodelay: bool,
     ) -> io::Result<()> {
-        let handler = Socks5TcpHandler::new(context, udp_bind_addr, nodelay, balancer, mode);
+        let handler = Socks5TcpHandler::new(context, udp_bind_addr, balancer, mode);
         handler.handle_socks5_client(stream, peer_addr).await
     }
 

crates/shadowsocks-service/src/local/socks/server/socks4/tcprelay.rs

@@ -24,16 +24,14 @@ use crate::local::socks::socks4::{Address, Command, HandshakeRequest, HandshakeR
 
 pub struct Socks4TcpHandler {
     context: Arc<ServiceContext>,
-    nodelay: bool,
     balancer: PingBalancer,
     mode: Mode,
 }
 
 impl Socks4TcpHandler {
-    pub fn new(context: Arc<ServiceContext>, nodelay: bool, balancer: PingBalancer, mode: Mode) -> Socks4TcpHandler {
+    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mode: Mode) -> Socks4TcpHandler {
         Socks4TcpHandler {
             context,
-            nodelay,
             balancer,
             mode,
         }
@@ -108,10 +106,6 @@ impl Socks4TcpHandler {
             }
         };
 
-        if self.nodelay {
-            remote.set_nodelay(true)?;
-        }
-
         // NOTE: Transfer all buffered data before unwrap, or these data will be lost
         let buffer = stream.buffer();
         if !buffer.is_empty() {

crates/shadowsocks-service/src/local/socks/server/socks5/tcprelay.rs

@@ -36,7 +36,6 @@ use crate::{
 pub struct Socks5TcpHandler {
     context: Arc<ServiceContext>,
     udp_bind_addr: Option<Arc<ServerAddr>>,
-    nodelay: bool,
     balancer: PingBalancer,
     mode: Mode,
 }
@@ -45,14 +44,12 @@ impl Socks5TcpHandler {
     pub fn new(
         context: Arc<ServiceContext>,
         udp_bind_addr: Option<Arc<ServerAddr>>,
-        nodelay: bool,
         balancer: PingBalancer,
         mode: Mode,
     ) -> Socks5TcpHandler {
         Socks5TcpHandler {
             context,
             udp_bind_addr,
-            nodelay,
             balancer,
             mode,
         }
@@ -163,10 +160,6 @@ impl Socks5TcpHandler {
             }
         };
 
-        if self.nodelay {
-            remote.set_nodelay(true)?;
-        }
-
         establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, &target_addr).await
     }