From 89e9b1400a8b90cba5dc42aba26ebcbfd5cd1cdc Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Mon, 21 Dec 2020 22:17:46 +0800 Subject: [PATCH] active cleanup of expired udp associations, fixed #351 --- .../src/local/redir/udprelay/mod.rs | 32 ++++++++++++++++--- .../src/local/redir/udprelay/sys/unix/bsd.rs | 14 +------- .../local/redir/udprelay/sys/unix/linux.rs | 2 +- .../src/local/socks/server/socks5/udprelay.rs | 32 ++++++++++++++++--- .../src/local/tunnel/udprelay.rs | 32 ++++++++++++++++--- .../src/server/udprelay.rs | 32 ++++++++++++++++--- 6 files changed, 114 insertions(+), 30 deletions(-) diff --git a/crates/shadowsocks-service/src/local/redir/udprelay/mod.rs b/crates/shadowsocks-service/src/local/redir/udprelay/mod.rs index aa8f0874..171674e9 100644 --- a/crates/shadowsocks-service/src/local/redir/udprelay/mod.rs +++ b/crates/shadowsocks-service/src/local/redir/udprelay/mod.rs @@ -40,6 +40,13 @@ pub struct UdpRedir { context: Arc, redir_ty: RedirType, assoc_map: Arc>>>, + cleanup_abortable: AbortHandle, +} + +impl Drop for UdpRedir { + fn drop(&mut self) { + self.cleanup_abortable.abort(); + } } impl UdpRedir { @@ -50,14 +57,31 @@ impl UdpRedir { capacity: Option, ) -> UdpRedir { let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION); + let assoc_map = Arc::new(Mutex::new(match capacity { + Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity), + None => LruCache::with_expiry_duration(time_to_live), + })); + + let cleanup_abortable = { + let assoc_map = assoc_map.clone(); + let (cleanup_task, cleanup_abortable) = future::abortable(async move { + let mut interval = time::interval(time_to_live); + loop { + interval.tick().await; + + // iter() will trigger a cleanup of expired associations + let _ = assoc_map.lock().await.iter(); + } + }); + tokio::spawn(cleanup_task); + cleanup_abortable + }; UdpRedir { context, redir_ty, - assoc_map: Arc::new(Mutex::new(match capacity { - Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity), - None => LruCache::with_expiry_duration(time_to_live), - })), + assoc_map, + cleanup_abortable, } } diff --git a/crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/bsd.rs b/crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/bsd.rs index 71d1318b..58d0c820 100644 --- a/crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/bsd.rs +++ b/crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/bsd.rs @@ -52,19 +52,7 @@ impl UdpRedirSocket { /// Send data to the socket to the given target address pub async fn send_to(&mut self, buf: &[u8], target: SocketAddr) -> io::Result { - // poll_fn(|cx| self.poll_send_to(cx, buf, target)).await - - loop { - let mut write_guard = self.io.writable().await?; - - match self.io.get_ref().send_to(buf, target) { - Err(ref e) if e.kind() == ErrorKind::WouldBlock => { - write_guard.clear_ready(); - // false positive - } - x => return x, - } - } + poll_fn(|cx| self.poll_send_to(cx, buf, target)).await } fn poll_send_to(&self, cx: &mut Context<'_>, buf: &[u8], target: SocketAddr) -> Poll> { diff --git a/crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/linux.rs b/crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/linux.rs index 86eb81f1..506c5c3f 100644 --- a/crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/linux.rs +++ b/crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/linux.rs @@ -8,7 +8,7 @@ use std::{ task::{Context, Poll}, }; -use futures::ready; +use futures::{future::poll_fn, ready}; use socket2::{Domain, Protocol, SockAddr, Socket, Type}; use tokio::io::unix::AsyncFd; diff --git a/crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs b/crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs index 5f366a77..0a1bf3af 100644 --- a/crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs +++ b/crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs @@ -37,6 +37,13 @@ use crate::{ pub struct Socks5UdpServer { context: Arc, assoc_map: Arc>>>, + cleanup_abortable: AbortHandle, +} + +impl Drop for Socks5UdpServer { + fn drop(&mut self) { + self.cleanup_abortable.abort(); + } } impl Socks5UdpServer { @@ -46,13 +53,30 @@ impl Socks5UdpServer { capacity: Option, ) -> Socks5UdpServer { let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION); + let assoc_map = Arc::new(Mutex::new(match capacity { + Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity), + None => LruCache::with_expiry_duration(time_to_live), + })); + + let cleanup_abortable = { + let assoc_map = assoc_map.clone(); + let (cleanup_task, cleanup_abortable) = future::abortable(async move { + let mut interval = time::interval(time_to_live); + loop { + interval.tick().await; + + // iter() will trigger a cleanup of expired associations + let _ = assoc_map.lock().await.iter(); + } + }); + tokio::spawn(cleanup_task); + cleanup_abortable + }; Socks5UdpServer { context, - assoc_map: Arc::new(Mutex::new(match capacity { - Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity), - None => LruCache::with_expiry_duration(time_to_live), - })), + assoc_map, + cleanup_abortable, } } diff --git a/crates/shadowsocks-service/src/local/tunnel/udprelay.rs b/crates/shadowsocks-service/src/local/tunnel/udprelay.rs index 72f0ea79..7a647c16 100644 --- a/crates/shadowsocks-service/src/local/tunnel/udprelay.rs +++ b/crates/shadowsocks-service/src/local/tunnel/udprelay.rs @@ -30,18 +30,42 @@ use crate::{ pub struct UdpTunnel { context: Arc, assoc_map: Arc>>, + cleanup_abortable: AbortHandle, +} + +impl Drop for UdpTunnel { + fn drop(&mut self) { + self.cleanup_abortable.abort(); + } } impl UdpTunnel { pub fn new(context: Arc, time_to_live: Option, capacity: Option) -> UdpTunnel { let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION); + let assoc_map = Arc::new(Mutex::new(match capacity { + Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity), + None => LruCache::with_expiry_duration(time_to_live), + })); + + let cleanup_abortable = { + let assoc_map = assoc_map.clone(); + let (cleanup_task, cleanup_abortable) = future::abortable(async move { + let mut interval = time::interval(time_to_live); + loop { + interval.tick().await; + + // iter() will trigger a cleanup of expired associations + let _ = assoc_map.lock().await.iter(); + } + }); + tokio::spawn(cleanup_task); + cleanup_abortable + }; UdpTunnel { context, - assoc_map: Arc::new(Mutex::new(match capacity { - Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity), - None => LruCache::with_expiry_duration(time_to_live), - })), + assoc_map, + cleanup_abortable, } } diff --git a/crates/shadowsocks-service/src/server/udprelay.rs b/crates/shadowsocks-service/src/server/udprelay.rs index ba9f39e6..ef786ec8 100644 --- a/crates/shadowsocks-service/src/server/udprelay.rs +++ b/crates/shadowsocks-service/src/server/udprelay.rs @@ -29,18 +29,42 @@ use super::context::ServiceContext; pub struct UdpServer { context: Arc, assoc_map: Arc>>>, + cleanup_abortable: AbortHandle, +} + +impl Drop for UdpServer { + fn drop(&mut self) { + self.cleanup_abortable.abort(); + } } impl UdpServer { pub fn new(context: Arc, time_to_live: Option, capacity: Option) -> UdpServer { let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION); + let assoc_map = Arc::new(Mutex::new(match capacity { + Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity), + None => LruCache::with_expiry_duration(time_to_live), + })); + + let cleanup_abortable = { + let assoc_map = assoc_map.clone(); + let (cleanup_task, cleanup_abortable) = future::abortable(async move { + let mut interval = time::interval(time_to_live); + loop { + interval.tick().await; + + // iter() will trigger a cleanup of expired associations + let _ = assoc_map.lock().await.iter(); + } + }); + tokio::spawn(cleanup_task); + cleanup_abortable + }; UdpServer { context, - assoc_map: Arc::new(Mutex::new(match capacity { - Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity), - None => LruCache::with_expiry_duration(time_to_live), - })), + assoc_map, + cleanup_abortable, } }