starred/shadowsocks-rust
1
0
Fork 0
mirror of https://github.com/shadowsocks/shadowsocks-rust.git synced 2026-11-06 16:14:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

Encapsulate ProxyAssociation management logic

Browse Source
This commit is contained in:
zonyitoo
2020-05-24 00:25:30 +08:00
parent 520018d30a
commit b3052199df
4 changed files with 169 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/relay/udprelay/association.rs
+86 -2
View File
@@ -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<K> {
map: Arc<Mutex<LruCache<K, ProxyAssociation>>>,
watcher: AbortHandle,
}
impl<K> Drop for ProxyAssociationManager<K> {
fn drop(&mut self) {
self.watcher.abort()
}
}
impl<K> ProxyAssociationManager<K>
where
K: Ord + Clone + Send + 'static,
{
/// Create a new ProxyAssociationManager based on Config
pub fn new(config: &Config) -> ProxyAssociationManager<K> {
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<F>(&self, key: K, target: Address, pkt: Vec<u8>, create_fut: F) -> io::Result<()>
where
F: Future<Output = io::Result<ProxyAssociation>>,
{
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(())
}
}
src/relay/udprelay/redir_local.rs
+28 -65
View File
@@ -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<String, ProxyAssociation>;
type SharedAssocMap = Arc<Mutex<AssocMap>>;
struct ProxyHandler {
ty: RedirType,
src_addr: SocketAddr,
cache_key: String,
assoc_map: SharedAssocMap,
assoc_map: ProxyAssociationManager<String>,
}
impl ProxyHandler {
@@ -39,7 +33,7 @@ impl ProxyHandler {
ty: RedirType,
src_addr: SocketAddr,
cache_key: String,
assoc_map: SharedAssocMap,
assoc_map: ProxyAssociationManager<String>,
) -> io::Result<ProxyHandler> {
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);
}
}
}
src/relay/udprelay/socks5_local.rs
+27 -66
View File
@@ -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<u8>)>(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);
}
}
}
src/relay/udprelay/tunnel_local.rs
+28 -66
View File
@@ -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<u8>)>(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);
}
}
}