From b3052199df35992499197eadffaf19ebcece2860 Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Sun, 24 May 2020 00:25:30 +0800 Subject: [PATCH] Encapsulate ProxyAssociation management logic --- src/relay/udprelay/association.rs | 88 +++++++++++++++++++++++++++- src/relay/udprelay/redir_local.rs | 93 +++++++++-------------------- src/relay/udprelay/socks5_local.rs | 93 +++++++++-------------------- src/relay/udprelay/tunnel_local.rs | 94 +++++++++--------------------- 4 files changed, 169 insertions(+), 199 deletions(-) diff --git a/src/relay/udprelay/association.rs b/src/relay/udprelay/association.rs index e4f8c8ce..9bcb67c5 100644 --- a/src/relay/udprelay/association.rs +++ b/src/relay/udprelay/association.rs @@ -5,22 +5,26 @@ #![allow(dead_code)] use std::{ + future::Future, io::{self, Cursor, Read}, net::{IpAddr, Ipv4Addr, SocketAddr}, + sync::Arc, }; use async_trait::async_trait; use bytes::BytesMut; use futures::future::{self, AbortHandle}; use log::{debug, error, warn}; +use lru_time_cache::{Entry, LruCache}; use tokio::{ self, net::udp::{RecvHalf, SendHalf}, - sync::mpsc, + sync::{mpsc, Mutex}, + time, }; use crate::{ - config::{ServerAddr, ServerConfig}, + config::{Config, ServerAddr, ServerConfig}, context::Context, relay::{ loadbalancing::server::{ServerData, SharedServerStatistic}, @@ -31,6 +35,7 @@ use crate::{ use super::{ crypto_io::{decrypt_payload, encrypt_payload}, + DEFAULT_TIMEOUT, MAXIMUM_UDP_PAYLOAD_SIZE, }; @@ -421,3 +426,82 @@ impl ProxyAssociation { Ok((addr, payload)) } } + +#[derive(Clone)] +pub struct ProxyAssociationManager { + map: Arc>>, + watcher: AbortHandle, +} + +impl Drop for ProxyAssociationManager { + fn drop(&mut self) { + self.watcher.abort() + } +} + +impl ProxyAssociationManager +where + K: Ord + Clone + Send + 'static, +{ + /// Create a new ProxyAssociationManager based on Config + pub fn new(config: &Config) -> ProxyAssociationManager { + let timeout = config.udp_timeout.unwrap_or(DEFAULT_TIMEOUT); + + // TODO: Set default capacity by getrlimit #262 + // Associations are only eliminated by expire time by default + // So it may exhaust all available file descriptors + let assoc_map = if let Some(max_assoc) = config.udp_max_associations { + LruCache::with_expiry_duration_and_capacity(timeout, max_assoc) + } else { + LruCache::with_expiry_duration(timeout) + }; + + let map = Arc::new(Mutex::new(assoc_map)); + + // Create a task for releasing timed out association + let map2 = map.clone(); + let (release_task, watcher) = future::abortable(async move { + let mut interval = time::interval(timeout); + loop { + interval.tick().await; + + let mut m = map2.lock().await; + // Cleanup expired association + // Do not consume this iterator, it will updates expire time of items that traversed + let _ = m.iter(); + } + }); + + tokio::spawn(release_task); + + ProxyAssociationManager { map, watcher } + } + + /// Try to reset ProxyAssociation's last used time by key + /// + /// Return true if ProxyAssociation is still exist + pub async fn keep_alive(&self, key: &K) -> bool { + let mut assoc = self.map.lock().await; + assoc.get(key).is_some() + } + + /// Send a packet to target address + /// + /// Create a new association by `create_fut` if association doesn't exist + pub async fn send_packet(&self, key: K, target: Address, pkt: Vec, create_fut: F) -> io::Result<()> + where + F: Future>, + { + let mut assoc_map = self.map.lock().await; + let assoc = match assoc_map.entry(key) { + Entry::Occupied(oc) => oc.into_mut(), + Entry::Vacant(vc) => vc.insert(create_fut.await?), + }; + + // FIXME: Lock is still kept for a mutable reference + // Send to local -> remote task + assoc.send(target, pkt).await; + + Ok(()) + } +} diff --git a/src/relay/udprelay/redir_local.rs b/src/relay/udprelay/redir_local.rs index eab476b8..881e79f0 100644 --- a/src/relay/udprelay/redir_local.rs +++ b/src/relay/udprelay/redir_local.rs @@ -1,11 +1,9 @@ //! UDP relay local server -use std::{io, net::SocketAddr, sync::Arc}; +use std::{io, net::SocketAddr}; use async_trait::async_trait; use log::{error, info, trace}; -use lru_time_cache::{Entry, LruCache}; -use tokio::{self, sync::Mutex, time}; use crate::{ config::RedirType, @@ -18,20 +16,16 @@ use crate::{ }; use super::{ - association::{ProxyAssociation, ProxySend}, + association::{ProxyAssociation, ProxyAssociationManager, ProxySend}, redir::sys::UdpRedirSocket, - DEFAULT_TIMEOUT, MAXIMUM_UDP_PAYLOAD_SIZE, }; -type AssocMap = LruCache; -type SharedAssocMap = Arc>; - struct ProxyHandler { ty: RedirType, src_addr: SocketAddr, cache_key: String, - assoc_map: SharedAssocMap, + assoc_map: ProxyAssociationManager, } impl ProxyHandler { @@ -39,7 +33,7 @@ impl ProxyHandler { ty: RedirType, src_addr: SocketAddr, cache_key: String, - assoc_map: SharedAssocMap, + assoc_map: ProxyAssociationManager, ) -> io::Result { Ok(ProxyHandler { ty, @@ -62,12 +56,7 @@ impl ProxySend for ProxyHandler { local_udp.send_to(&data, &self.src_addr).await?; // Update LRU - { - let mut amap = self.assoc_map.lock().await; - - // Check or update expire time - let _ = amap.get(&self.cache_key); - } + self.assoc_map.keep_alive(&self.cache_key).await; return Ok(()); } @@ -94,29 +83,12 @@ pub async fn run(context: SharedContext) -> io::Result<()> { info!("shadowsocks UDP redirect listening on {}", local_addr); - // NOTE: Associations are only eliminated by expire time by default - // So it may exhaust all available file descriptors - let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); - let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations { - LruCache::with_expiry_duration_and_capacity(timeout, max_assoc) - } else { - LruCache::with_expiry_duration(timeout) - }; - let assoc_map = Arc::new(Mutex::new(assoc_map)); + let assoc_manager = ProxyAssociationManager::new(context.config()); let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; loop { - let (recv_len, src, dst) = match time::timeout(timeout, l.recv_from_redir(&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; - } - }; + let (recv_len, src, dst) = l.recv_from_redir(&mut pkt_buf).await?; // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded. // Copy bytes, because udp_associate runs in another tokio Task @@ -145,41 +117,32 @@ pub async fn run(context: SharedContext) -> io::Result<()> { let is_bypassed = context.check_target_bypassed(&target).await; // Check or (re)create an association - { - // Locks the whole association map - let mut ref_assoc_map = assoc_map.lock().await; + let cache_key = format!("{}-{}", src, dst); + let cache_key_cloned = cache_key.clone(); + let res = assoc_manager + .send_packet(cache_key, target, pkt.to_vec(), async { + // Pick a server + let server = balancer.pick_server(); - let cache_key = format!("{}-{}", src, dst); - - // Get or create an association - let assoc = match ref_assoc_map.entry(cache_key.clone()) { - Entry::Occupied(oc) => oc.into_mut(), - Entry::Vacant(vc) => { - // Pick a server - let server = balancer.pick_server(); - - let sender = match ProxyHandler::new(ty, src, cache_key, assoc_map.clone()) { - Ok(s) => s, - Err(err) => { - error!("create UDP association for {} <-> {}, error: {}", src, dst, err); - continue; - } - }; - - let assoc = if is_bypassed { - ProxyAssociation::associate_bypassed(src, server, sender).await - } else { - ProxyAssociation::associate_proxied(src, server, sender).await + let sender = match ProxyHandler::new(ty, src, cache_key_cloned, assoc_manager.clone()) { + Ok(s) => s, + Err(err) => { + error!("create UDP association for {} <-> {}, error: {}", src, dst, err); + return Err(err); } - .expect("create UDP association"); + }; - vc.insert(assoc) + if is_bypassed { + ProxyAssociation::associate_bypassed(src, server, sender).await + } else { + ProxyAssociation::associate_proxied(src, server, sender).await } - }; + }) + .await; - // FIXME: Lock is still kept for a mutable reference - // Send to local -> remote task - assoc.send(target, pkt.to_vec()).await; + if let Err(err) = res { + error!("failed to create UDP association, {}", err); + return Err(err); } } } diff --git a/src/relay/udprelay/socks5_local.rs b/src/relay/udprelay/socks5_local.rs index a9438736..2593a634 100644 --- a/src/relay/udprelay/socks5_local.rs +++ b/src/relay/udprelay/socks5_local.rs @@ -3,18 +3,12 @@ use std::{ io::{self, Cursor, ErrorKind, Read}, net::SocketAddr, - sync::Arc, }; use async_trait::async_trait; use bytes::{Bytes, BytesMut}; use log::{debug, error, info, trace}; -use lru_time_cache::{Entry, LruCache}; -use tokio::{ - self, - sync::{mpsc, Mutex}, - time, -}; +use tokio::{self, sync::mpsc}; use crate::{ context::SharedContext, @@ -26,8 +20,7 @@ use crate::{ }; use super::{ - association::{ProxyAssociation, ProxySend}, - DEFAULT_TIMEOUT, + association::{ProxyAssociation, ProxyAssociationManager, ProxySend}, MAXIMUM_UDP_PAYLOAD_SIZE, }; @@ -95,38 +88,25 @@ pub async fn run(context: SharedContext) -> io::Result<()> { info!("shadowsocks SOCKS5 UDP listening on {}", local_addr); - // NOTE: Associations are only eliminated by expire time by default - // So it may exhaust all available file descriptors - let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); - let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations { - LruCache::with_expiry_duration_and_capacity(timeout, max_assoc) - } else { - LruCache::with_expiry_duration(timeout) - }; - let assoc_map = Arc::new(Mutex::new(assoc_map)); - let assoc_map_cloned = assoc_map.clone(); + let assoc_manager = ProxyAssociationManager::new(context.config()); + let assoc_manager_cloned = assoc_manager.clone(); let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::<(SocketAddr, Address, Vec)>(1024); tokio::spawn(async move { - let assoc_map = assoc_map_cloned; + let assoc_manager = assoc_manager_cloned; while let Some((src, target, 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 !assoc_manager.keep_alive(&cache_key).await { + debug!( + "UDP association {} <-> ... is already expired, throwing away packet {} bytes", + src, + pkt.len() + ); + continue; } let payload = assemble_packet(target, &pkt); @@ -141,16 +121,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { }); loop { - let (recv_len, src) = match time::timeout(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; - } - }; + let (recv_len, src) = r.recv_from(&mut pkt_buf).await?; // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded. // Copy bytes, because udp_associate runs in another tokio Task @@ -182,33 +153,23 @@ pub async fn run(context: SharedContext) -> io::Result<()> { }; // Check or (re)create an association - { - // Locks the whole association map - let mut assoc_map = assoc_map.lock().await; + let res = assoc_manager + .send_packet(src.to_string(), target, payload, async { + // Pick a server + let server = balancer.pick_server(); - // Get or create an association - let assoc = match assoc_map.entry(src.to_string()) { - Entry::Occupied(oc) => oc.into_mut(), - Entry::Vacant(vc) => { - // Pick a server - let server = balancer.pick_server(); + let sender = ProxyHandler { + src_addr: src, + response_tx: tx.clone(), + }; - let sender = ProxyHandler { - src_addr: src, - response_tx: tx.clone(), - }; + ProxyAssociation::associate_with_acl(src, server, sender).await + }) + .await; - vc.insert( - ProxyAssociation::associate_with_acl(src, server, sender) - .await - .expect("create UDP association"), - ) - } - }; - - // FIXME: Lock is still kept for a mutable reference - // Send to local -> remote task - assoc.send(target, payload).await; + if let Err(err) = res { + error!("failed to create UDP association, {}", err); + return Err(err); } } } diff --git a/src/relay/udprelay/tunnel_local.rs b/src/relay/udprelay/tunnel_local.rs index bca06b93..214e1d63 100644 --- a/src/relay/udprelay/tunnel_local.rs +++ b/src/relay/udprelay/tunnel_local.rs @@ -1,15 +1,10 @@ //! UDP relay local server -use std::{io, net::SocketAddr, sync::Arc}; +use std::{io, net::SocketAddr}; use async_trait::async_trait; use log::{debug, error, info, trace}; -use lru_time_cache::{Entry, LruCache}; -use tokio::{ - self, - sync::{mpsc, Mutex}, - time, -}; +use tokio::{self, sync::mpsc}; use crate::{ context::SharedContext, @@ -21,8 +16,7 @@ use crate::{ }; use super::{ - association::{ProxyAssociation, ProxySend}, - DEFAULT_TIMEOUT, + association::{ProxyAssociation, ProxyAssociationManager, ProxySend}, MAXIMUM_UDP_PAYLOAD_SIZE, }; @@ -61,38 +55,25 @@ pub async fn run(context: SharedContext) -> io::Result<()> { local_addr, forward_target ); - // NOTE: Associations are only eliminated by expire time by default - // So it may exhaust all available file descriptors - let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT); - let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations { - LruCache::with_expiry_duration_and_capacity(timeout, max_assoc) - } else { - LruCache::with_expiry_duration(timeout) - }; - let assoc_map = Arc::new(Mutex::new(assoc_map)); - let assoc_map_cloned = assoc_map.clone(); + let assoc_manager = ProxyAssociationManager::new(context.config()); + let assoc_manager_cloned = assoc_manager.clone(); let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::<(SocketAddr, Address, Vec)>(1024); tokio::spawn(async move { - let assoc_map = assoc_map_cloned; + let assoc_manager = assoc_manager_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 !assoc_manager.keep_alive(&cache_key).await { + debug!( + "UDP association {} <-> ... is already expired, throwing away packet {} bytes", + src, + pkt.len() + ); + continue; } if let Err(err) = w.send_to(&pkt, &src).await { @@ -105,16 +86,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { }); loop { - let (recv_len, src) = match time::timeout(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; - } - }; + let (recv_len, src) = r.recv_from(&mut pkt_buf).await?; // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded. // Copy bytes, because udp_associate runs in another tokio Task @@ -134,33 +106,23 @@ pub async fn run(context: SharedContext) -> io::Result<()> { } // Check or (re)create an association - { - // Locks the whole association map - let mut assoc_map = assoc_map.lock().await; + let res = assoc_manager + .send_packet(src.to_string(), forward_target.clone(), pkt.to_vec(), async { + // Pick a server + let server = balancer.pick_server(); - // Get or create an association - let assoc = match assoc_map.entry(src.to_string()) { - Entry::Occupied(oc) => oc.into_mut(), - Entry::Vacant(vc) => { - // Pick a server - let server = balancer.pick_server(); + let sender = ProxyHandler { + src_addr: src, + response_tx: tx.clone(), + }; - let sender = ProxyHandler { - src_addr: src, - response_tx: tx.clone(), - }; + ProxyAssociation::associate_with_acl(src, server, sender).await + }) + .await; - vc.insert( - ProxyAssociation::associate_with_acl(src, server, sender) - .await - .expect("create UDP association"), - ) - } - }; - - // FIXME: Lock is still kept for a mutable reference - // Send to local -> remote task - assoc.send(forward_target.clone(), pkt.to_vec()).await; + if let Err(err) = res { + error!("failed to create UDP association, {}", err); + return Err(err); } } }