copy pasted socks tun

This commit is contained in:
2026-03-05 12:18:07 +07:00
parent 8d3c7875cf
commit 9cc9b75123
15 changed files with 2074 additions and 140 deletions
+5 -1
View File
@@ -1,3 +1,7 @@
mod ip_packet;
pub mod tun_builder;
pub mod tcp;
pub mod tun;
pub mod udp;
pub mod virt_device;
pub use tun::Tun;
pub use tun::TunBuilder;
+669
View File
@@ -0,0 +1,669 @@
use std::{
collections::HashMap,
future::Future,
io, mem,
net::{IpAddr, SocketAddr},
pin::Pin,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
task::{Context, Poll, Waker},
thread::{self, JoinHandle, Thread},
time::Duration,
};
use log::{debug, error, trace};
use smoltcp::{
iface::{Config as InterfaceConfig, Interface, PollResult, SocketHandle, SocketSet},
phy::{Checksum, DeviceCapabilities, Medium},
socket::tcp::{
CongestionControl, Socket as TcpSocket, SocketBuffer as TcpSocketBuffer, State as TcpState,
},
storage::RingBuffer,
time::{Duration as SmolDuration, Instant as SmolInstant},
wire::{HardwareAddress, IpAddress, IpCidr, Ipv4Address, Ipv6Address, TcpPacket},
};
use spin::Mutex as SpinMutex;
use tokio::{
io::{AsyncRead, AsyncWrite, ReadBuf},
sync::{mpsc, oneshot},
};
use super::virt_device::{TokenBuffer, VirtTunDevice};
#[derive(Debug, Clone, Default)]
pub struct TcpSocketOpts {
/// TCP socket's `SO_SNDBUF`
pub send_buffer_size: Option<u32>,
/// TCP socket's `SO_RCVBUF`
pub recv_buffer_size: Option<u32>,
/// `TCP_NODELAY`
pub nodelay: bool,
/// `TCP_FASTOPEN`, enables TFO
pub fastopen: bool,
/// `SO_KEEPALIVE` and sets `TCP_KEEPIDLE`, `TCP_KEEPINTVL` and `TCP_KEEPCNT` respectively,
/// enables keep-alive messages on connection-oriented sockets
pub keepalive: Option<Duration>,
/// Enable Multipath-TCP (mptcp)
/// https://en.wikipedia.org/wiki/Multipath_TCP
///
/// Currently only supported on
/// - macOS (iOS, watchOS, ...) with Client Support only.
/// - Linux (>5.19)
pub mptcp: bool,
}
// NOTE: Default buffer could contain 5 AEAD packets
const DEFAULT_TCP_SEND_BUFFER_SIZE: u32 = (0x3FFFu32 * 5).next_power_of_two();
const DEFAULT_TCP_RECV_BUFFER_SIZE: u32 = (0x3FFFu32 * 5).next_power_of_two();
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
enum TcpSocketState {
Normal,
Close,
Closing,
Closed,
}
struct TcpSocketControl {
send_buffer: RingBuffer<'static, u8>,
send_waker: Option<Waker>,
recv_buffer: RingBuffer<'static, u8>,
recv_waker: Option<Waker>,
recv_state: TcpSocketState,
send_state: TcpSocketState,
}
struct ManagerNotify {
thread: Thread,
}
impl ManagerNotify {
fn new(thread: Thread) -> Self {
Self { thread }
}
fn notify(&self) {
self.thread.unpark();
}
}
struct TcpSocketManager {
device: VirtTunDevice,
iface: Interface,
sockets: HashMap<SocketHandle, SharedTcpConnectionControl>,
socket_creation_rx: mpsc::UnboundedReceiver<TcpSocketCreation>,
}
type SharedTcpConnectionControl = Arc<SpinMutex<TcpSocketControl>>;
struct TcpSocketCreation {
control: SharedTcpConnectionControl,
socket: TcpSocket<'static>,
socket_created_tx: oneshot::Sender<()>,
}
struct TcpConnection {
control: SharedTcpConnectionControl,
manager_notify: Arc<ManagerNotify>,
}
impl Drop for TcpConnection {
fn drop(&mut self) {
let mut control = self.control.lock();
if matches!(control.recv_state, TcpSocketState::Normal) {
control.recv_state = TcpSocketState::Close;
}
if matches!(control.send_state, TcpSocketState::Normal) {
control.send_state = TcpSocketState::Close;
}
self.manager_notify.notify();
}
}
impl TcpConnection {
fn new(
socket: TcpSocket<'static>,
socket_creation_tx: &mpsc::UnboundedSender<TcpSocketCreation>,
manager_notify: Arc<ManagerNotify>,
tcp_opts: &TcpSocketOpts,
) -> impl Future<Output = Self> + use<> {
let send_buffer_size = tcp_opts
.send_buffer_size
.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
let recv_buffer_size = tcp_opts
.recv_buffer_size
.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);
let control = Arc::new(SpinMutex::new(TcpSocketControl {
send_buffer: RingBuffer::new(vec![0u8; send_buffer_size as usize]),
send_waker: None,
recv_buffer: RingBuffer::new(vec![0u8; recv_buffer_size as usize]),
recv_waker: None,
recv_state: TcpSocketState::Normal,
send_state: TcpSocketState::Normal,
}));
let (tx, rx) = oneshot::channel();
let _ = socket_creation_tx.send(TcpSocketCreation {
control: control.clone(),
socket,
socket_created_tx: tx,
});
async move {
// waiting socket add to SocketSet
let _ = rx.await;
Self {
control,
manager_notify,
}
}
}
}
impl AsyncRead for TcpConnection {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let mut control = self.control.lock();
// Read from buffer
if control.recv_buffer.is_empty() {
// If socket is already closed / half closed, just return EOF directly.
if matches!(control.recv_state, TcpSocketState::Closed) {
return Ok(()).into();
}
// Nothing could be read. Wait for notify.
if let Some(old_waker) = control.recv_waker.replace(cx.waker().clone())
&& !old_waker.will_wake(cx.waker())
{
old_waker.wake();
}
return Poll::Pending;
}
let recv_buf =
unsafe { mem::transmute::<&mut [mem::MaybeUninit<u8>], &mut [u8]>(buf.unfilled_mut()) };
let n = control.recv_buffer.dequeue_slice(recv_buf);
buf.advance(n);
if n > 0 {
self.manager_notify.notify();
}
Ok(()).into()
}
}
impl AsyncWrite for TcpConnection {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
let mut control = self.control.lock();
// If state == Close | Closing | Closed, the TCP stream WR half is closed.
if !matches!(control.send_state, TcpSocketState::Normal) {
return Err(io::ErrorKind::BrokenPipe.into()).into();
}
// Write to buffer
if control.send_buffer.is_full() {
if let Some(old_waker) = control.send_waker.replace(cx.waker().clone())
&& !old_waker.will_wake(cx.waker())
{
old_waker.wake();
}
return Poll::Pending;
}
let n = control.send_buffer.enqueue_slice(buf);
if n > 0 {
self.manager_notify.notify();
}
Ok(n).into()
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Ok(()).into()
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
let mut control = self.control.lock();
if matches!(control.send_state, TcpSocketState::Closed) {
return Ok(()).into();
}
// SHUT_WR
if matches!(control.send_state, TcpSocketState::Normal) {
control.send_state = TcpSocketState::Close;
}
if let Some(old_waker) = control.send_waker.replace(cx.waker().clone())
&& !old_waker.will_wake(cx.waker())
{
old_waker.wake();
}
self.manager_notify.notify();
Poll::Pending
}
}
pub struct TcpTun {
context: Arc<ServiceContext>,
manager_handle: Option<JoinHandle<()>>,
manager_notify: Arc<ManagerNotify>,
manager_socket_creation_tx: mpsc::UnboundedSender<TcpSocketCreation>,
manager_running: Arc<AtomicBool>,
balancer: PingBalancer,
iface_rx: mpsc::UnboundedReceiver<TokenBuffer>,
iface_tx: mpsc::UnboundedSender<TokenBuffer>,
iface_tx_avail: Arc<AtomicBool>,
}
impl Drop for TcpTun {
fn drop(&mut self) {
self.manager_running.store(false, Ordering::Relaxed);
self.manager_notify.notify();
let _ = self.manager_handle.take().unwrap().join();
}
}
impl TcpTun {
pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mtu: u32) -> Self {
let mut capabilities = DeviceCapabilities::default();
capabilities.medium = Medium::Ip;
capabilities.max_transmission_unit = mtu as usize;
capabilities.checksum.ipv4 = Checksum::Tx;
capabilities.checksum.tcp = Checksum::Tx;
capabilities.checksum.udp = Checksum::Tx;
capabilities.checksum.icmpv4 = Checksum::Tx;
capabilities.checksum.icmpv6 = Checksum::Tx;
let (mut device, iface_rx, iface_tx, iface_tx_avail) = VirtTunDevice::new(capabilities);
let mut iface_config = InterfaceConfig::new(HardwareAddress::Ip);
iface_config.random_seed = rand::random();
let mut iface = Interface::new(iface_config, &mut device, SmolInstant::now());
iface.update_ip_addrs(|ip_addrs| {
ip_addrs
.push(IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0))
.expect("iface IPv4");
ip_addrs
.push(IpCidr::new(IpAddress::v6(0, 0, 0, 0, 0, 0, 0, 1), 0))
.expect("iface IPv6");
});
iface
.routes_mut()
.add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
.expect("IPv4 default route");
iface
.routes_mut()
.add_default_ipv6_route(Ipv6Address::new(0, 0, 0, 0, 0, 0, 0, 1))
.expect("IPv6 default route");
iface.set_any_ip(true);
let (manager_socket_creation_tx, manager_socket_creation_rx) = mpsc::unbounded_channel();
let mut manager = TcpSocketManager {
device,
iface,
sockets: HashMap::new(),
socket_creation_rx: manager_socket_creation_rx,
};
let manager_running = Arc::new(AtomicBool::new(true));
let manager_handle = {
let manager_running = manager_running.clone();
thread::Builder::new()
.name("smoltcp-poll".to_owned())
.spawn(move || {
let TcpSocketManager {
ref mut device,
ref mut iface,
ref mut sockets,
ref mut socket_creation_rx,
..
} = manager;
let mut socket_set = SocketSet::new(vec![]);
while manager_running.load(Ordering::Relaxed) {
while let Ok(TcpSocketCreation {
control,
socket,
socket_created_tx: socket_create_tx,
}) = socket_creation_rx.try_recv()
{
let handle = socket_set.add(socket);
let _ = socket_create_tx.send(());
sockets.insert(handle, control);
}
let before_poll = SmolInstant::now();
if let PollResult::SocketStateChanged =
iface.poll(before_poll, device, &mut socket_set)
{
trace!(
"VirtDevice::poll costed {}",
SmolInstant::now() - before_poll
);
}
// Check all the sockets' status
let mut sockets_to_remove = Vec::new();
for (socket_handle, control) in sockets.iter() {
let socket_handle = *socket_handle;
let socket = socket_set.get_mut::<TcpSocket>(socket_handle);
let mut control = control.lock();
// Remove the socket only when it is in the closed state.
if socket.state() == TcpState::Closed {
sockets_to_remove.push(socket_handle);
control.send_state = TcpSocketState::Closed;
control.recv_state = TcpSocketState::Closed;
if let Some(waker) = control.send_waker.take() {
waker.wake();
}
if let Some(waker) = control.recv_waker.take() {
waker.wake();
}
trace!("closed TCP connection");
continue;
}
// SHUT_WR
if matches!(control.send_state, TcpSocketState::Close)
&& socket.send_queue() == 0
&& control.send_buffer.is_empty()
{
trace!("closing TCP Write Half, {:?}", socket.state());
// Close the socket. Set to FIN state
socket.close();
control.send_state = TcpSocketState::Closing;
// We can still process the pending buffer.
}
// Check if readable
let mut wake_receiver = false;
while socket.can_recv() && !control.recv_buffer.is_full() {
let result = socket.recv(|buffer| {
let n = control.recv_buffer.enqueue_slice(buffer);
(n, ())
});
match result {
Ok(..) => {
wake_receiver = true;
}
Err(err) => {
error!(
"socket recv error: {:?}, {:?}",
err,
socket.state()
);
// Don't know why. Abort the connection.
socket.abort();
if matches!(control.recv_state, TcpSocketState::Normal) {
control.recv_state = TcpSocketState::Closed;
}
wake_receiver = true;
// The socket will be recycled in the next poll.
break;
}
}
}
// If socket is not in ESTABLISH, FIN-WAIT-1, FIN-WAIT-2,
// the local client have closed our receiver.
if matches!(control.recv_state, TcpSocketState::Normal)
&& !socket.may_recv()
&& !matches!(
socket.state(),
TcpState::Listen
| TcpState::SynReceived
| TcpState::Established
| TcpState::FinWait1
| TcpState::FinWait2
)
{
trace!("closed TCP Read Half, {:?}", socket.state());
// Let TcpConnection::poll_read returns EOF.
control.recv_state = TcpSocketState::Closed;
wake_receiver = true;
}
if wake_receiver
&& control.recv_waker.is_some()
&& let Some(waker) = control.recv_waker.take()
{
waker.wake();
}
// Check if writable
let mut wake_sender = false;
while socket.can_send() && !control.send_buffer.is_empty() {
let result = socket.send(|buffer| {
let n = control.send_buffer.dequeue_slice(buffer);
(n, ())
});
match result {
Ok(..) => {
wake_sender = true;
}
Err(err) => {
error!(
"socket send error: {:?}, {:?}",
err,
socket.state()
);
// Don't know why. Abort the connection.
socket.abort();
if matches!(control.send_state, TcpSocketState::Normal) {
control.send_state = TcpSocketState::Closed;
}
wake_sender = true;
// The socket will be recycled in the next poll.
break;
}
}
}
if wake_sender
&& control.send_waker.is_some()
&& let Some(waker) = control.send_waker.take()
{
waker.wake();
}
}
for socket_handle in sockets_to_remove {
sockets.remove(&socket_handle);
socket_set.remove(socket_handle);
}
if !device.recv_available() {
let next_duration = iface
.poll_delay(before_poll, &socket_set)
.unwrap_or(SmolDuration::from_millis(5));
if next_duration != SmolDuration::ZERO {
thread::park_timeout(Duration::from(next_duration));
}
}
}
trace!("VirtDevice::poll thread exited");
})
.unwrap()
};
let manager_notify = Arc::new(ManagerNotify::new(manager_handle.thread().clone()));
Self {
context,
manager_handle: Some(manager_handle),
manager_notify,
manager_socket_creation_tx,
manager_running,
balancer,
iface_rx,
iface_tx,
iface_tx_avail,
}
}
pub async fn handle_packet(
&mut self,
src_addr: SocketAddr,
dst_addr: SocketAddr,
tcp_packet: &TcpPacket<&[u8]>,
) -> io::Result<()> {
// TCP first handshake packet, create a new Connection
if tcp_packet.syn() && !tcp_packet.ack() {
let accept_opts = self.context.accept_opts();
let send_buffer_size = accept_opts
.tcp
.send_buffer_size
.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
let recv_buffer_size = accept_opts
.tcp
.recv_buffer_size
.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);
let mut socket = TcpSocket::new(
TcpSocketBuffer::new(vec![0u8; recv_buffer_size as usize]),
TcpSocketBuffer::new(vec![0u8; send_buffer_size as usize]),
);
socket.set_keep_alive(accept_opts.tcp.keepalive.map(From::from));
// FIXME: It should follow system's setting. 7200 is Linux's default.
socket.set_timeout(Some(SmolDuration::from_secs(7200)));
// NO ACK delay
// socket.set_ack_delay(None);
// Enable Cubic congestion control
socket.set_congestion_control(CongestionControl::Cubic);
if let Err(err) = socket.listen(dst_addr) {
return Err(io::Error::other(format!("listen error: {:?}", err)));
}
debug!("created TCP connection for {} <-> {}", src_addr, dst_addr);
let connection = TcpConnection::new(
socket,
&self.manager_socket_creation_tx,
self.manager_notify.clone(),
&accept_opts.tcp,
);
// establish a tunnel
let context = self.context.clone();
let balancer = self.balancer.clone();
tokio::spawn(async move {
let connection = connection.await;
if let Err(err) =
handle_redir_client(context, balancer, connection, src_addr, dst_addr).await
{
error!(
"TCP tunnel failure, {} <-> {}, error: {}",
src_addr, dst_addr, err
);
}
});
}
Ok(())
}
pub async fn drive_interface_state(&mut self, frame: TokenBuffer) {
if self.iface_tx.send(frame).is_err() {
panic!("interface send channel closed unexpectedly");
}
// Wake up and poll the interface.
self.iface_tx_avail.store(true, Ordering::Release);
self.manager_notify.notify();
}
pub async fn recv_packet(&mut self) -> TokenBuffer {
match self.iface_rx.recv().await {
Some(v) => v,
None => unreachable!("channel closed unexpectedly"),
}
}
}
/// Established Client Transparent Proxy
///
/// This method must be called after handshaking with client (for example, socks5 handshaking)
async fn establish_client_tcp_redir(
context: Arc<ServiceContext>,
balancer: PingBalancer,
mut stream: TcpConnection,
peer_addr: SocketAddr,
addr: &Address,
) -> io::Result<()> {
if balancer.is_empty() {
let mut remote = AutoProxyClientStream::connect_bypassed(context, addr).await?;
return establish_tcp_tunnel_bypassed(&mut stream, &mut remote, peer_addr, addr).await;
}
let server = balancer.best_tcp_server();
let svr_cfg = server.server_config();
let mut remote =
AutoProxyClientStream::connect_with_opts(context, &server, addr, server.connect_opts_ref())
.await?;
establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, addr).await
}
async fn handle_redir_client(
context: Arc<ServiceContext>,
balancer: PingBalancer,
s: TcpConnection,
peer_addr: SocketAddr,
mut daddr: SocketAddr,
) -> io::Result<()> {
// Get forward address from socket
//
// Try to convert IPv4 mapped IPv6 address for dual-stack mode.
if let SocketAddr::V6(ref a) = daddr
&& let Some(v4) = to_ipv4_mapped(a.ip())
{
daddr = SocketAddr::new(IpAddr::from(v4), a.port());
}
let target_addr = Address::from(daddr);
establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr).await
}
@@ -1,14 +1,17 @@
use std::net::SocketAddr;
use std::os::unix::io::RawFd;
use std::{io, mem, net::IpAddr, time::Duration};
use byte_string::ByteStr;
use ipnet::IpNet;
use log::{error, info, trace, warn};
use smoltcp::wire::{IpProtocol, TcpPacket, UdpPacket};
use tokio::io::AsyncReadExt;
use tracing::debug;
use tun::{
create_as_async, AbstractDevice, AsyncDevice, Configuration as TunConfiguration,
Error as TunError, Layer,
AbstractDevice, AsyncDevice, Configuration as TunConfiguration, Error as TunError, Layer,
create_as_async,
};
use crate::tun::ip_packet::IpPacket;
@@ -76,8 +79,9 @@ pub struct Tun {
device: AsyncDevice,
}
//Ther is a real tun device that creates by Os, it should transfer data to virtual device by itself or bridge
//Maybe bridge should be in interface
impl Tun {
/// Start serving
pub async fn run(mut self) -> io::Result<()> {
info!(
"tun device {}",
@@ -116,9 +120,7 @@ impl Tun {
trace!(
"[TUN] tun device network: {} (address: {}, netmask: {})",
address_net,
address,
netmask
address_net, address, netmask
);
let address_broadcast = address_net.broadcast();
@@ -185,13 +187,113 @@ impl Tun {
if src_non_unicast || dst_non_unicast {
trace!(
"[TUN] IP packet {} (unicast? {}) -> {} (unicast? {}) throwing away",
src_ip_addr,
!src_non_unicast,
dst_ip_addr,
!dst_non_unicast
src_ip_addr, !src_non_unicast, dst_ip_addr, !dst_non_unicast
);
return Ok(());
}
match packet.protocol() {
IpProtocol::Tcp => {
if !self.mode.enable_tcp() {
trace!(
"received TCP packet but mode is {}, throwing away",
self.mode
);
return Ok(());
}
let tcp_packet = match TcpPacket::new_checked(packet.payload()) {
Ok(p) => p,
Err(err) => {
error!(
"invalid TCP packet err: {}, src_ip: {}, dst_ip: {}, payload: {:?}",
err,
packet.src_addr(),
packet.dst_addr(),
ByteStr::new(packet.payload())
);
return Ok(());
}
};
let src_port = tcp_packet.src_port();
let dst_port = tcp_packet.dst_port();
let src_addr = SocketAddr::new(packet.src_addr(), src_port);
let dst_addr = SocketAddr::new(packet.dst_addr(), dst_port);
trace!(
"[TUN] TCP packet {} (unicast? {}) -> {} (unicast? {}) {}",
src_addr, !src_non_unicast, dst_addr, !dst_non_unicast, tcp_packet
);
// TCP first handshake packet.
if let Err(err) = self
.tcp
.handle_packet(src_addr, dst_addr, &tcp_packet)
.await
{
error!(
"handle TCP packet failed, error: {}, {} <-> {}, packet: {:?}",
err, src_addr, dst_addr, tcp_packet
);
}
self.tcp.drive_interface_state(frame).await;
}
IpProtocol::Udp => {
if !self.mode.enable_udp() {
trace!(
"received UDP packet but mode is {}, throwing away",
self.mode
);
return Ok(());
}
let udp_packet = match UdpPacket::new_checked(packet.payload()) {
Ok(p) => p,
Err(err) => {
error!(
"invalid UDP packet err: {}, src_ip: {}, dst_ip: {}, payload: {:?}",
err,
packet.src_addr(),
packet.dst_addr(),
ByteStr::new(packet.payload())
);
return Ok(());
}
};
let src_port = udp_packet.src_port();
let dst_port = udp_packet.dst_port();
let src_addr = SocketAddr::new(src_ip_addr, src_port);
let dst_addr = SocketAddr::new(packet.dst_addr(), dst_port);
let payload = udp_packet.payload();
trace!(
"[TUN] UDP packet {} (unicast? {}) -> {} (unicast? {}) {}",
src_addr, !src_non_unicast, dst_addr, !dst_non_unicast, udp_packet
);
if let Err(err) = self.udp.handle_packet(src_addr, dst_addr, payload).await {
error!(
"handle UDP packet failed, err: {}, packet: {:?}",
err, udp_packet
);
}
}
IpProtocol::Icmp | IpProtocol::Icmpv6 => {
// ICMP is handled by TCP's Interface.
// smoltcp's interface will always send replies to EchoRequest
self.tcp.drive_interface_state(frame).await;
}
_ => {
debug!("IP packet ignored (protocol: {:?})", packet.protocol());
return Ok(());
}
}
Ok(())
}
}
+171
View File
@@ -0,0 +1,171 @@
use std::{
io::{self, ErrorKind},
net::{IpAddr, SocketAddr},
sync::Arc,
time::Duration,
};
use bytes::{BufMut, BytesMut};
use etherparse::PacketBuilder;
use log::debug;
use tokio::sync::mpsc;
use crate::socks::socks5::Address;
pub struct UdpTun {
tun_rx: mpsc::Receiver<BytesMut>,
manager: UdpAssociationManager<UdpTunInboundWriter>,
}
impl UdpTun {
pub fn new(
context: Arc<ServiceContext>,
balancer: PingBalancer,
time_to_live: Option<Duration>,
capacity: Option<usize>,
) -> (Self, Duration, mpsc::Receiver<SocketAddr>) {
let (tun_tx, tun_rx) = mpsc::channel(64);
let (manager, cleanup_interval, keepalive_rx) = UdpAssociationManager::new(
context,
UdpTunInboundWriter::new(tun_tx),
time_to_live,
capacity,
balancer,
);
(Self { tun_rx, manager }, cleanup_interval, keepalive_rx)
}
pub async fn handle_packet(
&mut self,
src_addr: SocketAddr,
dst_addr: SocketAddr,
payload: &[u8],
) -> io::Result<()> {
debug!(
"UDP {} -> {} payload.size: {} bytes",
src_addr,
dst_addr,
payload.len()
);
if let Err(err) = self
.manager
.send_to(src_addr, dst_addr.into(), payload)
.await
{
debug!(
"UDP {} -> {} payload.size: {} bytes failed, error: {}",
src_addr,
dst_addr,
payload.len(),
err,
);
}
Ok(())
}
pub async fn recv_packet(&mut self) -> BytesMut {
match self.tun_rx.recv().await {
Some(b) => b,
None => unreachable!("channel closed unexpectedly"),
}
}
#[inline(always)]
pub async fn cleanup_expired(&mut self) {
self.manager.cleanup_expired().await;
}
#[inline(always)]
pub async fn keep_alive(&mut self, peer_addr: &SocketAddr) {
self.manager.keep_alive(peer_addr).await;
}
}
#[derive(Clone)]
struct UdpTunInboundWriter {
tun_tx: mpsc::Sender<BytesMut>,
}
impl UdpTunInboundWriter {
fn new(tun_tx: mpsc::Sender<BytesMut>) -> Self {
Self { tun_tx }
}
}
impl UdpInboundWrite for UdpTunInboundWriter {
async fn send_to(
&self,
peer_addr: SocketAddr,
remote_addr: &Address,
data: &[u8],
) -> io::Result<()> {
let addr = match *remote_addr {
Address::SocketAddress(sa) => {
// Try to convert IPv4 mapped IPv6 address if server is running on dual-stack mode
match (peer_addr, sa) {
(SocketAddr::V4(..), SocketAddr::V4(..))
| (SocketAddr::V6(..), SocketAddr::V6(..)) => sa,
(SocketAddr::V4(..), SocketAddr::V6(v6)) => {
// If peer is IPv4, then remote_addr can only be IPv4-mapped-IPv6
match to_ipv4_mapped(v6.ip()) {
Some(v4) => SocketAddr::new(IpAddr::from(v4), v6.port()),
None => {
return Err(io::Error::new(
ErrorKind::InvalidData,
"source and destination type unmatch",
));
}
}
}
(SocketAddr::V6(..), SocketAddr::V4(v4)) => {
// Convert remote_addr to IPv4-mapped-IPv6
SocketAddr::new(IpAddr::from(v4.ip().to_ipv6_mapped()), v4.port())
}
}
}
Address::DomainNameAddress(..) => {
let err = io::Error::new(
ErrorKind::InvalidInput,
"tun destination must not be an domain name address",
);
return Err(err);
}
};
let packet = match (peer_addr, addr) {
(SocketAddr::V4(peer), SocketAddr::V4(remote)) => {
let builder = PacketBuilder::ipv4(remote.ip().octets(), peer.ip().octets(), 20)
.udp(remote.port(), peer.port());
let packet = BytesMut::with_capacity(builder.size(data.len()));
let mut packet_writer = packet.writer();
builder
.write(&mut packet_writer, data)
.expect("PacketBuilder::write");
packet_writer.into_inner()
}
(SocketAddr::V6(peer), SocketAddr::V6(remote)) => {
let builder = PacketBuilder::ipv6(remote.ip().octets(), peer.ip().octets(), 20)
.udp(remote.port(), peer.port());
let packet = BytesMut::with_capacity(builder.size(data.len()));
let mut packet_writer = packet.writer();
builder
.write(&mut packet_writer, data)
.expect("PacketBuilder::write");
packet_writer.into_inner()
}
_ => {
return Err(io::Error::new(
ErrorKind::InvalidData,
"source and destination type unmatch",
));
}
};
self.tun_tx.send(packet).await.expect("tun_tx::send");
Ok(())
}
}