diff --git a/src/relay/udprelay/association.rs b/src/relay/udprelay/association.rs index 3b93350a..5ecf2939 100644 --- a/src/relay/udprelay/association.rs +++ b/src/relay/udprelay/association.rs @@ -30,6 +30,7 @@ use crate::{ context::{Context, SharedContext}, crypto::CipherCategory, relay::{ + flow::SharedServerFlowStatistic, loadbalancing::server::{ServerData, SharedServerStatistic}, socks5::Address, sys::create_outbound_udp_socket, @@ -753,6 +754,78 @@ where /// Association manager for local pub type ProxyAssociationManager = AssociationManager; +/// Server Association's Key type +pub type ServerAssociationKey = [u8; 18]; + +/// Handler for handing remote response +pub struct ServerProxyHandler { + src_addr: SocketAddr, + cache_key: ServerAssociationKey, + assoc_manager: ServerAssociationManager, + flow_stat: SharedServerFlowStatistic, + tx: Arc, +} + +impl ServerProxyHandler { + /// Create a new ServerProxyHandler + pub fn new( + src_addr: SocketAddr, + assoc_manager: ServerAssociationManager, + flow_stat: SharedServerFlowStatistic, + tx: Arc, + ) -> ServerProxyHandler { + ServerProxyHandler { + src_addr, + cache_key: ServerProxyHandler::association_key(&src_addr), + assoc_manager, + flow_stat, + tx, + } + } + + /// Make server association key + pub fn association_key(saddr: &SocketAddr) -> ServerAssociationKey { + let mut result = [0; 18]; + result[..16].copy_from_slice(&match saddr.ip() { + IpAddr::V4(ref ip) => ip.to_ipv6_mapped().octets(), + IpAddr::V6(ref ip) => ip.octets(), + }); + result[16..].copy_from_slice(&saddr.port().to_ne_bytes()); + result + } + + /// Send packet back to source client + pub async fn send_packet(&self, pkt: &[u8]) -> io::Result<()> { + if !self.assoc_manager.keep_alive(&self.cache_key).await { + debug!( + "UDP association {} <-> ... is already expired, throwing away packet {} bytes", + self.src_addr, + pkt.len() + ); + return Ok(()); + } + + let n = match self.tx.send_to(&pkt, &self.src_addr).await { + Ok(n) => { + if n < pkt.len() { + warn!( + "UDP association {} <- ... payload truncated, expecting {} bytes, but sent {} bytes", + self.src_addr, + pkt.len(), + n + ); + } + n + } + Err(err) => return Err(err), + }; + + self.flow_stat.udp().incr_tx(n); + + Ok(()) + } +} + // Represent a UDP association in server pub struct ServerAssociation { // local -> remote Queue @@ -777,7 +850,7 @@ impl ServerAssociation { context: SharedContext, svr_idx: usize, src_addr: SocketAddr, - mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>, + response_tx: ServerProxyHandler, ) -> io::Result { // Create a socket for receiving packets // Let system allocate an address for us (INADDR_ANY) @@ -844,7 +917,7 @@ impl ServerAssociation { &context, src_addr, &receiver, - &mut response_tx, + &response_tx, svr_cfg, &resolved_address_cache, ) @@ -969,7 +1042,7 @@ impl ServerAssociation { context: &Context, src_addr: SocketAddr, remote_udp: &UdpSocket, - response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>, + response_tx: &ServerProxyHandler, svr_cfg: &ServerConfig, resolved_address_cache: &SharedResolvedAddressCache, ) -> io::Result<()> { @@ -996,7 +1069,7 @@ impl ServerAssociation { if let CipherCategory::None = svr_cfg.method().category() { // Send back to src_addr - if let Err(err) = response_tx.send((src_addr, send_buf)).await { + if let Err(err) = response_tx.send_packet(&send_buf).await { error!("failed to send packet into response channel, error: {}", err); // FIXME: What to do? Ignore? @@ -1006,7 +1079,7 @@ impl ServerAssociation { encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?; // Send back to src_addr - if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await { + if let Err(err) = response_tx.send_packet(&encrypt_buf).await { error!("failed to send packet into response channel, error: {}", err); // FIXME: What to do? Ignore? diff --git a/src/relay/udprelay/redir_local.rs b/src/relay/udprelay/redir_local.rs index c965ecfd..8a922ead 100644 --- a/src/relay/udprelay/redir_local.rs +++ b/src/relay/udprelay/redir_local.rs @@ -3,7 +3,7 @@ use std::{io, net::SocketAddr, time::Duration}; use async_trait::async_trait; -use log::{debug, error, info, trace}; +use log::{debug, error, info, trace, warn}; use tokio::time; use crate::{ @@ -54,7 +54,19 @@ impl ProxySend for ProxyHandler { // This only works for systems that supports binding to non-local addresses let mut local_udp = UdpRedirSocket::bind(self.ty, dst_addr)?; - local_udp.send_to(&data, self.src_addr).await?; + match local_udp.send_to(&data, self.src_addr).await { + Ok(n) => { + if n < data.len() { + warn!( + "UDP association {} <- ... payload truncated, expecting {} bytes, but sent {} bytes", + self.src_addr, + data.len(), + n + ); + } + } + Err(err) => return Err(err), + } // Update LRU self.assoc_map.keep_alive(&self.cache_key).await; diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs index e13992c8..85126e4f 100644 --- a/src/relay/udprelay/server.rs +++ b/src/relay/udprelay/server.rs @@ -1,16 +1,10 @@ //! UDP relay proxy server -use std::{ - io, - net::{IpAddr, SocketAddr}, - sync::Arc, - time::Duration, -}; +use std::{io, sync::Arc, time::Duration}; -use bytes::BytesMut; use futures::{stream::FuturesUnordered, StreamExt}; use log::{debug, error, info, trace, warn}; -use tokio::{self, sync::mpsc, time}; +use tokio::{self, time}; use crate::{ context::SharedContext, @@ -21,20 +15,10 @@ use crate::{ }; use super::{ - association::{ServerAssociation, ServerAssociationManager}, + association::{ServerAssociation, ServerAssociationManager, ServerProxyHandler}, MAXIMUM_UDP_PAYLOAD_SIZE, }; -fn association_key(saddr: &SocketAddr) -> [u8; 18] { - let mut result = [0; 18]; - result[..16].copy_from_slice(&match saddr.ip() { - IpAddr::V4(ref ip) => ip.to_ipv6_mapped().octets(), - IpAddr::V6(ref ip) => ip.octets(), - }); - result[16..].copy_from_slice(&saddr.port().to_ne_bytes()); - result -} - async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> { let svr_cfg = context.server_config(svr_idx); let listen_addr = svr_cfg.addr().bind_addr(&context).await?; @@ -48,39 +32,6 @@ async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, sv let assoc_manager = ServerAssociationManager::new(context.config()); - // FIXME: Channel size 1024? - let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024); - - { - // Tokio task for sending data back to clients - - let assoc_manager = assoc_manager.clone(); - let flow_stat = flow_stat.clone(); - - tokio::spawn(async move { - while let Some((src, pkt)) = rx.recv().await { - let cache_key = association_key(&src); - 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 { - error!("UDP packet send failed, err: {:?}", err); - break; - } - - flow_stat.udp().incr_tx(pkt.len()); - } - - // FIXME: How to stop the outer listener Future? - }); - } - let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; loop { @@ -118,8 +69,9 @@ async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, sv // Check or (re)create an association let res = assoc_manager - .send_packet(association_key(&src), pkt.to_vec(), async { - ServerAssociation::associate(context.clone(), svr_idx, src, tx.clone()).await + .send_packet(ServerProxyHandler::association_key(&src), pkt.to_vec(), async { + let handler = ServerProxyHandler::new(src, assoc_manager.clone(), flow_stat.clone(), w.clone()); + ServerAssociation::associate(context.clone(), svr_idx, src, handler).await }) .await; diff --git a/src/relay/udprelay/socks5_local.rs b/src/relay/udprelay/socks5_local.rs index 3852c5a6..71bee94e 100644 --- a/src/relay/udprelay/socks5_local.rs +++ b/src/relay/udprelay/socks5_local.rs @@ -9,8 +9,8 @@ use std::{ use async_trait::async_trait; use bytes::{Bytes, BytesMut}; -use log::{debug, error, info, trace}; -use tokio::{self, sync::mpsc, time}; +use log::{debug, error, info, trace, warn}; +use tokio::{self, net::UdpSocket, time}; use crate::{ context::SharedContext, @@ -29,16 +29,50 @@ use super::{ #[derive(Clone)] struct ProxyHandler { src_addr: SocketAddr, - response_tx: mpsc::Sender<(SocketAddr, Address, Vec)>, + cache_key: String, + assoc_manager: ProxyAssociationManager, + tx: Arc, +} + +impl ProxyHandler { + fn new(src_addr: SocketAddr, assoc_manager: ProxyAssociationManager, tx: Arc) -> ProxyHandler { + ProxyHandler { + src_addr, + cache_key: src_addr.to_string(), + assoc_manager, + tx, + } + } } #[async_trait] impl ProxySend for ProxyHandler { async fn send_packet(&mut self, addr: Address, data: Vec) -> io::Result<()> { - if let Err(err) = self.response_tx.send((self.src_addr, addr, data)).await { - error!("UDP associate response channel error: {}", err); + if !self.assoc_manager.keep_alive(&self.cache_key).await { + debug!( + "UDP association {} <-> ... is already expired, throwing away packet {} bytes", + self.src_addr, + data.len() + ); + return Ok(()); + } + + let payload = assemble_packet(addr, &data); + + match self.tx.send_to(&payload, &self.src_addr).await { + Ok(n) => { + if n < data.len() { + warn!( + "UDP association {} <- ... payload truncated, expecting {} bytes, but sent {} bytes", + self.src_addr, + payload.len(), + n + ); + } + Ok(()) + } + Err(err) => return Err(err), } - Ok(()) } } @@ -97,37 +131,9 @@ pub async fn run(context: SharedContext) -> io::Result<()> { info!("shadowsocks SOCKS5 UDP listening on {}", local_addr); 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_manager = assoc_manager_cloned; - - while let Some((src, target, pkt)) = rx.recv().await { - let cache_key = src.to_string(); - 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); - - if let Err(err) = w.send_to(&payload, &src).await { - error!("UDP packet send failed, err: {:?}", err); - break; - } - } - - // FIXME: How to stop the outer listener Future? - }); - loop { let (recv_len, src) = match r.recv_from(&mut pkt_buf).await { Ok(o) => o, @@ -173,10 +179,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { // Pick a server let server = balancer.pick_server(); - let sender = ProxyHandler { - src_addr: src, - response_tx: tx.clone(), - }; + let sender = ProxyHandler::new(src, assoc_manager.clone(), w.clone()); ProxyAssociation::associate_with_acl(src, server, sender).await }) diff --git a/src/relay/udprelay/tunnel_local.rs b/src/relay/udprelay/tunnel_local.rs index 292d6521..23c0982a 100644 --- a/src/relay/udprelay/tunnel_local.rs +++ b/src/relay/udprelay/tunnel_local.rs @@ -3,8 +3,8 @@ use std::{io, net::SocketAddr, sync::Arc, time::Duration}; use async_trait::async_trait; -use log::{debug, error, info, trace}; -use tokio::{self, sync::mpsc, time}; +use log::{debug, error, info, trace, warn}; +use tokio::{self, net::UdpSocket, time}; use crate::{ context::SharedContext, @@ -23,16 +23,48 @@ use super::{ #[derive(Clone)] struct ProxyHandler { src_addr: SocketAddr, - response_tx: mpsc::Sender<(SocketAddr, Address, Vec)>, + cache_key: String, + assoc_manager: ProxyAssociationManager, + tx: Arc, +} + +impl ProxyHandler { + fn new(src_addr: SocketAddr, assoc_manager: ProxyAssociationManager, tx: Arc) -> ProxyHandler { + ProxyHandler { + src_addr, + cache_key: src_addr.to_string(), + assoc_manager, + tx, + } + } } #[async_trait] impl ProxySend for ProxyHandler { - async fn send_packet(&mut self, addr: Address, data: Vec) -> io::Result<()> { - if let Err(err) = self.response_tx.send((self.src_addr, addr, data)).await { - error!("UDP associate response channel error: {}", err); + async fn send_packet(&mut self, _addr: Address, data: Vec) -> io::Result<()> { + if !self.assoc_manager.keep_alive(&self.cache_key).await { + debug!( + "UDP association {} <-> ... is already expired, throwing away packet {} bytes", + self.src_addr, + data.len() + ); + return Ok(()); + } + + match self.tx.send_to(&data, &self.src_addr).await { + Ok(n) => { + if n < data.len() { + warn!( + "UDP association {} <- ... payload truncated, expecting {} bytes, but sent {} bytes", + self.src_addr, + data.len(), + n + ); + } + Ok(()) + } + Err(err) => return Err(err), } - Ok(()) } } @@ -57,35 +89,9 @@ pub async fn run(context: SharedContext) -> io::Result<()> { ); 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_manager = assoc_manager_cloned; - - while let Some((src, _, pkt)) = rx.recv().await { - let cache_key = src.to_string(); - 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 { - error!("UDP packet send failed, err: {:?}", err); - break; - } - } - - // FIXME: How to stop the outer listener Future? - }); - loop { let (recv_len, src) = match r.recv_from(&mut pkt_buf).await { Ok(o) => o, @@ -94,7 +100,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { time::sleep(Duration::from_secs(1)).await; continue; } - }; + }; // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded. // Copy bytes, because udp_associate runs in another tokio Task @@ -119,10 +125,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { // Pick a server let server = balancer.pick_server(); - let sender = ProxyHandler { - src_addr: src, - response_tx: tx.clone(), - }; + let sender = ProxyHandler::new(src, assoc_manager.clone(), w.clone()); ProxyAssociation::associate_with_acl(src, server, sender).await })