UDP relay local implements association properly (#168)

This commit is contained in:
zonyitoo
2019-12-01 19:36:10 +08:00
parent 9b6669d0bd
commit cdb3210434
2 changed files with 315 additions and 245 deletions
+220 -113
View File
@@ -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<ServerConfig>,
pkt: Vec<u8>,
src: SocketAddr,
) -> io::Result<Vec<u8>> {
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<u8>)> {
// 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<u8>)>,
}
debug!(
"UDP ASSOCIATE {} -> {}, payload length {} bytes",
src,
addr,
payload.len()
);
impl UdpAssociation {
/// Create an association with addr
async fn associate(
svr_cfg: Arc<ServerConfig>,
src_addr: SocketAddr,
mut response_tx: mpsc::Sender<(SocketAddr, Vec<u8>)>,
) -> io::Result<UdpAssociation> {
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<u8>)>(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<u8>)>,
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;
}
}
}
}
+95 -132
View File
@@ -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<ServerConfig>,
decrypted_pkt: Vec<u8>,
src: SocketAddr,
) -> io::Result<Vec<u8>> {
// 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<u8>)>,
tx: mpsc::Sender<Vec<u8>>,
}
impl UdpAssociation {
@@ -100,7 +34,7 @@ impl UdpAssociation {
async fn associate(
svr_cfg: Arc<ServerConfig>,
src_addr: SocketAddr,
mut response_tx: mpsc::Sender<(SocketAddr, Vec<u8>)>,
mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
) -> io::Result<UdpAssociation> {
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<u8>)>(1024);
let (tx, mut rx) = mpsc::channel::<Vec<u8>>(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<u8>)>,
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<ServerConfig>) -> io::Result<()> {
async fn listen(_context: SharedContext, svr_cfg: Arc<ServerConfig>) -> 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<u8>)>(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<ServerConfig>) -> 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<ServerConfig>) -> io::Resul
),
};
!assoc.send(&addr, body).await
!assoc.send(pkt).await
};
if retry {