From f6a9ee45f252fa756f4961596b4bde5a12d02a7a Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Sun, 22 Dec 2019 17:38:29 +0800 Subject: [PATCH] [#179] UDP association expire time update in both direction, removed UDP association limit count --- src/config.rs | 101 +++++++++++++---------------- src/relay/udprelay/mod.rs | 5 ++ src/relay/udprelay/server.rs | 41 ++++++++++-- src/relay/udprelay/socks5_local.rs | 41 ++++++++++-- src/relay/udprelay/tunnel_local.rs | 41 ++++++++++-- 5 files changed, 153 insertions(+), 76 deletions(-) diff --git a/src/config.rs b/src/config.rs index c6312049..42df0315 100644 --- a/src/config.rs +++ b/src/config.rs @@ -115,8 +115,6 @@ struct SSServerExtConfig { plugin_opts: Option, #[serde(skip_serializing_if = "Option::is_none")] timeout: Option, - #[serde(skip_serializing_if = "Option::is_none")] - udp_timeout: Option, } /// Server address @@ -203,8 +201,6 @@ pub struct ServerConfig { enc_key: Bytes, /// Plugin config plugin: Option, - /// UDP timeout - udp_timeout: Option, /// Plugin address plugin_addr: Option, } @@ -226,7 +222,6 @@ impl ServerConfig { timeout, enc_key, plugin, - udp_timeout: None, plugin_addr: None, } } @@ -283,11 +278,6 @@ impl ServerConfig { self.plugin.as_ref() } - /// Get UDP timeout - pub fn udp_timeout(&self) -> &Option { - &self.udp_timeout - } - /// Set plugin address pub fn set_plugin_addr(&mut self, a: ServerAddr) { self.plugin_addr = Some(a); @@ -545,6 +535,7 @@ pub struct Config { pub no_delay: bool, pub manager_address: Option, pub config_type: ConfigType, + pub udp_timeout: Option, } /// Configuration parsing error kind @@ -605,6 +596,7 @@ impl Config { no_delay: false, manager_address: None, config_type, + udp_timeout: None, } } @@ -681,11 +673,7 @@ impl Config { }; let timeout = config.timeout.map(Duration::from_secs); - let udp_timeout = config.udp_timeout.map(Duration::from_secs); - - let mut nsvr = ServerConfig::new(addr, pwd, method, timeout, plugin); - - nsvr.udp_timeout = udp_timeout; + let nsvr = ServerConfig::new(addr, pwd, method, timeout, plugin); nconfig.server.push(nsvr); } @@ -732,11 +720,7 @@ impl Config { }; let timeout = svr.timeout.map(Duration::from_secs); - let udp_timeout = config.udp_timeout.map(Duration::from_secs); - - let mut nsvr = ServerConfig::new(addr, svr.password, method, timeout, plugin); - - nsvr.udp_timeout = udp_timeout; + let nsvr = ServerConfig::new(addr, svr.password, method, timeout, plugin); nconfig.server.push(nsvr); } @@ -779,6 +763,9 @@ impl Config { nconfig.no_delay = b; } + // UDP + nconfig.udp_timeout = config.udp_timeout.map(Duration::from_secs); + Ok(nconfig) } @@ -852,43 +839,45 @@ impl fmt::Display for Config { // Servers // For 1 servers, uses standard configure format - if self.server.len() == 1 { - let svr = &self.server[0]; + match self.server.len() { + 0 => {} + 1 => { + let svr = &self.server[0]; - jconf.server = Some(match *svr.addr() { - ServerAddr::SocketAddr(ref sa) => sa.ip().to_string(), - ServerAddr::DomainName(ref dm, ..) => dm.to_string(), - }); - jconf.server_port = Some(match *svr.addr() { - ServerAddr::SocketAddr(ref sa) => sa.port(), - ServerAddr::DomainName(.., port) => port, - }); - jconf.method = Some(svr.method().to_string()); - jconf.password = Some(svr.password().to_string()); - jconf.plugin = svr.plugin().map(|p| p.plugin.to_string()); - jconf.plugin_opts = svr.plugin().and_then(|p| p.plugin_opt.clone()); - jconf.timeout = svr.timeout().map(|t| t.as_secs()); - jconf.udp_timeout = svr.udp_timeout().map(|t| t.as_secs()); - } else if self.server.len() > 1 { - let mut vsvr = Vec::new(); - - for svr in &self.server { - vsvr.push(SSServerExtConfig { - address: match *svr.addr() { - ServerAddr::SocketAddr(ref sa) => sa.ip().to_string(), - ServerAddr::DomainName(ref dm, ..) => dm.to_string(), - }, - port: match *svr.addr() { - ServerAddr::SocketAddr(ref sa) => sa.port(), - ServerAddr::DomainName(.., port) => port, - }, - password: svr.password().to_string(), - method: svr.method().to_string(), - plugin: svr.plugin().map(|p| p.plugin.to_string()), - plugin_opts: svr.plugin().and_then(|p| p.plugin_opt.clone()), - timeout: svr.timeout().map(|t| t.as_secs()), - udp_timeout: svr.udp_timeout().map(|t| t.as_secs()), + jconf.server = Some(match *svr.addr() { + ServerAddr::SocketAddr(ref sa) => sa.ip().to_string(), + ServerAddr::DomainName(ref dm, ..) => dm.to_string(), }); + jconf.server_port = Some(match *svr.addr() { + ServerAddr::SocketAddr(ref sa) => sa.port(), + ServerAddr::DomainName(.., port) => port, + }); + jconf.method = Some(svr.method().to_string()); + jconf.password = Some(svr.password().to_string()); + jconf.plugin = svr.plugin().map(|p| p.plugin.to_string()); + jconf.plugin_opts = svr.plugin().and_then(|p| p.plugin_opt.clone()); + jconf.timeout = svr.timeout().map(|t| t.as_secs()); + } + _ => { + let mut vsvr = Vec::new(); + + for svr in &self.server { + vsvr.push(SSServerExtConfig { + address: match *svr.addr() { + ServerAddr::SocketAddr(ref sa) => sa.ip().to_string(), + ServerAddr::DomainName(ref dm, ..) => dm.to_string(), + }, + port: match *svr.addr() { + ServerAddr::SocketAddr(ref sa) => sa.port(), + ServerAddr::DomainName(.., port) => port, + }, + password: svr.password().to_string(), + method: svr.method().to_string(), + plugin: svr.plugin().map(|p| p.plugin.to_string()), + plugin_opts: svr.plugin().and_then(|p| p.plugin_opt.clone()), + timeout: svr.timeout().map(|t| t.as_secs()), + }); + } } } @@ -910,6 +899,8 @@ impl fmt::Display for Config { jconf.dns = Some(dns.to_string()); } + jconf.udp_timeout = self.udp_timeout.map(|t| t.as_secs()); + write!(f, "{}", json5::to_string(&jconf).unwrap()) } } diff --git a/src/relay/udprelay/mod.rs b/src/relay/udprelay/mod.rs index 64559c06..0f1c83b3 100644 --- a/src/relay/udprelay/mod.rs +++ b/src/relay/udprelay/mod.rs @@ -37,6 +37,8 @@ //! | Fixed | Variable | //! +-------+--------------+ +use std::time::Duration; + pub mod local; pub mod server; pub(crate) mod socks5_local; @@ -50,3 +52,6 @@ mod crypto_io; /// Socks5 UDP ASSOCIATE packet size. The only thing I can find is /// [here](http://support.microsoft.com/kb/822061/)* pub const MAXIMUM_UDP_PAYLOAD_SIZE: usize = 65536; + +/// Default association expire time +pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5 * 60); diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs index f8f63427..748addab 100644 --- a/src/relay/udprelay/server.rs +++ b/src/relay/udprelay/server.rs @@ -20,7 +20,7 @@ use tokio::{ udp::{RecvHalf, SendHalf}, UdpSocket, }, - sync::mpsc, + sync::{mpsc, Mutex}, time, }; @@ -32,11 +32,10 @@ use crate::{ use super::{ crypto_io::{decrypt_payload, encrypt_payload}, + DEFAULT_TIMEOUT, MAXIMUM_UDP_PAYLOAD_SIZE, }; -const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30); - // Represent a UDP association struct UdpAssociation { tx: mpsc::Sender>, @@ -74,7 +73,7 @@ impl UdpAssociation { // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); - let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); + let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); // local -> remote let c_svr_cfg = svr_cfg.clone(); @@ -260,10 +259,33 @@ async fn listen(context: SharedContext, svr_cfg: Arc) -> io::Resul let listener = UdpSocket::bind(&listen_addr).await?; let (mut r, mut w) = listener.split(); + // NOTE: Associations are only eliminated by expire time + // So it may exhaust all available file descriptors + let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); + let assoc_map = Arc::new(Mutex::new(LruCache::with_expiry_duration(timeout))); + let assoc_map_cloned = assoc_map.clone(); + // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024); tokio::spawn(async move { + let assoc_map = assoc_map_cloned; + while let Some((src, pkt)) = rx.recv().await { + let cache_key = src.to_string(); + { + let mut amap = assoc_map.lock().await; + + // Check or update expire time + if amap.get(&cache_key).is_none() { + debug!( + "UDP association {} <-> ... is already expired, throwing away packet {} bytes", + src, + pkt.len() + ); + continue; + } + } + if let Err(err) = w.send_to(&pkt, &src).await { error!("UDP packet send failed, err: {:?}", err); break; @@ -273,9 +295,6 @@ async fn listen(context: SharedContext, svr_cfg: Arc) -> io::Resul // FIXME: How to stop the outer listener Future? }); - let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); - let mut assoc_map = LruCache::with_expiry_duration_and_capacity(timeout, 1024 /* Conservative, ulimit */); - let mut pkt_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; loop { @@ -284,6 +303,7 @@ async fn listen(context: SharedContext, svr_cfg: Arc) -> io::Resul Err(..) => { // Cleanup expired association // Do not consume this iterator, it will updates expire time of items that traversed + let mut assoc_map = assoc_map.lock().await; let _ = assoc_map.iter(); continue; } @@ -294,8 +314,15 @@ async fn listen(context: SharedContext, svr_cfg: Arc) -> io::Resul trace!("received UDP packet from {}, length {} bytes", src, recv_len); + if recv_len == 0 { + // Interestingly some clients may send a empty packet to server + continue; + } + // Check or (re)create an association loop { + let mut assoc_map = assoc_map.lock().await; + let retry = { let assoc = match assoc_map.entry(src.to_string()) { Entry::Occupied(oc) => oc.into_mut(), diff --git a/src/relay/udprelay/socks5_local.rs b/src/relay/udprelay/socks5_local.rs index 33d45cce..db759c1c 100644 --- a/src/relay/udprelay/socks5_local.rs +++ b/src/relay/udprelay/socks5_local.rs @@ -19,7 +19,7 @@ use tokio::{ udp::{RecvHalf, SendHalf}, UdpSocket, }, - sync::mpsc, + sync::{mpsc, Mutex}, time, }; @@ -35,11 +35,10 @@ use crate::{ use super::{ crypto_io::{decrypt_payload, encrypt_payload}, + DEFAULT_TIMEOUT, MAXIMUM_UDP_PAYLOAD_SIZE, }; -const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30); - async fn parse_packet(pkt: &[u8]) -> io::Result<(Address, Vec)> { // PKT = UdpAssociateHeader + PAYLOAD let mut cur = Cursor::new(pkt); @@ -98,7 +97,7 @@ impl UdpAssociation { // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); - let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); + let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); // local -> remote let c_svr_cfg = svr_cfg.clone(); @@ -264,12 +263,35 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { let (mut r, mut w) = l.split(); + // NOTE: Associations are only eliminated by expire time + // So it may exhaust all available file descriptors + let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); + let assoc_map = Arc::new(Mutex::new(LruCache::with_expiry_duration(timeout))); + let assoc_map_cloned = assoc_map.clone(); + let mut pkt_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::<(SocketAddr, Vec)>(1024); tokio::spawn(async move { + let assoc_map = assoc_map_cloned; + while let Some((src, pkt)) = rx.recv().await { + let cache_key = src.to_string(); + { + let mut amap = assoc_map.lock().await; + + // Check or update expire time + if amap.get(&cache_key).is_none() { + debug!( + "UDP association {} <-> ... is already expired, throwing away packet {} bytes", + src, + pkt.len() + ); + continue; + } + } + if let Err(err) = w.send_to(&pkt, &src).await { error!("UDP packet send failed, err: {:?}", err); break; @@ -279,15 +301,13 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { // FIXME: How to stop the outer listener Future? }); - let mut assoc_map = - LruCache::with_expiry_duration_and_capacity(DEFAULT_TIMEOUT, 1024 /* Conservative, ulimit */); - loop { let (recv_len, src) = match time::timeout(DEFAULT_TIMEOUT, r.recv_from(&mut pkt_buf)).await { Ok(r) => r?, Err(..) => { // Cleanup expired association // Do not consume this iterator, it will updates expire time of items that traversed + let mut assoc_map = assoc_map.lock().await; let _ = assoc_map.iter(); continue; } @@ -299,11 +319,18 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { trace!("received UDP packet from {}, length {} bytes", src, recv_len); + if recv_len == 0 { + // Interestingly some clients may send a empty packet to server + continue; + } + // Pick a server let svr_cfg = balancer.pick_server(); // Check or (re)create an association loop { + let mut assoc_map = assoc_map.lock().await; + let retry = { let assoc = match assoc_map.entry(src.to_string()) { Entry::Occupied(oc) => oc.into_mut(), diff --git a/src/relay/udprelay/tunnel_local.rs b/src/relay/udprelay/tunnel_local.rs index f428db21..4c1234b8 100644 --- a/src/relay/udprelay/tunnel_local.rs +++ b/src/relay/udprelay/tunnel_local.rs @@ -19,7 +19,7 @@ use tokio::{ udp::{RecvHalf, SendHalf}, UdpSocket, }, - sync::mpsc, + sync::{mpsc, Mutex}, time, }; @@ -35,11 +35,10 @@ use crate::{ use super::{ crypto_io::{decrypt_payload, encrypt_payload}, + DEFAULT_TIMEOUT, MAXIMUM_UDP_PAYLOAD_SIZE, }; -const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30); - // Represent a UDP association struct UdpAssociation { tx: mpsc::Sender>, @@ -77,7 +76,7 @@ impl UdpAssociation { // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); - let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); + let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); // local -> remote let c_svr_cfg = svr_cfg.clone(); @@ -236,12 +235,35 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { let (mut r, mut w) = l.split(); + // NOTE: Associations are only eliminated by expire time + // So it may exhaust all available file descriptors + let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); + let assoc_map = Arc::new(Mutex::new(LruCache::with_expiry_duration(timeout))); + let assoc_map_cloned = assoc_map.clone(); + let mut pkt_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::<(SocketAddr, Vec)>(1024); tokio::spawn(async move { + let assoc_map = assoc_map_cloned; + while let Some((src, pkt)) = rx.recv().await { + let cache_key = src.to_string(); + { + let mut amap = assoc_map.lock().await; + + // Check or update expire time + if amap.get(&cache_key).is_none() { + debug!( + "UDP association {} <-> ... is already expired, throwing away packet {} bytes", + src, + pkt.len() + ); + continue; + } + } + if let Err(err) = w.send_to(&pkt, &src).await { error!("UDP packet send failed, err: {:?}", err); break; @@ -251,15 +273,13 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { // FIXME: How to stop the outer listener Future? }); - let mut assoc_map = - LruCache::with_expiry_duration_and_capacity(DEFAULT_TIMEOUT, 1024 /* Conservative, ulimit */); - loop { let (recv_len, src) = match time::timeout(DEFAULT_TIMEOUT, r.recv_from(&mut pkt_buf)).await { Ok(r) => r?, Err(..) => { // Cleanup expired association // Do not consume this iterator, it will updates expire time of items that traversed + let mut assoc_map = assoc_map.lock().await; let _ = assoc_map.iter(); continue; } @@ -271,11 +291,18 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { trace!("received UDP packet from {}, length {} bytes", src, recv_len); + if recv_len == 0 { + // Interestingly some clients may send a empty packet to server + continue; + } + // Pick a server let svr_cfg = balancer.pick_server(); // Check or (re)create an association loop { + let mut assoc_map = assoc_map.lock().await; + let retry = { let assoc = match assoc_map.entry(src.to_string()) { Entry::Occupied(oc) => oc.into_mut(),