diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs index f93a0039..be6d22b1 100644 --- a/src/relay/udprelay/server.rs +++ b/src/relay/udprelay/server.rs @@ -34,13 +34,23 @@ use super::{ MAXIMUM_UDP_PAYLOAD_SIZE, }; -// Represent a UDP association -struct UdpAssociation { - tx: mpsc::Sender>, - closed: Arc, +struct UdpAssociationState { + closed: AtomicBool, } -impl Drop for UdpAssociation { +impl UdpAssociationState { + fn new() -> UdpAssociationState { + UdpAssociationState { + closed: AtomicBool::new(false), + } + } + + fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } +} + +impl Drop for UdpAssociationState { fn drop(&mut self) { // 1. Drops tx, will close local -> remote task // 2. Drops closed, will close local <- remote task @@ -48,6 +58,18 @@ impl Drop for UdpAssociation { } } +// Represent a UDP association +#[derive(Clone)] +struct UdpAssociation { + // local -> remote Queue + // Drops tx, will close local -> remote task + tx: mpsc::Sender>, + + // local <- remote relay task + // Drops state, will close local <- remote task + state: Arc, +} + impl UdpAssociation { /// Create an association with addr async fn associate( @@ -66,7 +88,7 @@ impl UdpAssociation { // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::>(1024); - let close_flag = Arc::new(AtomicBool::new(false)); + let close_flag = Arc::new(UdpAssociationState::new()); // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); @@ -93,7 +115,7 @@ impl UdpAssociation { // local <- remote let closed = close_flag.clone(); tokio::spawn(async move { - while !closed.load(Ordering::Acquire) { + while !closed.is_closed() { use std::io::ErrorKind; // Read and send back to source @@ -114,7 +136,7 @@ impl UdpAssociation { debug!("UDP ASSOCIATE {} <- .. finished", src_addr); }); - Ok(UdpAssociation { tx, closed: close_flag }) + Ok(UdpAssociation { tx, state: close_flag }) } /// Relay packets from local to remote @@ -239,13 +261,13 @@ impl UdpAssociation { Ok(()) } - async fn send(&mut self, pkt: &[u8]) -> bool { - match self.tx.send(pkt.to_vec()).await { - Ok(..) => true, - Err(err) => { - error!("failed to send packet, error: {}", err); - false - } + // Send packet to remote + // + // Return `Err` if receiver have been closed + async fn send(&mut self, pkt: Vec) { + if let Err(..) = self.tx.send(pkt).await { + // SHOULDn't HAPPEN + unreachable!("UDP Association local -> remote Queue closed unexpectly"); } } } @@ -324,28 +346,27 @@ async fn listen(context: SharedContext, svr_cfg: Arc) -> io::Resul } // Check or (re)create an association - loop { + let mut assoc = { + // Locks the whole association map 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(), - Entry::Vacant(vc) => vc.insert( - UdpAssociation::associate(context.clone(), svr_cfg.clone(), src, tx.clone()) - .await - .expect("Failed to create udp association"), - ), - }; - - !assoc.send(pkt).await + // Get or create an association + let assoc = match assoc_map.entry(src.to_string()) { + Entry::Occupied(oc) => oc.into_mut(), + Entry::Vacant(vc) => vc.insert( + UdpAssociation::associate(context.clone(), svr_cfg.clone(), src, tx.clone()) + .await + .expect("Failed to create udp association"), + ), }; - if retry { - assoc_map.remove(&src.to_string()); - } else { - break; - } - } + // Clone the handle and release the lock. + // Make sure we keep the critical section small + assoc.clone() + }; + + // Send to local -> remote task + assoc.send(pkt.to_vec()).await; } } diff --git a/src/relay/udprelay/socks5_local.rs b/src/relay/udprelay/socks5_local.rs index ccbc1017..80628bb5 100644 --- a/src/relay/udprelay/socks5_local.rs +++ b/src/relay/udprelay/socks5_local.rs @@ -61,13 +61,23 @@ async fn parse_packet(pkt: &[u8]) -> io::Result<(Address, Vec)> { Ok((addr, payload)) } -// Represent a UDP association -struct UdpAssociation { - tx: mpsc::Sender>, - closed: Arc, +struct UdpAssociationState { + closed: AtomicBool, } -impl Drop for UdpAssociation { +impl UdpAssociationState { + fn new() -> UdpAssociationState { + UdpAssociationState { + closed: AtomicBool::new(false), + } + } + + fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } +} + +impl Drop for UdpAssociationState { fn drop(&mut self) { // 1. Drops tx, will close local -> remote task // 2. Drops closed, will close local <- remote task @@ -75,6 +85,18 @@ impl Drop for UdpAssociation { } } +// Represent a UDP association +#[derive(Clone)] +struct UdpAssociation { + // local -> remote Queue + // Drops tx, will close local -> remote task + tx: mpsc::Sender>, + + // local <- remote relay task + // Drops state, will close local <- remote task + state: Arc, +} + impl UdpAssociation { /// Create an association with addr async fn associate( @@ -93,7 +115,7 @@ impl UdpAssociation { // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::>(1024); - let close_flag = Arc::new(AtomicBool::new(false)); + let close_flag = Arc::new(UdpAssociationState::new()); // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); @@ -120,7 +142,7 @@ impl UdpAssociation { // local <- remote let closed = close_flag.clone(); tokio::spawn(async move { - while !closed.load(Ordering::Acquire) { + while !closed.is_closed() { // Read and send back to source match UdpAssociation::relay_r2l(src_addr, &mut receiver, timeout, &mut response_tx, &*svr_cfg).await { Ok(..) => {} @@ -139,7 +161,7 @@ impl UdpAssociation { debug!("UDP ASSOCIATE {} <- .. finished", src_addr); }); - Ok(UdpAssociation { tx, closed: close_flag }) + Ok(UdpAssociation { tx, state: close_flag }) } /// Relay packets from local to remote @@ -248,13 +270,13 @@ impl UdpAssociation { Ok(()) } - async fn send(&mut self, pkt: &[u8]) -> bool { - match self.tx.send(pkt.to_vec()).await { - Ok(..) => true, - Err(err) => { - error!("failed to send packet, error: {}", err); - false - } + // Send packet to remote + // + // Return `Err` if receiver have been closed + async fn send(&mut self, pkt: Vec) { + if let Err(..) = self.tx.send(pkt).await { + // SHOULDn't HAPPEN + unreachable!("UDP Association local -> remote Queue closed unexpectly"); } } } @@ -335,28 +357,27 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { let svr_cfg = balancer.pick_server(); // Check or (re)create an association - loop { + let mut assoc = { + // Locks the whole association map 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(), - Entry::Vacant(vc) => vc.insert( - UdpAssociation::associate(context.clone(), svr_cfg.clone(), src, tx.clone()) - .await - .expect("Failed to create udp association"), - ), - }; - - !assoc.send(pkt).await + // Get or create an association + let assoc = match assoc_map.entry(src.to_string()) { + Entry::Occupied(oc) => oc.into_mut(), + Entry::Vacant(vc) => vc.insert( + UdpAssociation::associate(context.clone(), svr_cfg.clone(), src, tx.clone()) + .await + .expect("Failed to create udp association"), + ), }; - if retry { - assoc_map.remove(&src.to_string()); - } else { - break; - } - } + // Clone the handle and release the lock. + // Make sure we keep the critical section small + assoc.clone() + }; + + // Send to local -> remote task + assoc.send(pkt.to_vec()).await; } } diff --git a/src/relay/udprelay/tunnel_local.rs b/src/relay/udprelay/tunnel_local.rs index fa2def9b..e4acc30f 100644 --- a/src/relay/udprelay/tunnel_local.rs +++ b/src/relay/udprelay/tunnel_local.rs @@ -40,13 +40,23 @@ use super::{ MAXIMUM_UDP_PAYLOAD_SIZE, }; -// Represent a UDP association -struct UdpAssociation { - tx: mpsc::Sender>, - closed: Arc, +struct UdpAssociationState { + closed: AtomicBool, } -impl Drop for UdpAssociation { +impl UdpAssociationState { + fn new() -> UdpAssociationState { + UdpAssociationState { + closed: AtomicBool::new(false), + } + } + + fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } +} + +impl Drop for UdpAssociationState { fn drop(&mut self) { // 1. Drops tx, will close local -> remote task // 2. Drops closed, will close local <- remote task @@ -54,6 +64,18 @@ impl Drop for UdpAssociation { } } +// Represent a UDP association +#[derive(Clone)] +struct UdpAssociation { + // local -> remote Queue + // Drops tx, will close local -> remote task + tx: mpsc::Sender>, + + // local <- remote relay task + // Drops state, will close local <- remote task + state: Arc, +} + impl UdpAssociation { /// Create an association with addr async fn associate( @@ -72,7 +94,7 @@ impl UdpAssociation { // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::>(1024); - let close_flag = Arc::new(AtomicBool::new(false)); + let close_flag = Arc::new(UdpAssociationState::new()); // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); @@ -99,7 +121,7 @@ impl UdpAssociation { // local <- remote let closed = close_flag.clone(); tokio::spawn(async move { - while !closed.load(Ordering::Acquire) { + while !closed.is_closed() { use std::io::ErrorKind; // Read and send back to source @@ -120,7 +142,7 @@ impl UdpAssociation { debug!("UDP TUNNEL {} <- .. finished", src_addr); }); - Ok(UdpAssociation { tx, closed: close_flag }) + Ok(UdpAssociation { tx, state: close_flag }) } /// Relay packets from local to remote @@ -220,13 +242,13 @@ impl UdpAssociation { Ok(()) } - async fn send(&mut self, pkt: &[u8]) -> bool { - match self.tx.send(pkt.to_vec()).await { - Ok(..) => true, - Err(err) => { - error!("failed to send packet, error: {}", err); - false - } + // Send packet to remote + // + // Return `Err` if receiver have been closed + async fn send(&mut self, pkt: Vec) { + if let Err(..) = self.tx.send(pkt).await { + // SHOULDn't HAPPEN + unreachable!("UDP Association local -> remote Queue closed unexpectly"); } } } @@ -307,28 +329,27 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { let svr_cfg = balancer.pick_server(); // Check or (re)create an association - loop { + let mut assoc = { + // Locks the whole association map 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(), - Entry::Vacant(vc) => vc.insert( - UdpAssociation::associate(context.clone(), svr_cfg.clone(), src, tx.clone()) - .await - .expect("Failed to create udp association"), - ), - }; - - !assoc.send(pkt).await + // Get or create an association + let assoc = match assoc_map.entry(src.to_string()) { + Entry::Occupied(oc) => oc.into_mut(), + Entry::Vacant(vc) => vc.insert( + UdpAssociation::associate(context.clone(), svr_cfg.clone(), src, tx.clone()) + .await + .expect("Failed to create udp association"), + ), }; - if retry { - assoc_map.remove(&src.to_string()); - } else { - break; - } - } + // Clone the handle and release the lock. + // Make sure we keep the critical section small + assoc.clone() + }; + + // Send to local -> remote task + assoc.send(pkt.to_vec()).await; } }