Skip to content

Commit

Permalink
allow customizing keep-alive seconds on TCP sockets
Browse files Browse the repository at this point in the history
- ref #490
  • Loading branch information
zonyitoo committed Jun 5, 2021
1 parent f1d312f commit df87125
Show file tree
Hide file tree
Showing 16 changed files with 109 additions and 133 deletions.
3 changes: 3 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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,

Expand Down
16 changes: 16 additions & 0 deletions crates/shadowsocks-service/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Expand Down Expand Up @@ -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>,
Expand Down Expand Up @@ -883,6 +888,8 @@ impl Config {

no_delay: false,
fast_open: false,
keep_alive: None,

#[cfg(all(unix, not(target_os = "android")))]
nofile: None,

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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")]
Expand Down
12 changes: 3 additions & 9 deletions crates/shadowsocks-service/src/local/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down Expand Up @@ -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());
}
Expand All @@ -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());
Expand All @@ -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);

Expand Down
16 changes: 1 addition & 15 deletions crates/shadowsocks-service/src/local/redir/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand All @@ -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(),
}
Expand All @@ -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;
Expand Down Expand Up @@ -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<()> {
Expand Down
23 changes: 2 additions & 21 deletions crates/shadowsocks-service/src/local/redir/tcprelay/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand All @@ -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.
Expand All @@ -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?,
Expand Down Expand Up @@ -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);
}
});
Expand Down
21 changes: 3 additions & 18 deletions crates/shadowsocks-service/src/local/socks/server/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -48,7 +47,6 @@ impl Socks {
udp_expiry_duration: None,
udp_capacity: None,
udp_bind_addr: None,
nodelay: false,
}
}

Expand All @@ -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();
Expand Down Expand Up @@ -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;

Expand All @@ -150,7 +138,6 @@ impl Socks {
balancer,
peer_addr,
mode,
nodelay,
));
}
}
Expand All @@ -163,7 +150,6 @@ impl Socks {
balancer: PingBalancer,
peer_addr: SocketAddr,
mode: Mode,
nodelay: bool,
) -> io::Result<()> {
use std::io::ErrorKind;

Expand All @@ -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
}

Expand All @@ -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
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand Down Expand Up @@ -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() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ use crate::{
pub struct Socks5TcpHandler {
context: Arc<ServiceContext>,
udp_bind_addr: Option<Arc<ServerAddr>>,
nodelay: bool,
balancer: PingBalancer,
mode: Mode,
}
Expand All @@ -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,
}
Expand Down Expand Up @@ -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
}

Expand Down
Loading

0 comments on commit df87125

Please sign in to comment.