starred/shadowsocks-rust
1
0
Fork 0
mirror of https://github.com/shadowsocks/shadowsocks-rust.git synced 2026-11-15 08:15:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

UDP associations sending response directly to socket instead of channel

Browse Source
This commit is contained in:
zonyitoo
2020-11-15 15:55:58 +08:00
parent 5bf36c9ef9
commit 174e50e18e
5 changed files with 180 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/relay/udprelay/association.rs
+78 -5
View File
@@ -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<K> = AssociationManager<K, ProxyAssociation>;
/// 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<ServerAssociationKey>,
flow_stat: SharedServerFlowStatistic,
tx: Arc<UdpSocket>,
}
impl ServerProxyHandler {
/// Create a new ServerProxyHandler
pub fn new(
src_addr: SocketAddr,
assoc_manager: ServerAssociationManager<ServerAssociationKey>,
flow_stat: SharedServerFlowStatistic,
tx: Arc<UdpSocket>,
) -> 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<ServerAssociation> {
// 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?
src/relay/udprelay/redir_local.rs
+14 -2
View File
@@ -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;
src/relay/udprelay/server.rs
+6 -54
View File
@@ -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;
src/relay/udprelay/socks5_local.rs
+41 -38
View File
@@ -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<u8>)>,
cache_key: String,
assoc_manager: ProxyAssociationManager<String>,
tx: Arc<UdpSocket>,
}
impl ProxyHandler {
fn new(src_addr: SocketAddr, assoc_manager: ProxyAssociationManager<String>, tx: Arc<UdpSocket>) -> 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<u8>) -> 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<u8>)>(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
})
src/relay/udprelay/tunnel_local.rs
+41 -38
View File
@@ -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<u8>)>,
cache_key: String,
assoc_manager: ProxyAssociationManager<String>,
tx: Arc<UdpSocket>,
}
impl ProxyHandler {
fn new(src_addr: SocketAddr, assoc_manager: ProxyAssociationManager<String>, tx: Arc<UdpSocket>) -> 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<u8>) -> 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<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,
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<u8>)>(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
})