sujiacong
2 years ago
committed by
GitHub
GitHub
6 changed files with 956 additions and 7 deletions
@ -0,0 +1,210 @@ |
|||||
|
use socks5_impl::protocol::{AddressType, AsyncStreamOperation}; |
||||
|
use std::{net::SocketAddr, sync::Arc}; |
||||
|
use tokio::{ |
||||
|
io::AsyncWriteExt, |
||||
|
net::{ |
||||
|
tcp::{ReadHalf, WriteHalf}, |
||||
|
UdpSocket, |
||||
|
}, |
||||
|
sync::mpsc::{self, Receiver, Sender}, |
||||
|
}; |
||||
|
use tun2proxy::{ |
||||
|
udpgw::{Packet, UDPGW_FLAG_KEEPALIVE}, |
||||
|
ArgVerbosity, Result, |
||||
|
}; |
||||
|
|
||||
|
pub(crate) const CLIENT_DISCONNECT_TIMEOUT: tokio::time::Duration = std::time::Duration::from_secs(60); |
||||
|
|
||||
|
#[derive(Debug, Clone)] |
||||
|
pub struct Client { |
||||
|
addr: SocketAddr, |
||||
|
last_activity: std::time::Instant, |
||||
|
} |
||||
|
|
||||
|
impl Client { |
||||
|
pub fn new(addr: SocketAddr) -> Self { |
||||
|
let last_activity = std::time::Instant::now(); |
||||
|
Self { addr, last_activity } |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
#[derive(Debug, Clone, clap::Parser)] |
||||
|
pub struct UdpGwArgs { |
||||
|
/// UDP gateway listen address
|
||||
|
#[arg(short, long, value_name = "IP:PORT", default_value = "127.0.0.1:7300")] |
||||
|
pub listen_addr: SocketAddr, |
||||
|
|
||||
|
/// UDP mtu
|
||||
|
#[arg(short = 'm', long, value_name = "udp mtu", default_value = "10240")] |
||||
|
pub udp_mtu: u16, |
||||
|
|
||||
|
/// UDP timeout in seconds
|
||||
|
#[arg(short = 't', long, value_name = "seconds", default_value = "3")] |
||||
|
pub udp_timeout: u64, |
||||
|
|
||||
|
/// Daemonize for unix family or run as Windows service
|
||||
|
#[cfg(unix)] |
||||
|
#[arg(long)] |
||||
|
pub daemonize: bool, |
||||
|
|
||||
|
/// Verbosity level
|
||||
|
#[arg(short, long, value_name = "level", value_enum, default_value = "info")] |
||||
|
pub verbosity: ArgVerbosity, |
||||
|
} |
||||
|
|
||||
|
impl UdpGwArgs { |
||||
|
#[allow(clippy::let_and_return)] |
||||
|
pub fn parse_args() -> Self { |
||||
|
use clap::Parser; |
||||
|
Self::parse() |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
async fn send_error(tx: Sender<Packet>, conn_id: u16) { |
||||
|
let error_packet = Packet::build_error_packet(conn_id); |
||||
|
if let Err(e) = tx.send(error_packet).await { |
||||
|
log::error!("send error response error {:?}", e); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
async fn send_keepalive_response(tx: Sender<Packet>, conn_id: u16) { |
||||
|
let keepalive_packet = Packet::build_keepalive_packet(conn_id); |
||||
|
if let Err(e) = tx.send(keepalive_packet).await { |
||||
|
log::error!("send keepalive response error {:?}", e); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
/// Send data field of packet from client to destination server and receive response,
|
||||
|
/// then wrap response data to the packet's data field and send packet back to client.
|
||||
|
async fn process_udp(client: SocketAddr, udp_mtu: u16, udp_timeout: u64, tx: Sender<Packet>, mut packet: Packet) -> Result<()> { |
||||
|
let Some(dst_addr) = &packet.address else { |
||||
|
log::error!("client {} udp request address is None", client); |
||||
|
return Ok(()); |
||||
|
}; |
||||
|
let std_sock = if dst_addr.get_type() == AddressType::IPv6 { |
||||
|
std::net::UdpSocket::bind("[::]:0")? |
||||
|
} else { |
||||
|
std::net::UdpSocket::bind("0.0.0.0:0")? |
||||
|
}; |
||||
|
std_sock.set_nonblocking(true)?; |
||||
|
#[cfg(unix)] |
||||
|
nix::sys::socket::setsockopt(&std_sock, nix::sys::socket::sockopt::ReuseAddr, &true)?; |
||||
|
let socket = UdpSocket::from_std(std_sock)?; |
||||
|
use std::net::ToSocketAddrs; |
||||
|
let Some(dst_addr) = dst_addr.to_socket_addrs()?.next() else { |
||||
|
log::error!("client {} udp request address to_socket_addrs", client); |
||||
|
return Ok(()); |
||||
|
}; |
||||
|
// 1. send udp data to destination server
|
||||
|
socket.send_to(&packet.data, &dst_addr).await?; |
||||
|
packet.data.resize(udp_mtu as usize, 0); |
||||
|
// 2. receive response from destination server
|
||||
|
let (len, _addr) = tokio::time::timeout(tokio::time::Duration::from_secs(udp_timeout), socket.recv_from(&mut packet.data)) |
||||
|
.await |
||||
|
.map_err(std::io::Error::from)??; |
||||
|
packet.data.truncate(len); |
||||
|
// 3. send response back to client
|
||||
|
use std::io::{Error, ErrorKind::BrokenPipe}; |
||||
|
tx.send(packet).await.map_err(|e| Error::new(BrokenPipe, e))?; |
||||
|
Ok(()) |
||||
|
} |
||||
|
|
||||
|
async fn process_client_udp_req(args: &UdpGwArgs, tx: Sender<Packet>, mut client: Client, mut reader: ReadHalf<'_>) -> std::io::Result<()> { |
||||
|
let udp_timeout = args.udp_timeout; |
||||
|
let udp_mtu = args.udp_mtu; |
||||
|
|
||||
|
loop { |
||||
|
// 1. read udpgw packet from client
|
||||
|
let res = tokio::time::timeout(tokio::time::Duration::from_secs(2), Packet::retrieve_from_async_stream(&mut reader)).await; |
||||
|
let packet = match res { |
||||
|
Ok(Ok(packet)) => packet, |
||||
|
Ok(Err(e)) => { |
||||
|
log::error!("client {} retrieve_from_async_stream \"{}\"", client.addr, e); |
||||
|
break; |
||||
|
} |
||||
|
Err(e) => { |
||||
|
if client.last_activity.elapsed() >= CLIENT_DISCONNECT_TIMEOUT { |
||||
|
log::debug!("client {} last_activity elapsed \"{e}\"", client.addr); |
||||
|
break; |
||||
|
} |
||||
|
continue; |
||||
|
} |
||||
|
}; |
||||
|
client.last_activity = std::time::Instant::now(); |
||||
|
|
||||
|
let flags = packet.header.flags; |
||||
|
let conn_id = packet.header.conn_id; |
||||
|
if flags & UDPGW_FLAG_KEEPALIVE != 0 { |
||||
|
log::trace!("client {} send keepalive", client.addr); |
||||
|
// 2. if keepalive packet, do nothing, send keepalive response to client
|
||||
|
send_keepalive_response(tx.clone(), conn_id).await; |
||||
|
continue; |
||||
|
} |
||||
|
log::trace!("client {} received udp data {}", client.addr, packet); |
||||
|
|
||||
|
// 3. process client udpgw packet in a new task
|
||||
|
let tx = tx.clone(); |
||||
|
tokio::spawn(async move { |
||||
|
if let Err(e) = process_udp(client.addr, udp_mtu, udp_timeout, tx.clone(), packet).await { |
||||
|
send_error(tx, conn_id).await; |
||||
|
log::error!("client {} process udp function {}", client.addr, e); |
||||
|
} |
||||
|
}); |
||||
|
} |
||||
|
Ok(()) |
||||
|
} |
||||
|
|
||||
|
async fn write_to_client(addr: SocketAddr, mut writer: WriteHalf<'_>, mut rx: Receiver<Packet>) -> std::io::Result<()> { |
||||
|
loop { |
||||
|
use std::io::{Error, ErrorKind::BrokenPipe}; |
||||
|
let packet = rx.recv().await.ok_or(Error::new(BrokenPipe, "recv error"))?; |
||||
|
log::trace!("send response to client {} with {}", addr, packet); |
||||
|
let data: Vec<u8> = packet.into(); |
||||
|
let _r = writer.write(&data).await?; |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
#[tokio::main] |
||||
|
async fn main() -> Result<()> { |
||||
|
let args = Arc::new(UdpGwArgs::parse_args()); |
||||
|
|
||||
|
let tcp_listener = tokio::net::TcpListener::bind(args.listen_addr).await?; |
||||
|
|
||||
|
let default = format!("{:?}", args.verbosity); |
||||
|
|
||||
|
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or(default)).init(); |
||||
|
|
||||
|
log::info!("{} {} starting...", module_path!(), env!("CARGO_PKG_VERSION")); |
||||
|
log::info!("UDP Gateway Server running at {}", args.listen_addr); |
||||
|
|
||||
|
#[cfg(unix)] |
||||
|
if args.daemonize { |
||||
|
let stdout = std::fs::File::create("/tmp/udpgw.out")?; |
||||
|
let stderr = std::fs::File::create("/tmp/udpgw.err")?; |
||||
|
let daemonize = daemonize::Daemonize::new() |
||||
|
.working_directory("/tmp") |
||||
|
.umask(0o777) |
||||
|
.stdout(stdout) |
||||
|
.stderr(stderr) |
||||
|
.privileged_action(|| "Executed before drop privileges"); |
||||
|
let _ = daemonize |
||||
|
.start() |
||||
|
.map_err(|e| format!("Failed to daemonize process, error:{:?}", e))?; |
||||
|
} |
||||
|
|
||||
|
loop { |
||||
|
let (mut tcp_stream, addr) = tcp_listener.accept().await?; |
||||
|
let client = Client::new(addr); |
||||
|
log::info!("client {} connected", addr); |
||||
|
let params = args.clone(); |
||||
|
tokio::spawn(async move { |
||||
|
let (tx, rx) = mpsc::channel::<Packet>(100); |
||||
|
let (tcp_read_stream, tcp_write_stream) = tcp_stream.split(); |
||||
|
let res = tokio::select! { |
||||
|
v = process_client_udp_req(¶ms, tx, client, tcp_read_stream) => v, |
||||
|
v = write_to_client(addr, tcp_write_stream, rx) => v, |
||||
|
}; |
||||
|
log::info!("client {} disconnected with {:?}", addr, res); |
||||
|
}); |
||||
|
} |
||||
|
} |
||||
@ -0,0 +1,542 @@ |
|||||
|
use crate::error::Result; |
||||
|
use socks5_impl::protocol::{Address, AsyncStreamOperation, BufMut, StreamOperation}; |
||||
|
use std::{collections::VecDeque, hash::Hash, net::SocketAddr, sync::atomic::Ordering::Relaxed}; |
||||
|
use tokio::{ |
||||
|
io::{AsyncReadExt, AsyncWriteExt}, |
||||
|
net::{ |
||||
|
tcp::{OwnedReadHalf, OwnedWriteHalf}, |
||||
|
TcpStream, |
||||
|
}, |
||||
|
sync::Mutex, |
||||
|
time::{sleep, Duration}, |
||||
|
}; |
||||
|
|
||||
|
pub(crate) const UDPGW_LENGTH_FIELD_SIZE: usize = std::mem::size_of::<u16>(); |
||||
|
pub(crate) const UDPGW_MAX_CONNECTIONS: u16 = 100; |
||||
|
pub(crate) const UDPGW_KEEPALIVE_TIME: tokio::time::Duration = std::time::Duration::from_secs(10); |
||||
|
|
||||
|
pub const UDPGW_FLAG_KEEPALIVE: u8 = 0x01; |
||||
|
pub const UDPGW_FLAG_ERR: u8 = 0x20; |
||||
|
pub const UDPGW_FLAG_DATA: u8 = 0x02; |
||||
|
|
||||
|
static TCP_COUNTER: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0); |
||||
|
|
||||
|
/// UDP Gateway Packet Format
|
||||
|
///
|
||||
|
/// The format is referenced from SOCKS5 packet format, with additional flags and connection ID fields.
|
||||
|
///
|
||||
|
/// `LEN`: This field is indicated the length of the packet, not including the length field itself.
|
||||
|
///
|
||||
|
/// `FLAGS`: This field is used to indicate the packet type. The flags are defined as follows:
|
||||
|
/// - `0x01`: Keepalive packet without address and data
|
||||
|
/// - `0x20`: Error packet without address and data
|
||||
|
/// - `0x02`: Data packet with address and data
|
||||
|
///
|
||||
|
/// `CONN_ID`: This field is used to indicate the unique connection ID for the packet.
|
||||
|
///
|
||||
|
/// `ATYP` & `DST.ADDR` & `DST.PORT`: This fields are used to indicate the remote address and port.
|
||||
|
/// It can be either an IPv4 address, an IPv6 address, or a domain name, depending on the `ATYP` field.
|
||||
|
/// The address format directly uses the address format of the [SOCKS5](https://datatracker.ietf.org/doc/html/rfc1928#section-4) protocol.
|
||||
|
/// - `ATYP`: Address Type, 1 byte, indicating the type of address ( 0x01-IPv4, 0x04-IPv6, or 0x03-domain name )
|
||||
|
/// - `DST.ADDR`: Destination Address. If `ATYP` is 0x01 or 0x04, it is 4 or 16 bytes of IP address;
|
||||
|
/// If `ATYP` is 0x03, it is a domain name, `DST.ADDR` is a variable length field,
|
||||
|
/// it begins with a 1-byte length field and then the domain name without null-termination,
|
||||
|
/// since the length field is 1 byte, the maximum length of the domain name is 255 bytes.
|
||||
|
/// - `DST.PORT`: Destination Port, 2 bytes, the port number of the destination address.
|
||||
|
///
|
||||
|
/// `DATA`: The data field, a variable length field, the length is determined by the `LEN` field.
|
||||
|
///
|
||||
|
/// All the digits fields are in big-endian byte order.
|
||||
|
///
|
||||
|
/// ```plain
|
||||
|
/// +-----+ +-------+---------+ +------+----------+----------+ +----------+
|
||||
|
/// | LEN | | FLAGS | CONN_ID | | ATYP | DST.ADDR | DST.PORT | | DATA |
|
||||
|
/// +-----+ +-------+---------+ +------+----------+----------+ +----------+
|
||||
|
/// | 2 | | 1 | 2 | | 1 | Variable | 2 | | Variable |
|
||||
|
/// +-----+ +-------+---------+ +------+----------+----------+ +----------+
|
||||
|
/// ```
|
||||
|
#[derive(Debug, Clone, PartialEq, Eq, Hash)] |
||||
|
pub struct Packet { |
||||
|
pub header: UdpgwHeader, |
||||
|
pub address: Option<Address>, |
||||
|
pub data: Vec<u8>, |
||||
|
} |
||||
|
|
||||
|
impl std::fmt::Display for Packet { |
||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
||||
|
let addr = self.address.as_ref().map_or("None".to_string(), |addr| addr.to_string()); |
||||
|
let len = self.data.len(); |
||||
|
write!(f, "Packet {{ {}, address: {}, payload length: {} }}", self.header, addr, len) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl From<Packet> for Vec<u8> { |
||||
|
fn from(packet: Packet) -> Vec<u8> { |
||||
|
(&packet).into() |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl From<&Packet> for Vec<u8> { |
||||
|
fn from(packet: &Packet) -> Vec<u8> { |
||||
|
let mut bytes: Vec<u8> = vec![]; |
||||
|
packet.write_to_buf(&mut bytes); |
||||
|
bytes |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl TryFrom<&[u8]> for Packet { |
||||
|
type Error = std::io::Error; |
||||
|
|
||||
|
fn try_from(value: &[u8]) -> std::result::Result<Self, Self::Error> { |
||||
|
if value.len() < UDPGW_LENGTH_FIELD_SIZE { |
||||
|
return Err(std::io::ErrorKind::InvalidData.into()); |
||||
|
} |
||||
|
let mut iter = std::io::Cursor::new(value); |
||||
|
use tokio_util::bytes::Buf; |
||||
|
let length = iter.get_u16(); |
||||
|
if value.len() < length as usize + UDPGW_LENGTH_FIELD_SIZE { |
||||
|
return Err(std::io::ErrorKind::InvalidData.into()); |
||||
|
} |
||||
|
let header = UdpgwHeader::retrieve_from_stream(&mut iter)?; |
||||
|
let address = if header.flags & UDPGW_FLAG_DATA != 0 { |
||||
|
Some(Address::retrieve_from_stream(&mut iter)?) |
||||
|
} else { |
||||
|
None |
||||
|
}; |
||||
|
Ok(Packet::new(header, address, iter.chunk())) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl Packet { |
||||
|
pub fn new(header: UdpgwHeader, address: Option<Address>, data: &[u8]) -> Self { |
||||
|
let data = data.to_vec(); |
||||
|
Packet { header, address, data } |
||||
|
} |
||||
|
|
||||
|
pub fn build_keepalive_packet(conn_id: u16) -> Self { |
||||
|
Packet::new(UdpgwHeader::new(UDPGW_FLAG_KEEPALIVE, conn_id), None, &[]) |
||||
|
} |
||||
|
|
||||
|
pub fn build_error_packet(conn_id: u16) -> Self { |
||||
|
Packet::new(UdpgwHeader::new(UDPGW_FLAG_ERR, conn_id), None, &[]) |
||||
|
} |
||||
|
|
||||
|
pub fn build_packet_from_address(conn_id: u16, remote_addr: &Address, data: &[u8]) -> std::io::Result<Self> { |
||||
|
use socks5_impl::protocol::Address::{DomainAddress, SocketAddress}; |
||||
|
let packet = match remote_addr { |
||||
|
SocketAddress(addr) => Packet::build_ip_packet(conn_id, *addr, data), |
||||
|
DomainAddress(domain, port) => Packet::build_domain_packet(conn_id, *port, domain, data)?, |
||||
|
}; |
||||
|
Ok(packet) |
||||
|
} |
||||
|
|
||||
|
pub fn build_ip_packet(conn_id: u16, remote_addr: SocketAddr, data: &[u8]) -> Self { |
||||
|
let addr: Address = remote_addr.into(); |
||||
|
Packet::new(UdpgwHeader::new(UDPGW_FLAG_DATA, conn_id), Some(addr), data) |
||||
|
} |
||||
|
|
||||
|
pub fn build_domain_packet(conn_id: u16, port: u16, domain: &str, data: &[u8]) -> std::io::Result<Self> { |
||||
|
if domain.len() > 255 { |
||||
|
return Err(std::io::ErrorKind::InvalidInput.into()); |
||||
|
} |
||||
|
let addr = Address::from((domain, port)); |
||||
|
Ok(Packet::new(UdpgwHeader::new(UDPGW_FLAG_DATA, conn_id), Some(addr), data)) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl StreamOperation for Packet { |
||||
|
fn retrieve_from_stream<R>(stream: &mut R) -> std::io::Result<Self> |
||||
|
where |
||||
|
R: std::io::Read, |
||||
|
Self: Sized, |
||||
|
{ |
||||
|
let mut buf = [0; UDPGW_LENGTH_FIELD_SIZE]; |
||||
|
stream.read_exact(&mut buf)?; |
||||
|
let length = u16::from_be_bytes(buf) as usize; |
||||
|
let header = UdpgwHeader::retrieve_from_stream(stream)?; |
||||
|
let address = if header.flags & UDPGW_FLAG_DATA != 0 { |
||||
|
Some(Address::retrieve_from_stream(stream)?) |
||||
|
} else { |
||||
|
None |
||||
|
}; |
||||
|
let read_len = header.len() + address.as_ref().map_or(0, |addr| addr.len()); |
||||
|
if length < read_len { |
||||
|
return Err(std::io::ErrorKind::InvalidData.into()); |
||||
|
} |
||||
|
let mut data = vec![0; length - read_len]; |
||||
|
stream.read_exact(&mut data)?; |
||||
|
Ok(Packet::new(header, address, &data)) |
||||
|
} |
||||
|
|
||||
|
fn write_to_buf<B: BufMut>(&self, buf: &mut B) { |
||||
|
let len = self.len() - UDPGW_LENGTH_FIELD_SIZE; |
||||
|
buf.put_u16(len as u16); |
||||
|
self.header.write_to_buf(buf); |
||||
|
if let Some(addr) = &self.address { |
||||
|
addr.write_to_buf(buf); |
||||
|
} |
||||
|
buf.put_slice(&self.data); |
||||
|
} |
||||
|
|
||||
|
fn len(&self) -> usize { |
||||
|
UDPGW_LENGTH_FIELD_SIZE + self.header.len() + self.address.as_ref().map_or(0, |addr| addr.len()) + self.data.len() |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
#[async_trait::async_trait] |
||||
|
impl AsyncStreamOperation for Packet { |
||||
|
async fn retrieve_from_async_stream<R>(r: &mut R) -> std::io::Result<Self> |
||||
|
where |
||||
|
R: tokio::io::AsyncRead + Unpin + Send, |
||||
|
Self: Sized, |
||||
|
{ |
||||
|
let mut buf = [0; UDPGW_LENGTH_FIELD_SIZE]; |
||||
|
r.read_exact(&mut buf).await?; |
||||
|
let length = u16::from_be_bytes(buf) as usize; |
||||
|
let header = UdpgwHeader::retrieve_from_async_stream(r).await?; |
||||
|
let address = if header.flags & UDPGW_FLAG_DATA != 0 { |
||||
|
Some(Address::retrieve_from_async_stream(r).await?) |
||||
|
} else { |
||||
|
None |
||||
|
}; |
||||
|
let read_len = header.len() + address.as_ref().map_or(0, |addr| addr.len()); |
||||
|
if length < read_len { |
||||
|
return Err(std::io::ErrorKind::InvalidData.into()); |
||||
|
} |
||||
|
let mut data = vec![0; length - read_len]; |
||||
|
r.read_exact(&mut data).await?; |
||||
|
Ok(Packet::new(header, address, &data)) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] |
||||
|
pub struct UdpgwHeader { |
||||
|
pub flags: u8, |
||||
|
pub conn_id: u16, |
||||
|
} |
||||
|
|
||||
|
impl std::fmt::Display for UdpgwHeader { |
||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
||||
|
let id = self.conn_id; |
||||
|
write!(f, "flags: 0x{:02x}, conn_id: {}", self.flags, id) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl StreamOperation for UdpgwHeader { |
||||
|
fn retrieve_from_stream<R>(stream: &mut R) -> std::io::Result<Self> |
||||
|
where |
||||
|
R: std::io::Read, |
||||
|
Self: Sized, |
||||
|
{ |
||||
|
let mut buf = [0; UdpgwHeader::static_len()]; |
||||
|
stream.read_exact(&mut buf)?; |
||||
|
UdpgwHeader::try_from(&buf[..]) |
||||
|
} |
||||
|
|
||||
|
fn write_to_buf<B: BufMut>(&self, buf: &mut B) { |
||||
|
let bytes: Vec<u8> = self.into(); |
||||
|
buf.put_slice(&bytes); |
||||
|
} |
||||
|
|
||||
|
fn len(&self) -> usize { |
||||
|
Self::static_len() |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
#[async_trait::async_trait] |
||||
|
impl AsyncStreamOperation for UdpgwHeader { |
||||
|
async fn retrieve_from_async_stream<R>(r: &mut R) -> std::io::Result<Self> |
||||
|
where |
||||
|
R: tokio::io::AsyncRead + Unpin + Send, |
||||
|
Self: Sized, |
||||
|
{ |
||||
|
let mut buf = [0; UdpgwHeader::static_len()]; |
||||
|
r.read_exact(&mut buf).await?; |
||||
|
UdpgwHeader::try_from(&buf[..]) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl UdpgwHeader { |
||||
|
pub fn new(flags: u8, conn_id: u16) -> Self { |
||||
|
UdpgwHeader { flags, conn_id } |
||||
|
} |
||||
|
|
||||
|
pub const fn static_len() -> usize { |
||||
|
std::mem::size_of::<u8>() + std::mem::size_of::<u16>() |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl TryFrom<&[u8]> for UdpgwHeader { |
||||
|
type Error = std::io::Error; |
||||
|
|
||||
|
fn try_from(value: &[u8]) -> std::result::Result<Self, Self::Error> { |
||||
|
if value.len() < UdpgwHeader::static_len() { |
||||
|
return Err(std::io::ErrorKind::InvalidData.into()); |
||||
|
} |
||||
|
let conn_id = u16::from_be_bytes([value[1], value[2]]); |
||||
|
Ok(UdpgwHeader { flags: value[0], conn_id }) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl From<&UdpgwHeader> for Vec<u8> { |
||||
|
fn from(header: &UdpgwHeader) -> Vec<u8> { |
||||
|
let mut bytes = vec![0; header.len()]; |
||||
|
bytes[0] = header.flags; |
||||
|
bytes[1..3].copy_from_slice(&header.conn_id.to_be_bytes()); |
||||
|
bytes |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
#[allow(dead_code)] |
||||
|
#[derive(Debug)] |
||||
|
pub(crate) enum UdpGwResponse { |
||||
|
KeepAlive, |
||||
|
Error, |
||||
|
TcpClose, |
||||
|
Data(Packet), |
||||
|
} |
||||
|
|
||||
|
#[derive(Debug)] |
||||
|
pub(crate) struct UdpGwClientStream { |
||||
|
local_addr: String, |
||||
|
writer: Option<OwnedWriteHalf>, |
||||
|
reader: Option<OwnedReadHalf>, |
||||
|
conn_id: u16, |
||||
|
closed: bool, |
||||
|
last_activity: std::time::Instant, |
||||
|
} |
||||
|
|
||||
|
impl Drop for UdpGwClientStream { |
||||
|
fn drop(&mut self) { |
||||
|
TCP_COUNTER.fetch_sub(1, Relaxed); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl UdpGwClientStream { |
||||
|
pub fn close(&mut self) { |
||||
|
self.closed = true; |
||||
|
} |
||||
|
|
||||
|
pub fn get_reader(&mut self) -> Option<OwnedReadHalf> { |
||||
|
self.reader.take() |
||||
|
} |
||||
|
|
||||
|
pub fn set_reader(&mut self, reader: Option<OwnedReadHalf>) { |
||||
|
self.reader = reader; |
||||
|
} |
||||
|
|
||||
|
pub fn set_writer(&mut self, writer: Option<OwnedWriteHalf>) { |
||||
|
self.writer = writer; |
||||
|
} |
||||
|
|
||||
|
pub fn get_writer(&mut self) -> Option<OwnedWriteHalf> { |
||||
|
self.writer.take() |
||||
|
} |
||||
|
|
||||
|
pub fn local_addr(&self) -> &String { |
||||
|
&self.local_addr |
||||
|
} |
||||
|
|
||||
|
pub fn update_activity(&mut self) { |
||||
|
self.last_activity = std::time::Instant::now(); |
||||
|
} |
||||
|
|
||||
|
pub fn is_closed(&mut self) -> bool { |
||||
|
self.closed |
||||
|
} |
||||
|
|
||||
|
pub fn id(&mut self) -> u16 { |
||||
|
self.conn_id |
||||
|
} |
||||
|
|
||||
|
pub fn new_id(&mut self) -> u16 { |
||||
|
self.conn_id += 1; |
||||
|
self.conn_id |
||||
|
} |
||||
|
|
||||
|
pub fn new(tcp_server_stream: TcpStream) -> Self { |
||||
|
let default = "0.0.0.0:0".parse::<SocketAddr>().unwrap(); |
||||
|
let local_addr = tcp_server_stream.local_addr().unwrap_or(default).to_string(); |
||||
|
let (rx, tx) = tcp_server_stream.into_split(); |
||||
|
let writer = tx; |
||||
|
let reader = rx; |
||||
|
TCP_COUNTER.fetch_add(1, Relaxed); |
||||
|
UdpGwClientStream { |
||||
|
local_addr, |
||||
|
reader: Some(reader), |
||||
|
writer: Some(writer), |
||||
|
last_activity: std::time::Instant::now(), |
||||
|
closed: false, |
||||
|
conn_id: 0, |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
#[derive(Debug)] |
||||
|
pub(crate) struct UdpGwClient { |
||||
|
udp_mtu: u16, |
||||
|
max_connections: u16, |
||||
|
udp_timeout: u64, |
||||
|
keepalive_time: Duration, |
||||
|
server_addr: SocketAddr, |
||||
|
server_connections: Mutex<VecDeque<UdpGwClientStream>>, |
||||
|
} |
||||
|
|
||||
|
impl UdpGwClient { |
||||
|
pub fn new(udp_mtu: u16, max_connections: u16, keepalive_time: Duration, udp_timeout: u64, server_addr: SocketAddr) -> Self { |
||||
|
let server_connections = Mutex::new(VecDeque::with_capacity(max_connections as usize)); |
||||
|
UdpGwClient { |
||||
|
udp_mtu, |
||||
|
max_connections, |
||||
|
udp_timeout, |
||||
|
server_addr, |
||||
|
keepalive_time, |
||||
|
server_connections, |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
pub(crate) fn get_udp_mtu(&self) -> u16 { |
||||
|
self.udp_mtu |
||||
|
} |
||||
|
|
||||
|
pub(crate) fn get_udp_timeout(&self) -> u64 { |
||||
|
self.udp_timeout |
||||
|
} |
||||
|
|
||||
|
pub(crate) fn is_full(&self) -> bool { |
||||
|
TCP_COUNTER.load(Relaxed) >= self.max_connections as u32 |
||||
|
} |
||||
|
|
||||
|
pub(crate) async fn get_server_connection(&self) -> Option<UdpGwClientStream> { |
||||
|
self.server_connections.lock().await.pop_front() |
||||
|
} |
||||
|
|
||||
|
pub(crate) async fn release_server_connection(&self, stream: UdpGwClientStream) { |
||||
|
if self.server_connections.lock().await.len() < self.max_connections as usize { |
||||
|
self.server_connections.lock().await.push_back(stream); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
pub(crate) async fn release_server_connection_full( |
||||
|
&self, |
||||
|
mut stream: UdpGwClientStream, |
||||
|
reader: OwnedReadHalf, |
||||
|
writer: OwnedWriteHalf, |
||||
|
) { |
||||
|
if self.server_connections.lock().await.len() < self.max_connections as usize { |
||||
|
stream.set_reader(Some(reader)); |
||||
|
stream.set_writer(Some(writer)); |
||||
|
self.server_connections.lock().await.push_back(stream); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
pub(crate) fn get_server_addr(&self) -> SocketAddr { |
||||
|
self.server_addr |
||||
|
} |
||||
|
|
||||
|
/// Heartbeat task asynchronous function to periodically check and maintain the active state of the server connection.
|
||||
|
pub(crate) async fn heartbeat_task(&self) { |
||||
|
loop { |
||||
|
sleep(self.keepalive_time).await; |
||||
|
if let Some(mut stream) = self.get_server_connection().await { |
||||
|
if stream.last_activity.elapsed() < self.keepalive_time { |
||||
|
self.release_server_connection(stream).await; |
||||
|
continue; |
||||
|
} |
||||
|
|
||||
|
let Some(mut stream_reader) = stream.get_reader() else { |
||||
|
continue; |
||||
|
}; |
||||
|
|
||||
|
let Some(mut stream_writer) = stream.get_writer() else { |
||||
|
continue; |
||||
|
}; |
||||
|
let local_addr = stream_writer.local_addr(); |
||||
|
log::debug!("{:?}:{} send keepalive", local_addr, stream.id()); |
||||
|
let keepalive_packet: Vec<u8> = Packet::build_keepalive_packet(stream.id()).into(); |
||||
|
if let Err(e) = stream_writer.write_all(&keepalive_packet).await { |
||||
|
log::warn!("{:?}:{} send keepalive failed: {}", local_addr, stream.id(), e); |
||||
|
continue; |
||||
|
} |
||||
|
match UdpGwClient::recv_udpgw_packet(self.udp_mtu, 10, &mut stream_reader).await { |
||||
|
Ok(UdpGwResponse::KeepAlive) => { |
||||
|
stream.update_activity(); |
||||
|
self.release_server_connection_full(stream, stream_reader, stream_writer).await; |
||||
|
} |
||||
|
Ok(v) => log::warn!("{:?}:{} keepalive unexpected response: {:?}", local_addr, stream.id(), v), |
||||
|
Err(e) => log::warn!("{:?}:{} keepalive no response, error \"{}\"", local_addr, stream.id(), e), |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
/// Parses the UDP response data.
|
||||
|
pub(crate) fn parse_udp_response(udp_mtu: u16, data: &[u8]) -> Result<UdpGwResponse> { |
||||
|
let packet = Packet::try_from(data)?; |
||||
|
let flags = packet.header.flags; |
||||
|
if flags & UDPGW_FLAG_ERR != 0 { |
||||
|
return Ok(UdpGwResponse::Error); |
||||
|
} |
||||
|
if flags & UDPGW_FLAG_KEEPALIVE != 0 { |
||||
|
return Ok(UdpGwResponse::KeepAlive); |
||||
|
} |
||||
|
if packet.data.len() > udp_mtu as usize { |
||||
|
return Err("too much data".into()); |
||||
|
} |
||||
|
Ok(UdpGwResponse::Data(packet)) |
||||
|
} |
||||
|
|
||||
|
/// Receives a UDP gateway packet.
|
||||
|
///
|
||||
|
/// This function is responsible for receiving packets from the UDP gateway
|
||||
|
///
|
||||
|
/// # Arguments
|
||||
|
/// - `udp_mtu`: The maximum transmission unit size for UDP packets.
|
||||
|
/// - `udp_timeout`: The timeout in seconds for receiving UDP packets.
|
||||
|
/// - `stream`: A mutable reference to the UDP gateway client stream reader.
|
||||
|
///
|
||||
|
/// # Returns
|
||||
|
/// - `Result<UdpGwResponse>`: Returns a result type containing the parsed UDP gateway response, or an error if one occurs.
|
||||
|
pub(crate) async fn recv_udpgw_packet(udp_mtu: u16, udp_timeout: u64, stream: &mut OwnedReadHalf) -> Result<UdpGwResponse> { |
||||
|
let mut data = vec![0; udp_mtu.into()]; |
||||
|
let data_len = tokio::time::timeout(tokio::time::Duration::from_secs(udp_timeout + 2), stream.read(&mut data)) |
||||
|
.await |
||||
|
.map_err(std::io::Error::from)??; |
||||
|
if data_len == 0 { |
||||
|
return Ok(UdpGwResponse::TcpClose); |
||||
|
} |
||||
|
UdpGwClient::parse_udp_response(udp_mtu, &data[..data_len]) |
||||
|
} |
||||
|
|
||||
|
/// Sends a UDP gateway packet.
|
||||
|
///
|
||||
|
/// This function constructs and sends a UDP gateway packet based on the IPv6 enabled status, data length,
|
||||
|
/// remote address, domain (if any), connection ID, and the UDP gateway client writer stream.
|
||||
|
///
|
||||
|
/// # Arguments
|
||||
|
///
|
||||
|
/// * `ipv6_enabled` - Whether IPv6 is enabled
|
||||
|
/// * `data` - The data packet
|
||||
|
/// * `remote_addr` - Remote address
|
||||
|
/// * `conn_id` - Connection ID
|
||||
|
/// * `stream` - UDP gateway client writer stream
|
||||
|
///
|
||||
|
/// # Returns
|
||||
|
///
|
||||
|
/// Returns `Ok(())` if the packet is sent successfully, otherwise returns an error.
|
||||
|
pub(crate) async fn send_udpgw_packet( |
||||
|
ipv6_enabled: bool, |
||||
|
data: &[u8], |
||||
|
remote_addr: &socks5_impl::protocol::Address, |
||||
|
conn_id: u16, |
||||
|
stream: &mut OwnedWriteHalf, |
||||
|
) -> Result<()> { |
||||
|
if !ipv6_enabled && remote_addr.get_type() == socks5_impl::protocol::AddressType::IPv6 { |
||||
|
return Err("ipv6 not support".into()); |
||||
|
} |
||||
|
let out_data: Vec<u8> = Packet::build_packet_from_address(conn_id, remote_addr, data)?.into(); |
||||
|
stream.write_all(&out_data).await?; |
||||
|
|
||||
|
Ok(()) |
||||
|
} |
||||
|
} |
||||
Loading…
Reference in new issue