diff --git a/src/relay/udprelay/local.rs b/src/relay/udprelay/local.rs index 11274297..7bd88bd7 100644 --- a/src/relay/udprelay/local.rs +++ b/src/relay/udprelay/local.rs @@ -1,40 +1,30 @@ //! UDP relay local server -use std::{ - io::{self, Cursor, ErrorKind, Read}, - net::{IpAddr, Ipv4Addr, SocketAddr}, - sync::Arc, - time::Duration, -}; +use std::io::{self, Cursor, ErrorKind, Read}; +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; +use std::time::Duration; use bytes::BytesMut; -use log::{debug, error, info, warn}; -use tokio::{self, net::UdpSocket, sync::mpsc}; +use log::{debug, error, info, trace}; +use lru_time_cache::{Entry, LruCache}; +use tokio; +use tokio::net::udp::{RecvHalf, SendHalf}; +use tokio::net::UdpSocket; +use tokio::sync::mpsc; -use crate::{ - config::{ServerAddr, ServerConfig}, - context::SharedContext, - relay::{ - loadbalancing::server::{LoadBalancer, RoundRobin}, - socks5::{Address, UdpAssociateHeader}, - utils::try_timeout, - }, -}; +use crate::config::{ServerAddr, ServerConfig}; +use crate::context::SharedContext; +use crate::relay::loadbalancing::server::{LoadBalancer, RoundRobin}; +use crate::relay::socks5::{Address, UdpAssociateHeader}; +use crate::relay::utils::try_timeout; -use super::{ - crypto_io::{decrypt_payload, encrypt_payload}, - MAXIMUM_UDP_PAYLOAD_SIZE, -}; +use super::crypto_io::{decrypt_payload, encrypt_payload}; +use super::MAXIMUM_UDP_PAYLOAD_SIZE; -#[allow(unused_variables)] // `context` is only used if trust-dns is enabled -async fn udp_associate( - context: SharedContext, - svr_cfg: Arc, - pkt: Vec, - src: SocketAddr, -) -> io::Result> { - const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); +const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); +async fn parse_packet(pkt: &[u8]) -> io::Result<(Address, Vec)> { // PKT = UdpAssociateHeader + PAYLOAD let mut cur = Cursor::new(pkt); @@ -50,89 +40,181 @@ async fn udp_associate( // The remaining is PAYLOAD let mut payload = Vec::new(); - Read::read_to_end(&mut cur, &mut payload)?; + cur.read_to_end(&mut payload)?; - // Binds to 0.0.0.0:0 (let system choose a random port) - // FIXME: Create a UdpSocket for every UDP associate requests - let local_addr = SocketAddr::new(IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)), 0); - let mut remote_udp = UdpSocket::bind(&local_addr).await?; + Ok((addr, payload)) +} - // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD - let mut send_buf = Vec::new(); - addr.write_to_buf(&mut send_buf); - send_buf.extend_from_slice(&payload); +// Represent a UDP association +struct UdpAssociation { + tx: mpsc::Sender<(Address, Vec)>, +} - debug!( - "UDP ASSOCIATE {} -> {}, payload length {} bytes", - src, - addr, - payload.len() - ); +impl UdpAssociation { + /// Create an association with addr + async fn associate( + svr_cfg: Arc, + src_addr: SocketAddr, + mut response_tx: mpsc::Sender<(SocketAddr, Vec)>, + ) -> io::Result { + debug!("Created UDP Association for {}", src_addr); - let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); + // Create a socket for receiving packets + let local_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0); + let remote_udp = UdpSocket::bind(&local_addr).await?; - // Write to remote socket CLIENT -> SERVER - let mut encrypt_buf = BytesMut::new(); - encrypt_payload(svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?; + // Create a channel for sending packets to remote + // FIXME: Channel size 1024? + let (tx, mut rx) = mpsc::channel::<(Address, Vec)>(1024); - let send_n = match *svr_cfg.addr() { - ServerAddr::SocketAddr(ref addr) => try_timeout(remote_udp.send_to(&encrypt_buf, addr), Some(timeout)).await?, - #[cfg(feature = "trust-dns")] - ServerAddr::DomainName(ref dname, port) => { - use crate::relay::dns_resolver::resolve; + // Splits socket into sender and receiver + let (mut receiver, mut sender) = remote_udp.split(); - let vec_ipaddr = resolve(context, dname, port, false).await?; - assert!(!vec_ipaddr.is_empty()); + let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); - try_timeout(remote_udp.send_to(&encrypt_buf, &vec_ipaddr[0]), Some(timeout)).await? - } - #[cfg(not(feature = "trust-dns"))] - ServerAddr::DomainName(ref dname, port) => { - try_timeout(remote_udp.send_to(&encrypt_buf, (dname.as_str(), port)), Some(timeout)).await? - } - }; + // local -> remote + let c_svr_cfg = svr_cfg.clone(); + tokio::spawn(async move { + while let Some((addr, pkt)) = rx.recv().await { + // pkt is already a raw packet, so just send it + if let Err(err) = + UdpAssociation::relay_l2r(src_addr, &addr, &mut sender, &pkt[..], timeout, &*c_svr_cfg).await + { + error!("Failed to send packet to {}, error: {}", addr, err); - if send_n != encrypt_buf.len() { - warn!( - "Sent packet length {}, but expected length {}", - send_n, - encrypt_buf.len() - ); + // FIXME: Ignore? Or how to deal with it? + } + } + + debug!("UDP ASSOCIATE {} -> .. finished", src_addr); + }); + + // local <- remote + tokio::spawn(async move { + loop { + // Read and send back to source + if let Err(err) = + UdpAssociation::relay_r2l(src_addr, &mut receiver, timeout, &mut response_tx, &*svr_cfg).await + { + error!("Failed to receive packet, {} <- .., error: {}", src_addr, err); + break; + } + } + + debug!("UDP ASSOCIATE {} <- .. finished", src_addr); + }); + + Ok(UdpAssociation { tx }) } - // 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 = try_timeout(remote_udp.recv(&mut recv_buf), Some(timeout)).await?; + /// Relay packets from local to remote + async fn relay_l2r( + src: SocketAddr, + addr: &Address, + remote_udp: &mut SendHalf, + payload: &[u8], + timeout: Duration, + svr_cfg: &ServerConfig, + ) -> io::Result<()> { + debug!( + "UDP ASSOCIATE {} -> {}, payload length {} bytes", + src, + addr, + payload.len() + ); - let decrypt_buf = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), &recv_buf[..recv_n])? { - None => { - error!("UDP packet too short, received length {}", recv_n); - let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short"); - return Err(err); + // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD + let mut send_buf = Vec::new(); + addr.write_to_buf(&mut send_buf); + send_buf.extend_from_slice(&payload); + + let mut encrypt_buf = BytesMut::new(); + encrypt_payload(svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?; + + let send_len = match svr_cfg.addr() { + ServerAddr::SocketAddr(ref remote_addr) => { + try_timeout(remote_udp.send_to(&encrypt_buf[..], remote_addr), Some(timeout)).await? + } + #[cfg(feature = "trust-dns")] + ServerAddr::DomainName(ref dname, port) => { + use crate::relay::dns_resolver::resolve; + + let vec_ipaddr = resolve(context, dname, port, false).await?; + assert!(!vec_ipaddr.is_empty()); + + try_timeout(remote_udp.send_to(&encrypt_buf[..], &vec_ipaddr[0]), Some(timeout)).await? + } + #[cfg(not(feature = "trust-dns"))] + ServerAddr::DomainName(ref dname, port) => { + // try_timeout(remote_udp.send_to(&encrypt_buf[..], (dname.as_str(), port)), Some(timeout)).await? + unimplemented!("tokio's UdpSocket SendHalf doesn't support ToSocketAddrs"); + } + }; + + assert_eq!(encrypt_buf.len(), send_len); + + Ok(()) + } + + /// Relay packets from remote to local + 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 decrypt_buf = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), &recv_buf[..recv_n])? { + None => { + error!("UDP packet too short, received length {}", recv_n); + let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short"); + return Err(err); + } + Some(b) => b, + }; + // SERVER -> CLIENT protocol: ADDRESS + PAYLOAD + let mut cur = Cursor::new(decrypt_buf); + // FIXME: Address is ignored. Maybe useful in the future if we uses one common UdpSocket for communicate with remote server + let _ = Address::read_from(&mut cur).await?; + + let mut payload = Vec::new(); + + let header = UdpAssociateHeader::new(0, Address::SocketAddress(src_addr)); + + header.write_to_buf(&mut payload); + cur.read_to_end(&mut payload)?; + + debug!( + "UDP ASSOCIATE {} <- {}, payload length {} bytes", + src_addr, + remote_addr, + payload.len() + ); + + // Send back to src_addr + if let Err(err) = response_tx.send((src_addr, payload)).await { + error!("Failed to send packet into response channel, error: {}", err); + + // FIXME: What to do? Ignore? } - Some(b) => b, - }; - // SERVER -> CLIENT protocol: ADDRESS + PAYLOAD - let mut cur = Cursor::new(decrypt_buf); + Ok(()) + } - // FIXME: Address is ignored. Maybe useful in the future if we uses one common UdpSocket for communicate with remote server - let _ = Address::read_from(&mut cur).await?; - - let payload_len = cur.get_ref().len() - cur.position() as usize; - debug!( - "UDP ASSOCIATE {} <- {}, payload length {} bytes", - src, addr, payload_len - ); - - let mut data = Vec::new(); - UdpAssociateHeader::new(0, Address::SocketAddress(src)).write_to_buf(&mut data); - - // Copy payload directly - Read::read_to_end(&mut cur, &mut data)?; - - Ok(data) + async fn send(&mut self, remote_addr: &Address, pkt: &[u8]) -> bool { + match self.tx.send((remote_addr.clone(), pkt.to_vec())).await { + Ok(..) => true, + Err(err) => { + error!("Failed to send packet, error: {}", err); + false + } + } + } } async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { @@ -155,32 +237,57 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> { // FIXME: How to stop the outer listener Future? }); + // let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); + let mut assoc_map = LruCache::with_expiry_duration(DEFAULT_TIMEOUT); + loop { 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 - let pkt = pkt_buf[..recv_len].to_vec(); + let pkt = &pkt_buf[..recv_len]; - let svr_cfg = balancer.pick_server(); + trace!("Received UDP packet from {}, length {} bytes", src, recv_len); - let context = context.clone(); - let svr_cfg = svr_cfg.clone(); - let mut tx = tx.clone(); - let assoc = async move { - match udp_associate(context, svr_cfg, pkt, src).await { - Ok(pkt) => { - if let Err(..) = tx.send((src, pkt)).await { - error!("UDP packet channel closed"); - } - } - Err(err) => { - error!("Error occurs in UDP relay: {}", err); - } + let (addr, body) = match parse_packet(pkt).await { + Ok((addr, body)) => (addr, body), + Err(err) => { + error!("Failed to parse UDP associate packet, error: {}", err); + continue; } }; - tokio::spawn(assoc); + trace!( + "UDP relay request {} -> {}, payload length {} bytes", + src, + addr, + body.len() + ); + + // Pick a server + let svr_cfg = balancer.pick_server(); + + // Check or (re)create an association + loop { + let retry = { + let assoc = match assoc_map.entry(src.to_string()) { + Entry::Occupied(oc) => oc.into_mut(), + Entry::Vacant(vc) => vc.insert( + UdpAssociation::associate(svr_cfg.clone(), src, tx.clone()) + .await + .expect("Failed to create udp association"), + ), + }; + + !assoc.send(&addr, &body[..]).await + }; + + if retry { + assoc_map.remove(&src.to_string()); + } else { + break; + } + } } } diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs index 1e036b4a..10b062df 100644 --- a/src/relay/udprelay/server.rs +++ b/src/relay/udprelay/server.rs @@ -8,7 +8,7 @@ use std::time::Duration; use bytes::BytesMut; use futures::stream::FuturesUnordered; use futures::{self, StreamExt}; -use log::{debug, error, info}; +use log::{debug, error, info, trace}; use lru_time_cache::{Entry, LruCache}; use tokio; use tokio::net::udp::{RecvHalf, SendHalf}; @@ -24,75 +24,9 @@ use super::MAXIMUM_UDP_PAYLOAD_SIZE; const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); -#[allow(unused_variables)] // `context` is only used if trust-dns is enabled -async fn udp_associate( - context: SharedContext, - svr_cfg: Arc, - decrypted_pkt: Vec, - src: SocketAddr, -) -> io::Result> { - // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD - let mut cur = Cursor::new(decrypted_pkt); - - let addr = Address::read_from(&mut cur).await?; - - // Take out internal buffer for optimizing one byte copy - let header_len = cur.position() as usize; - let decrypted_pkt = cur.into_inner(); - let body = &decrypted_pkt[header_len..]; - - debug!("UDP ASSOCIATE {} -> {}, payload length {} bytes", src, addr, body.len()); - - // FIXME: Create one UdpSocket for one associate - let local_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0); - let mut remote_udp = UdpSocket::bind(&local_addr).await?; - - let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); - - // Writes body to remote - let send_len = match addr { - Address::SocketAddress(ref remote_addr) => { - try_timeout(remote_udp.send_to(&body, remote_addr), Some(timeout)).await? - } - #[cfg(feature = "trust-dns")] - Address::DomainNameAddress(ref dname, port) => { - use crate::relay::dns_resolver::resolve; - - let vec_ipaddr = resolve(context, dname, port, false).await?; - assert!(!vec_ipaddr.is_empty()); - - try_timeout(remote_udp.send_to(&body, &vec_ipaddr[0]), Some(timeout)).await? - } - #[cfg(not(feature = "trust-dns"))] - Address::DomainNameAddress(ref dname, port) => { - try_timeout(remote_udp.send_to(&body, (dname.as_str(), port)), Some(timeout)).await? - } - }; - assert_eq!(body.len(), send_len); - - // 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 = try_timeout(remote_udp.recv(&mut remote_buf), Some(timeout)).await?; - - // Making response packet, SERVER -> CLIENT: ADDRESS + PAYLOAD - let mut send_buf = Vec::new(); - addr.write_to_buf(&mut send_buf); - send_buf.extend_from_slice(&remote_buf[..remote_recv_len]); - - debug!( - "UDP ASSOCIATE {} <- {}, payload length {} bytes", - src, - addr, - send_buf.len() - ); - - Ok(send_buf) -} - // Represent a UDP association struct UdpAssociation { - tx: mpsc::Sender<(Address, Vec)>, + tx: mpsc::Sender>, } impl UdpAssociation { @@ -100,7 +34,7 @@ impl UdpAssociation { async fn associate( svr_cfg: Arc, src_addr: SocketAddr, - mut response_tx: mpsc::Sender<(SocketAddr, Vec)>, + mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>, ) -> io::Result { debug!("Created UDP Association for {}", src_addr); @@ -110,7 +44,7 @@ impl UdpAssociation { // Create a channel for sending packets to remote // FIXME: Channel size 1024? - let (tx, mut rx) = mpsc::channel::<(Address, Vec)>(1024); + let (tx, mut rx) = mpsc::channel::>(1024); // Splits socket into sender and receiver let (mut receiver, mut sender) = remote_udp.split(); @@ -118,18 +52,13 @@ impl UdpAssociation { let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT); // local -> remote + let c_svr_cfg = svr_cfg.clone(); tokio::spawn(async move { - while let Some((addr, pkt)) = rx.recv().await { - debug!( - "UDP ASSOCIATE {} -> {}, payload length {} bytes", - src_addr, - addr, - pkt.len(), - ); - + while let Some(pkt) = rx.recv().await { // pkt is already a raw packet, so just send it - if let Err(err) = UdpAssociation::relay_l2r(&addr, &mut sender, &pkt[..], timeout).await { - error!("Failed to send packet to {}, error: {}", addr, err); + if let Err(err) = UdpAssociation::relay_l2r(src_addr, &mut sender, &pkt[..], timeout, &*c_svr_cfg).await + { + error!("Failed to relay packet, {} -> ..., error: {}", src_addr, err); // FIXME: Ignore? Or how to deal with it? } @@ -142,8 +71,10 @@ impl UdpAssociation { tokio::spawn(async move { loop { // Read and send back to source - if let Err(err) = UdpAssociation::relay_r2l(src_addr, &mut receiver, timeout, &mut response_tx).await { - error!("Failed to receive packet, {} <- .., error: {}", src_addr, err); + if let Err(err) = + UdpAssociation::relay_r2l(src_addr, &mut receiver, timeout, &mut response_tx, &*svr_cfg).await + { + error!("Failed to relay packet, {} <- .., error: {}", src_addr, err); break; } } @@ -155,10 +86,48 @@ impl UdpAssociation { } /// Relay packets from local to remote - async fn relay_l2r(addr: &Address, remote_udp: &mut SendHalf, pkt: &[u8], timeout: Duration) -> io::Result<()> { + async fn relay_l2r( + src: SocketAddr, + remote_udp: &mut SendHalf, + pkt: &[u8], + timeout: Duration, + svr_cfg: &ServerConfig, + ) -> io::Result<()> { + // First of all, decrypt payload CLIENT -> SERVER + let decrypted_pkt = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), pkt) { + Ok(Some(pkt)) => pkt, + Ok(None) => { + error!("Failed to decrypt pkt in UDP relay, packet too short"); + let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short"); + return Err(err); + } + Err(err) => { + error!("Failed to decrypt pkt in UDP relay: {}", err); + let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed"); + return Err(err); + } + }; + + // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD + let mut cur = Cursor::new(decrypted_pkt); + + let addr = Address::read_from(&mut cur).await?; + + // Take out internal buffer for optimizing one byte copy + let header_len = cur.position() as usize; + let decrypted_pkt = cur.into_inner(); + let body = &decrypted_pkt[header_len..]; + let send_len = match addr { Address::SocketAddress(ref remote_addr) => { - try_timeout(remote_udp.send_to(pkt, remote_addr), Some(timeout)).await? + debug!( + "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes", + src, + addr, + remote_addr, + body.len() + ); + try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await? } #[cfg(feature = "trust-dns")] Address::DomainNameAddress(ref dname, port) => { @@ -167,16 +136,28 @@ impl UdpAssociation { let vec_ipaddr = resolve(context, dname, port, false).await?; assert!(!vec_ipaddr.is_empty()); - try_timeout(remote_udp.send_to(pkt, &vec_ipaddr[0]), Some(timeout)).await? + let remote_addr = &vec_ipaddr[0]; + + debug!( + "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes", + src, + addr, + remote_addr, + body.len() + ); + + try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await? } #[cfg(not(feature = "trust-dns"))] Address::DomainNameAddress(ref dname, port) => { - // try_timeout(remote_udp.send_to(pkt, (dname.as_str(), port)), Some(timeout)).await? + debug!("UDP ASSOCIATE {} -> {}, payload length {} bytes", src, addr, body.len()); + + // try_timeout(remote_udp.send_to(body, (dname.as_str(), port)), Some(timeout)).await? unimplemented!("tokio's UdpSocket SendHalf doesn't support ToSocketAddrs"); } }; - assert_eq!(pkt.len(), send_len); + assert_eq!(body.len(), send_len); Ok(()) } @@ -186,18 +167,32 @@ impl UdpAssociation { src_addr: SocketAddr, remote_udp: &mut RecvHalf, timeout: Duration, - response_tx: &mut mpsc::Sender<(SocketAddr, Vec)>, + 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 = try_timeout(remote_udp.recv(&mut remote_buf), Some(timeout)).await?; + let (remote_recv_len, remote_addr) = try_timeout(remote_udp.recv_from(&mut remote_buf), Some(timeout)).await?; + + debug!( + "UDP ASSOCIATE {} <- {}, payload length {} bytes", + src_addr, remote_addr, remote_recv_len + ); + + // FIXME: The Address should be the Address that client sent + let addr = Address::SocketAddress(remote_addr); + + // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD + let mut send_buf = Vec::new(); + addr.write_to_buf(&mut send_buf); + send_buf.extend_from_slice(&remote_buf[..remote_recv_len]); + + let mut encrypt_buf = BytesMut::new(); + encrypt_payload(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, remote_buf[..remote_recv_len].to_vec())) - .await - { + if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await { error!("Failed to send packet into response channel, error: {}", err); // FIXME: What to do? Ignore? @@ -206,8 +201,8 @@ impl UdpAssociation { Ok(()) } - async fn send(&mut self, remote_addr: &Address, pkt: &[u8]) -> bool { - match self.tx.send((remote_addr.clone(), pkt.to_vec())).await { + 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); @@ -217,29 +212,18 @@ impl UdpAssociation { } } -async fn listen(context: SharedContext, svr_cfg: Arc) -> io::Result<()> { +async fn listen(_context: SharedContext, svr_cfg: Arc) -> io::Result<()> { let listen_addr = *svr_cfg.addr().listen_addr(); info!("ShadowSocks UDP listening on {}", listen_addr); let listener = UdpSocket::bind(&listen_addr).await?; let (mut r, mut w) = listener.split(); - let svr_cfg_cloned = svr_cfg.clone(); - // FIXME: Channel size 1024? - let (tx, mut rx) = mpsc::channel::<(SocketAddr, Vec)>(1024); + let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024); tokio::spawn(async move { - let svr_cfg = svr_cfg_cloned; - while let Some((src, pkt)) = rx.recv().await { - // Encrypts - let mut response_pkt = BytesMut::new(); - if let Err(err) = encrypt_payload(svr_cfg.method(), svr_cfg.key(), &pkt, &mut response_pkt) { - error!("UDP packet encrypt failed, err: {:?}", err); - continue; - } - - if let Err(err) = w.send_to(&response_pkt, &src).await { + if let Err(err) = w.send_to(&pkt, &src).await { error!("UDP packet send failed, err: {:?}", err); break; } @@ -259,28 +243,7 @@ async fn listen(context: SharedContext, svr_cfg: Arc) -> io::Resul // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded. let pkt = &pkt_buf[..recv_len]; - // First of all, decrypt payload CLIENT -> SERVER - let decrypted_pkt = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), pkt) { - Ok(Some(pkt)) => pkt, - Ok(None) => { - error!("Failed to decrypt pkt in UDP relay, packet too short"); - continue; - } - Err(err) => { - error!("Failed to decrypt pkt in UDP relay: {}", err); - continue; - } - }; - - // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD - let mut cur = Cursor::new(decrypted_pkt); - - let addr = Address::read_from(&mut cur).await?; - - // Take out internal buffer for optimizing one byte copy - let header_len = cur.position() as usize; - let decrypted_pkt = cur.into_inner(); - let body = &decrypted_pkt[header_len..]; + trace!("Received UDP packet from {}, length {} bytes", src, recv_len); // Check or (re)create an association loop { @@ -294,7 +257,7 @@ async fn listen(context: SharedContext, svr_cfg: Arc) -> io::Resul ), }; - !assoc.send(&addr, body).await + !assoc.send(pkt).await }; if retry {