diff --git a/crates/shadowsocks-service/src/local/redir/udprelay/mod.rs b/crates/shadowsocks-service/src/local/redir/udprelay/mod.rs index 151229b6..f7da3b30 100644 --- a/crates/shadowsocks-service/src/local/redir/udprelay/mod.rs +++ b/crates/shadowsocks-service/src/local/redir/udprelay/mod.rs @@ -39,7 +39,7 @@ mod sys; pub struct UdpRedir { context: Arc, redir_ty: RedirType, - assoc_map: Arc>>>, + assoc_map: Arc>>, cleanup_abortable: AbortHandle, } @@ -150,50 +150,33 @@ impl UdpRedir { data: &[u8], ) -> io::Result<()> { // Check or (re)create an association - let assoc = match self.assoc_map.lock().await.entry(peer_addr) { - Entry::Occupied(occ) => occ.into_mut().clone(), + match self.assoc_map.lock().await.entry(peer_addr) { + Entry::Occupied(occ) => { + let assoc = occ.into_mut(); + assoc.try_send((target_addr, Bytes::copy_from_slice(data))) + } Entry::Vacant(vac) => { - let assoc = UdpAssociation::new( + let assoc = vac.insert(UdpAssociation::new( self.context.clone(), self.redir_ty, peer_addr, self.assoc_map.clone(), balancer.clone(), - ); + )); trace!("created udp association for {}", peer_addr); - vac.insert(assoc.clone()); - assoc + assoc.try_send((target_addr, Bytes::copy_from_slice(data))) } - }; - - if let Err(..) = assoc.sender.try_send((target_addr, Bytes::copy_from_slice(data))) { - let err = io::Error::new(ErrorKind::Other, "udp relay channel full"); - return Err(err); } - - Ok(()) } } struct UdpAssociation { - context: Arc, - redir_ty: RedirType, - peer_addr: SocketAddr, - sender: mpsc::Sender<(SocketAddr, Bytes)>, - bypassed_ipv4_socket: SpinMutex>>, - bypassed_ipv6_socket: SpinMutex>>, - proxied_socket: SpinMutex>>, - assoc_map: Arc>>>, - balancer: PingBalancer, - abortables: SpinMutex>, + assoc: Arc, } impl Drop for UdpAssociation { fn drop(&mut self) { - for ab in self.abortables.lock().iter() { - ab.abort(); - } - trace!("udp association for {} is closed", self.peer_addr); + self.assoc.abortables.lock().abort_all(); } } @@ -202,15 +185,87 @@ impl UdpAssociation { context: Arc, redir_ty: RedirType, peer_addr: SocketAddr, - assoc_map: Arc>>>, + assoc_map: Arc>>, balancer: PingBalancer, - ) -> Arc { + ) -> UdpAssociation { + UdpAssociation { + assoc: UdpAssociationContext::new(context, redir_ty, peer_addr, assoc_map, balancer), + } + } + + fn try_send(&self, data: (SocketAddr, Bytes)) -> io::Result<()> { + if let Err(..) = self.assoc.sender.try_send(data) { + let err = io::Error::new(ErrorKind::Other, "udp relay channel full"); + return Err(err); + } + Ok(()) + } +} + +struct UdpAssociationTaskHandle { + abortables: Vec, + finished: bool, +} + +impl UdpAssociationTaskHandle { + fn new() -> UdpAssociationTaskHandle { + UdpAssociationTaskHandle { + abortables: Vec::new(), + finished: false, + } + } + + fn push_abortable(&mut self, abortable: AbortHandle) { + if self.finished { + // Association is already finished. Kill it immediately. + abortable.abort(); + } else { + self.abortables.push(abortable); + } + } + + fn abort_all(&mut self) { + self.finished = true; + for abortable in &self.abortables { + abortable.abort(); + } + self.abortables.clear(); + } +} + +struct UdpAssociationContext { + context: Arc, + redir_ty: RedirType, + peer_addr: SocketAddr, + sender: mpsc::Sender<(SocketAddr, Bytes)>, + bypassed_ipv4_socket: SpinMutex>>, + bypassed_ipv6_socket: SpinMutex>>, + proxied_socket: SpinMutex>>, + assoc_map: Arc>>, + balancer: PingBalancer, + abortables: SpinMutex, +} + +impl Drop for UdpAssociationContext { + fn drop(&mut self) { + trace!("udp association for {} is closed", self.peer_addr); + } +} + +impl UdpAssociationContext { + fn new( + context: Arc, + redir_ty: RedirType, + peer_addr: SocketAddr, + assoc_map: Arc>>, + balancer: PingBalancer, + ) -> Arc { // Pending packets 1024 should be good enough for a server. // If there are plenty of packets stuck in the channel, dropping exccess packets is a good way to protect the server from // being OOM. let (sender, receiver) = mpsc::channel(1024); - let assoc = Arc::new(UdpAssociation { + let assoc = Arc::new(UdpAssociationContext { context, redir_ty, peer_addr, @@ -220,7 +275,7 @@ impl UdpAssociation { proxied_socket: SpinMutex::new(None), assoc_map, balancer, - abortables: SpinMutex::new(Vec::new()), + abortables: SpinMutex::new(UdpAssociationTaskHandle::new()), }); let (l2r_task, l2r_abortable) = { @@ -229,7 +284,7 @@ impl UdpAssociation { }; tokio::spawn(l2r_task); - assoc.abortables.lock().push(l2r_abortable); + assoc.abortables.lock().push_abortable(l2r_abortable); assoc } @@ -299,7 +354,7 @@ impl UdpAssociation { self.context.connect_opts_ref() ); *bypassed_socket = Some(socket); - self.abortables.lock().push(r2l_abortable); + self.abortables.lock().push_abortable(r2l_abortable); } let socket = bypassed_socket.as_ref().unwrap(); @@ -340,7 +395,7 @@ impl UdpAssociation { self.context.connect_opts_ref() ); *bypassed_socket = Some(socket); - self.abortables.lock().push(r2l_abortable); + self.abortables.lock().push_abortable(r2l_abortable); } let socket = bypassed_socket.as_ref().unwrap(); @@ -388,7 +443,7 @@ impl UdpAssociation { self.context.connect_opts_ref() ); *proxied_socket = Some(socket); - self.abortables.lock().push(r2l_abortable); + self.abortables.lock().push_abortable(r2l_abortable); } let socket = proxied_socket.as_ref().unwrap(); 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 a30f7e60..cd97b554 100644 --- a/crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs +++ b/crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs @@ -36,7 +36,7 @@ use crate::{ pub struct Socks5UdpServer { context: Arc, - assoc_map: Arc>>>, + assoc_map: Arc>>, cleanup_abortable: AbortHandle, } @@ -155,50 +155,33 @@ impl Socks5UdpServer { balancer: &PingBalancer, data: &[u8], ) -> io::Result<()> { - let assoc = match self.assoc_map.lock().await.entry(peer_addr) { - Entry::Occupied(occ) => occ.into_mut().clone(), + match self.assoc_map.lock().await.entry(peer_addr) { + Entry::Occupied(occ) => { + let assoc = occ.into_mut(); + assoc.try_send((target_addr, Bytes::copy_from_slice(data))) + } Entry::Vacant(vac) => { - let assoc = UdpAssociation::new( + let assoc = vac.insert(UdpAssociation::new( self.context.clone(), listener.clone(), peer_addr, self.assoc_map.clone(), balancer.clone(), - ); + )); trace!("created udp association for {}", peer_addr); - vac.insert(assoc.clone()); - assoc + assoc.try_send((target_addr, Bytes::copy_from_slice(data))) } - }; - - if let Err(..) = assoc.sender.try_send((target_addr, Bytes::copy_from_slice(data))) { - let err = io::Error::new(ErrorKind::Other, "udp relay channel full"); - return Err(err); } - - Ok(()) } } struct UdpAssociation { - context: Arc, - inbound: Arc, - peer_addr: SocketAddr, - sender: mpsc::Sender<(Address, Bytes)>, - bypassed_ipv4_socket: SpinMutex>>, - bypassed_ipv6_socket: SpinMutex>>, - proxied_socket: SpinMutex>>, - assoc_map: Arc>>>, - balancer: PingBalancer, - abortables: SpinMutex>, + assoc: Arc, } impl Drop for UdpAssociation { fn drop(&mut self) { - for ab in self.abortables.lock().iter() { - ab.abort(); - } - trace!("udp association for {} is closed", self.peer_addr); + self.assoc.abortables.lock().abort_all(); } } @@ -207,15 +190,87 @@ impl UdpAssociation { context: Arc, inbound: Arc, peer_addr: SocketAddr, - assoc_map: Arc>>>, + assoc_map: Arc>>, balancer: PingBalancer, - ) -> Arc { + ) -> UdpAssociation { + UdpAssociation { + assoc: UdpAssociationContext::new(context, inbound, peer_addr, assoc_map, balancer), + } + } + + fn try_send(&self, data: (Address, Bytes)) -> io::Result<()> { + if let Err(..) = self.assoc.sender.try_send(data) { + let err = io::Error::new(ErrorKind::Other, "udp relay channel full"); + return Err(err); + } + Ok(()) + } +} + +struct UdpAssociationTaskHandle { + abortables: Vec, + finished: bool, +} + +impl UdpAssociationTaskHandle { + fn new() -> UdpAssociationTaskHandle { + UdpAssociationTaskHandle { + abortables: Vec::new(), + finished: false, + } + } + + fn push_abortable(&mut self, abortable: AbortHandle) { + if self.finished { + // Association is already finished. Kill it immediately. + abortable.abort(); + } else { + self.abortables.push(abortable); + } + } + + fn abort_all(&mut self) { + self.finished = true; + for abortable in &self.abortables { + abortable.abort(); + } + self.abortables.clear(); + } +} + +struct UdpAssociationContext { + context: Arc, + inbound: Arc, + peer_addr: SocketAddr, + sender: mpsc::Sender<(Address, Bytes)>, + bypassed_ipv4_socket: SpinMutex>>, + bypassed_ipv6_socket: SpinMutex>>, + proxied_socket: SpinMutex>>, + assoc_map: Arc>>, + balancer: PingBalancer, + abortables: SpinMutex, +} + +impl Drop for UdpAssociationContext { + fn drop(&mut self) { + trace!("udp association for {} is closed", self.peer_addr); + } +} + +impl UdpAssociationContext { + fn new( + context: Arc, + inbound: Arc, + peer_addr: SocketAddr, + assoc_map: Arc>>, + balancer: PingBalancer, + ) -> Arc { // Pending packets 1024 should be good enough for a server. // If there are plenty of packets stuck in the channel, dropping exccess packets is a good way to protect the server from // being OOM. let (sender, receiver) = mpsc::channel(1024); - let assoc = Arc::new(UdpAssociation { + let assoc = Arc::new(UdpAssociationContext { context, inbound, peer_addr, @@ -225,7 +280,7 @@ impl UdpAssociation { proxied_socket: SpinMutex::new(None), assoc_map, balancer, - abortables: SpinMutex::new(Vec::new()), + abortables: SpinMutex::new(UdpAssociationTaskHandle::new()), }); let (l2r_task, l2r_abortable) = { @@ -234,7 +289,7 @@ impl UdpAssociation { }; tokio::spawn(l2r_task); - assoc.abortables.lock().push(l2r_abortable); + assoc.abortables.lock().push_abortable(l2r_abortable); assoc } @@ -315,7 +370,7 @@ impl UdpAssociation { self.context.connect_opts_ref() ); *bypassed_socket = Some(socket); - self.abortables.lock().push(r2l_abortable); + self.abortables.lock().push_abortable(r2l_abortable); } let socket = bypassed_socket.as_ref().unwrap(); @@ -356,7 +411,7 @@ impl UdpAssociation { self.context.connect_opts_ref() ); *bypassed_socket = Some(socket); - self.abortables.lock().push(r2l_abortable); + self.abortables.lock().push_abortable(r2l_abortable); } let socket = bypassed_socket.as_ref().unwrap(); @@ -404,7 +459,7 @@ impl UdpAssociation { self.context.connect_opts_ref() ); *proxied_socket = Some(socket); - self.abortables.lock().push(r2l_abortable); + self.abortables.lock().push_abortable(r2l_abortable); } let socket = proxied_socket.as_ref().unwrap(); diff --git a/crates/shadowsocks-service/src/server/udprelay.rs b/crates/shadowsocks-service/src/server/udprelay.rs index 37e61b0f..82a8e19e 100644 --- a/crates/shadowsocks-service/src/server/udprelay.rs +++ b/crates/shadowsocks-service/src/server/udprelay.rs @@ -28,7 +28,7 @@ use super::context::ServiceContext; pub struct UdpServer { context: Arc, - assoc_map: Arc>>>, + assoc_map: Arc>>, cleanup_abortable: AbortHandle, } @@ -112,47 +112,32 @@ impl UdpServer { target_addr: Address, data: &[u8], ) -> io::Result<()> { - let assoc = match self.assoc_map.lock().await.entry(peer_addr) { - Entry::Occupied(occ) => occ.into_mut().clone(), + match self.assoc_map.lock().await.entry(peer_addr) { + Entry::Occupied(occ) => { + let assoc = occ.into_mut(); + assoc.try_send((target_addr, Bytes::copy_from_slice(data))) + } Entry::Vacant(vac) => { - let assoc = UdpAssociation::new( + let assoc = vac.insert(UdpAssociation::new( self.context.clone(), listener.clone(), peer_addr, self.assoc_map.clone(), - ); - vac.insert(assoc.clone()); - assoc + )); + trace!("created udp association for {}", peer_addr); + assoc.try_send((target_addr, Bytes::copy_from_slice(data))) } - }; - - if let Err(..) = assoc.sender.try_send((target_addr, Bytes::copy_from_slice(data))) { - let err = io::Error::new(ErrorKind::Other, "udp relay channel full"); - return Err(err); } - - Ok(()) } } struct UdpAssociation { - context: Arc, - inbound: Arc, - peer_addr: SocketAddr, - sender: mpsc::Sender<(Address, Bytes)>, - outbound_ipv4_socket: SpinMutex>>, - outbound_ipv6_socket: SpinMutex>>, - assoc_map: Arc>>>, - abortables: SpinMutex>, - target_cache: Mutex>, + assoc: Arc, } impl Drop for UdpAssociation { fn drop(&mut self) { - for ab in self.abortables.lock().iter() { - ab.abort(); - } - trace!("udp association for {} is closed", self.peer_addr); + self.assoc.abortables.lock().abort_all(); } } @@ -161,14 +146,84 @@ impl UdpAssociation { context: Arc, inbound: Arc, peer_addr: SocketAddr, - assoc_map: Arc>>>, - ) -> Arc { + assoc_map: Arc>>, + ) -> UdpAssociation { + UdpAssociation { + assoc: UdpAssociationContext::new(context, inbound, peer_addr, assoc_map), + } + } + + fn try_send(&self, data: (Address, Bytes)) -> io::Result<()> { + if let Err(..) = self.assoc.sender.try_send(data) { + let err = io::Error::new(ErrorKind::Other, "udp relay channel full"); + return Err(err); + } + Ok(()) + } +} + +struct UdpAssociationTaskHandle { + abortables: Vec, + finished: bool, +} + +impl UdpAssociationTaskHandle { + fn new() -> UdpAssociationTaskHandle { + UdpAssociationTaskHandle { + abortables: Vec::new(), + finished: false, + } + } + + fn push_abortable(&mut self, abortable: AbortHandle) { + if self.finished { + // Association is already finished. Kill it immediately. + abortable.abort(); + } else { + self.abortables.push(abortable); + } + } + + fn abort_all(&mut self) { + self.finished = true; + for abortable in &self.abortables { + abortable.abort(); + } + self.abortables.clear(); + } +} + +struct UdpAssociationContext { + context: Arc, + inbound: Arc, + peer_addr: SocketAddr, + sender: mpsc::Sender<(Address, Bytes)>, + outbound_ipv4_socket: SpinMutex>>, + outbound_ipv6_socket: SpinMutex>>, + assoc_map: Arc>>, + abortables: SpinMutex, + target_cache: Mutex>, +} + +impl Drop for UdpAssociationContext { + fn drop(&mut self) { + trace!("udp association for {} is closed", self.peer_addr); + } +} + +impl UdpAssociationContext { + fn new( + context: Arc, + inbound: Arc, + peer_addr: SocketAddr, + assoc_map: Arc>>, + ) -> Arc { // Pending packets 1024 should be good enough for a server. // If there are plenty of packets stuck in the channel, dropping exccess packets is a good way to protect the server from // being OOM. let (sender, receiver) = mpsc::channel(1024); - let assoc = Arc::new(UdpAssociation { + let assoc = Arc::new(UdpAssociationContext { context, inbound, peer_addr, @@ -176,7 +231,7 @@ impl UdpAssociation { outbound_ipv4_socket: SpinMutex::new(None), outbound_ipv6_socket: SpinMutex::new(None), assoc_map, - abortables: SpinMutex::new(Vec::new()), + abortables: SpinMutex::new(UdpAssociationTaskHandle::new()), // Cache for remembering the original Address of target, // when recv_from a SocketAddr, we have to know whch Address that client was originally requested. // @@ -190,7 +245,7 @@ impl UdpAssociation { }; tokio::spawn(l2r_task); - assoc.abortables.lock().push(l2r_abortable); + assoc.abortables.lock().push_abortable(l2r_abortable); assoc } @@ -261,7 +316,7 @@ impl UdpAssociation { self.context.connect_opts_ref() ); *outbound = Some(socket); - self.abortables.lock().push(r2l_abortable); + self.abortables.lock().push_abortable(r2l_abortable); } let socket = outbound.as_ref().unwrap(); @@ -301,7 +356,7 @@ impl UdpAssociation { self.context.connect_opts_ref() ); *outbound = Some(socket); - self.abortables.lock().push(r2l_abortable); + self.abortables.lock().push_abortable(r2l_abortable); } let socket = outbound.as_ref().unwrap();