[#179] UDP association expire time update in both direction, removed UDP association limit count

src/config.rs

@@ -115,8 +115,6 @@ struct SSServerExtConfig {
     plugin_opts: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     timeout: Option<u64>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    udp_timeout: Option<u64>,
 }
 
 /// Server address
@@ -203,8 +201,6 @@ pub struct ServerConfig {
     enc_key: Bytes,
     /// Plugin config
     plugin: Option<PluginConfig>,
-    /// UDP timeout
-    udp_timeout: Option<Duration>,
     /// Plugin address
     plugin_addr: Option<ServerAddr>,
 }
@@ -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<Duration> {
-        &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<ServerAddr>,
     pub config_type: ConfigType,
+    pub udp_timeout: Option<Duration>,
 }
 
 /// 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())
     }
 }

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);

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<Vec<u8>>,
@@ -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<ServerConfig>) -> 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<ServerConfig>) -> 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<ServerConfig>) -> 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<ServerConfig>) -> 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(),

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<u8>)> {
     // 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<u8>)>(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(),

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<Vec<u8>>,
@@ -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<u8>)>(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(),