From e529d1541b1171ee74f17242f8e46bf9226cf331 Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Fri, 27 Dec 2019 22:27:52 +0800 Subject: [PATCH] Better strategy to drop UDP local <- remote task --- src/relay/local.rs | 4 +- src/relay/server.rs | 4 +- src/relay/udprelay/server.rs | 80 +++++++++++------------------ src/relay/udprelay/socks5_local.rs | 79 +++++++++++----------------- src/relay/udprelay/tunnel_local.rs | 82 +++++++++++------------------- 5 files changed, 96 insertions(+), 153 deletions(-) diff --git a/src/relay/local.rs b/src/relay/local.rs index 1fd1e982..18f0eca5 100644 --- a/src/relay/local.rs +++ b/src/relay/local.rs @@ -3,7 +3,7 @@ use std::io; use futures::{future::select_all, FutureExt}; -use log::error; +use log::{error, trace}; use tokio::runtime::Handle; use crate::{ @@ -15,6 +15,8 @@ use crate::{ /// Relay server running under local environment. pub async fn run(mut config: Config, rt: Handle) -> io::Result<()> { + trace!("{:?}", config); + // Create a context containing a DNS resolver and server running state flag. let state = SharedServerState::new(&config, rt).await?; diff --git a/src/relay/server.rs b/src/relay/server.rs index e5a46a39..8027e3bf 100644 --- a/src/relay/server.rs +++ b/src/relay/server.rs @@ -3,7 +3,7 @@ use std::io; use futures::future::{select_all, FutureExt}; -use log::error; +use log::{error, trace}; use tokio::runtime::Handle; use crate::{ @@ -15,6 +15,8 @@ use crate::{ /// Relay server running on server side. pub async fn run(mut config: Config, rt: Handle) -> io::Result<()> { + trace!("{:?}", config); + // Create a context containing a DNS resolver and server running state flag. let state = SharedServerState::new(&config, rt).await?; diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs index be6d22b1..9a6845a6 100644 --- a/src/relay/udprelay/server.rs +++ b/src/relay/udprelay/server.rs @@ -3,21 +3,18 @@ use std::{ io::{self, Cursor}, net::{IpAddr, Ipv4Addr, SocketAddr}, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, + sync::Arc, time::Duration, }; use bytes::BytesMut; -use futures::{self, stream::FuturesUnordered, StreamExt}; +use futures::{self, future, stream::FuturesUnordered, FutureExt, StreamExt}; use log::{debug, error, info, trace}; use lru_time_cache::{Entry, LruCache}; use tokio::{ self, net::udp::{RecvHalf, SendHalf}, - sync::{mpsc, Mutex}, + sync::{mpsc, oneshot, Mutex}, time, }; @@ -34,29 +31,7 @@ use super::{ MAXIMUM_UDP_PAYLOAD_SIZE, }; -struct UdpAssociationState { - closed: AtomicBool, -} - -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 - self.closed.store(true, Ordering::Release); - } -} +struct UdpAssociationWatcher(oneshot::Sender<()>); // Represent a UDP association #[derive(Clone)] @@ -65,9 +40,8 @@ struct UdpAssociation { // Drops tx, will close local -> remote task tx: mpsc::Sender>, - // local <- remote relay task - // Drops state, will close local <- remote task - state: Arc, + // local <- remote task life watcher + watcher: Arc, } impl UdpAssociation { @@ -88,7 +62,10 @@ impl UdpAssociation { // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::>(1024); - let close_flag = Arc::new(UdpAssociationState::new()); + // Create a watcher for local <- remote task + let (watcher_tx, watcher_rx) = oneshot::channel::<()>(); + + let close_flag = Arc::new(UdpAssociationWatcher(watcher_tx)); // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); @@ -113,30 +90,32 @@ impl UdpAssociation { }); // local <- remote - let closed = close_flag.clone(); tokio::spawn(async move { - while !closed.is_closed() { - use std::io::ErrorKind; + let transfer_fut = async move { + loop { + // Read and send back to source + match UdpAssociation::relay_r2l(src_addr, &mut receiver, &mut response_tx, &*svr_cfg).await { + Ok(..) => {} + Err(err) => { + error!("failed to receive packet, {} <- .., error: {}", src_addr, err); - // Read and send back to source - match UdpAssociation::relay_r2l(src_addr, &mut receiver, timeout, &mut response_tx, &*svr_cfg).await { - Ok(..) => {} - Err(ref err) if err.kind() == ErrorKind::TimedOut => { - trace!("receive packet timeout, {} <- ...", src_addr); - } - Err(err) => { - error!("railed to receive packet, {} <- .., error: {}", src_addr, err); - - // FIXME: Don't break, or if you can find a way to drop the UdpAssociation - // break; + // FIXME: Don't break, or if you can find a way to drop the UdpAssociation + // break; + } } } - } + }; + + // Resolved only if watcher_rx resolved + let _ = future::select(transfer_fut.boxed(), watcher_rx.boxed()).await; debug!("UDP ASSOCIATE {} <- .. finished", src_addr); }); - Ok(UdpAssociation { tx, state: close_flag }) + Ok(UdpAssociation { + tx, + watcher: close_flag, + }) } /// Relay packets from local to remote @@ -226,14 +205,13 @@ impl UdpAssociation { async fn relay_r2l( src_addr: SocketAddr, remote_udp: &mut RecvHalf, - timeout: Duration, response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>, svr_cfg: &ServerConfig, ) -> io::Result<()> { // Waiting for response from server SERVER -> CLIENT // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded. let mut remote_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; - let (remote_recv_len, remote_addr) = try_timeout(remote_udp.recv_from(&mut remote_buf), Some(timeout)).await?; + let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?; debug!( "UDP ASSOCIATE {} <- {}, payload length {} bytes", diff --git a/src/relay/udprelay/socks5_local.rs b/src/relay/udprelay/socks5_local.rs index f31fb3d4..c85ab10a 100644 --- a/src/relay/udprelay/socks5_local.rs +++ b/src/relay/udprelay/socks5_local.rs @@ -3,20 +3,18 @@ use std::{ io::{self, Cursor, ErrorKind, Read}, net::{IpAddr, Ipv4Addr, SocketAddr}, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, + sync::Arc, time::Duration, }; use bytes::BytesMut; +use futures::{future, FutureExt}; use log::{debug, error, info, trace}; use lru_time_cache::{Entry, LruCache}; use tokio::{ self, net::udp::{RecvHalf, SendHalf}, - sync::{mpsc, Mutex}, + sync::{mpsc, oneshot, Mutex}, time, }; @@ -58,29 +56,7 @@ async fn parse_packet(pkt: &[u8]) -> io::Result<(Address, Vec)> { Ok((addr, payload)) } -struct UdpAssociationState { - closed: AtomicBool, -} - -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 - self.closed.store(true, Ordering::Release); - } -} +struct UdpAssociationWatcher(oneshot::Sender<()>); // Represent a UDP association #[derive(Clone)] @@ -89,9 +65,8 @@ struct UdpAssociation { // Drops tx, will close local -> remote task tx: mpsc::Sender>, - // local <- remote relay task - // Drops state, will close local <- remote task - state: Arc, + // local <- remote task life watcher + watcher: Arc, } impl UdpAssociation { @@ -112,7 +87,10 @@ impl UdpAssociation { // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::>(1024); - let close_flag = Arc::new(UdpAssociationState::new()); + // Create a watcher for local <- remote task + let (watcher_tx, watcher_rx) = oneshot::channel::<()>(); + + let close_flag = Arc::new(UdpAssociationWatcher(watcher_tx)); // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); @@ -137,28 +115,32 @@ impl UdpAssociation { }); // local <- remote - let closed = close_flag.clone(); tokio::spawn(async move { - 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(..) => {} - Err(ref err) if err.kind() == ErrorKind::TimedOut => { - trace!("receive packet timeout, {} <- ...", src_addr); - } - Err(err) => { - error!("failed to receive packet, {} <- .., error: {}", src_addr, err); + let transfer_fut = async move { + loop { + // Read and send back to source + match UdpAssociation::relay_r2l(src_addr, &mut receiver, &mut response_tx, &*svr_cfg).await { + Ok(..) => {} + Err(err) => { + error!("failed to receive packet, {} <- .., error: {}", src_addr, err); - // FIXME: Don't break, or if you can find a way to drop the UdpAssociation - // break; + // FIXME: Don't break, or if you can find a way to drop the UdpAssociation + // break; + } } } - } + }; + + // Resolved only if watcher_rx resolved + let _ = future::select(transfer_fut.boxed(), watcher_rx.boxed()).await; debug!("UDP ASSOCIATE {} <- .. finished", src_addr); }); - Ok(UdpAssociation { tx, state: close_flag }) + Ok(UdpAssociation { + tx, + watcher: close_flag, + }) } /// Relay packets from local to remote @@ -221,14 +203,13 @@ impl UdpAssociation { async fn relay_r2l( src_addr: SocketAddr, remote_udp: &mut RecvHalf, - timeout: Duration, response_tx: &mut mpsc::Sender<(SocketAddr, Vec)>, svr_cfg: &ServerConfig, ) -> io::Result<()> { // Waiting for response from server SERVER -> CLIENT // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded. let mut recv_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; - let (recv_n, remote_addr) = try_timeout(remote_udp.recv_from(&mut recv_buf), Some(timeout)).await?; + let (recv_n, remote_addr) = remote_udp.recv_from(&mut recv_buf).await?; let decrypt_buf = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), &recv_buf[..recv_n])? { None => { @@ -329,7 +310,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { }); loop { - let (recv_len, src) = match time::timeout(DEFAULT_TIMEOUT, r.recv_from(&mut pkt_buf)).await { + let (recv_len, src) = match time::timeout(timeout, r.recv_from(&mut pkt_buf)).await { Ok(r) => r?, Err(..) => { // Cleanup expired association diff --git a/src/relay/udprelay/tunnel_local.rs b/src/relay/udprelay/tunnel_local.rs index 055f0a6c..c4aebefe 100644 --- a/src/relay/udprelay/tunnel_local.rs +++ b/src/relay/udprelay/tunnel_local.rs @@ -3,20 +3,18 @@ use std::{ io::{self, Cursor, Read}, net::{IpAddr, Ipv4Addr, SocketAddr}, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, + sync::Arc, time::Duration, }; use bytes::BytesMut; +use futures::{future, FutureExt}; use log::{debug, error, info, trace}; use lru_time_cache::{Entry, LruCache}; use tokio::{ self, net::udp::{RecvHalf, SendHalf}, - sync::{mpsc, Mutex}, + sync::{mpsc, oneshot, Mutex}, time, }; @@ -37,29 +35,8 @@ use super::{ MAXIMUM_UDP_PAYLOAD_SIZE, }; -struct UdpAssociationState { - closed: AtomicBool, -} - -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 - self.closed.store(true, Ordering::Release); - } -} +// Drop the oneshot::Sender<()> will trigger local <- remote task to finish +struct UdpAssociationWatcher(oneshot::Sender<()>); // Represent a UDP association #[derive(Clone)] @@ -68,9 +45,8 @@ struct UdpAssociation { // Drops tx, will close local -> remote task tx: mpsc::Sender>, - // local <- remote relay task - // Drops state, will close local <- remote task - state: Arc, + // local <- remote task life watcher + watcher: Arc, } impl UdpAssociation { @@ -91,7 +67,10 @@ impl UdpAssociation { // FIXME: Channel size 1024? let (tx, mut rx) = mpsc::channel::>(1024); - let close_flag = Arc::new(UdpAssociationState::new()); + // Create a watcher for local <- remote task + let (watcher_tx, watcher_rx) = oneshot::channel::<()>(); + + let close_flag = Arc::new(UdpAssociationWatcher(watcher_tx)); // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); @@ -116,30 +95,32 @@ impl UdpAssociation { }); // local <- remote - let closed = close_flag.clone(); tokio::spawn(async move { - while !closed.is_closed() { - use std::io::ErrorKind; + let transfer_fut = async move { + loop { + // Read and send back to source + match UdpAssociation::relay_r2l(src_addr, &mut receiver, &mut response_tx, &*svr_cfg).await { + Ok(..) => {} + Err(err) => { + error!("failed to receive packet, {} <- .., error: {}", src_addr, err); - // Read and send back to source - match UdpAssociation::relay_r2l(src_addr, &mut receiver, timeout, &mut response_tx, &*svr_cfg).await { - Ok(..) => {} - Err(ref err) if err.kind() == ErrorKind::TimedOut => { - trace!("receive packet timeout, {} <- ...", src_addr); - } - Err(err) => { - error!("failed to receive packet, {} <- .., error: {}", src_addr, err); - - // FIXME: Don't break, or if you can find a way to drop the UdpAssociation - // break; + // FIXME: Don't break, or if you can find a way to drop the UdpAssociation + // break; + } } } - } + }; + + // Resolved only if watcher_rx resolved + let _ = future::select(transfer_fut.boxed(), watcher_rx.boxed()).await; debug!("UDP TUNNEL {} <- .. finished", src_addr); }); - Ok(UdpAssociation { tx, state: close_flag }) + Ok(UdpAssociation { + tx, + watcher: close_flag, + }) } /// Relay packets from local to remote @@ -196,7 +177,6 @@ impl UdpAssociation { async fn relay_r2l( src_addr: SocketAddr, remote_udp: &mut RecvHalf, - timeout: Duration, response_tx: &mut mpsc::Sender<(SocketAddr, Vec)>, svr_cfg: &ServerConfig, ) -> io::Result<()> { @@ -204,7 +184,7 @@ impl UdpAssociation { // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded. let mut recv_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; - let (recv_n, remote_addr) = try_timeout(remote_udp.recv_from(&mut recv_buf), Some(timeout)).await?; + let (recv_n, remote_addr) = remote_udp.recv_from(&mut recv_buf).await?; let decrypt_buf = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), &recv_buf[..recv_n])? { None => { @@ -305,7 +285,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { }); loop { - let (recv_len, src) = match time::timeout(DEFAULT_TIMEOUT, r.recv_from(&mut pkt_buf)).await { + let (recv_len, src) = match time::timeout(timeout, r.recv_from(&mut pkt_buf)).await { Ok(r) => r?, Err(..) => { // Cleanup expired association