cliet tun core

This commit is contained in:
2026-03-05 19:55:05 +07:00
parent 9cc9b75123
commit 592005cbac
20 changed files with 525 additions and 2401 deletions
+213
View File
@@ -0,0 +1,213 @@
use smoltcp::{
iface::{SocketHandle, SocketSet},
socket::{AnySocket, icmp, tcp, udp},
wire::IpListenEndpoint,
};
use std::{
collections::HashMap,
time::{Duration, Instant as StdInstant},
};
use tracing::{debug, info};
use crate::tun::engine::START_TIME;
pub struct ConnectionManager {
last_activity: HashMap<SocketHandle, StdInstant>,
proxy_ip: String,
}
impl ConnectionManager {
pub fn new(ip: String) -> Self {
Self {
last_activity: HashMap::new(),
proxy_ip: ip,
}
}
/// Основной метод, который вызывается в цикле Engine
pub fn process_sockets(&mut self, socket_set: &mut SocketSet) {
for (handle, socket) in socket_set.iter_mut() {
// 1. Пытаемся даункастить сокет до TCP
if let Some(tcp) = tcp::Socket::downcast_mut(socket) {
self.handle_tcp(handle, tcp);
continue;
}
// 2. До UDP
if let Some(udp) = udp::Socket::downcast_mut(socket) {
self.handle_udp(handle, udp);
continue;
}
// 3. До ICMP
if let Some(icmp) = icmp::Socket::downcast_mut(socket) {
self.handle_icmp(handle, icmp);
continue;
}
}
}
fn handle_tcp(&mut self, handle: SocketHandle, socket: &mut tcp::Socket) {
// Важно: проверяем, есть ли что принимать
if socket.can_recv() {
let mut data_processed = false;
// Используем recv_slice или проверяем длину в recv
let result = socket.recv(|data| {
if !data.is_empty() {
debug!(handle=%handle, len=data.len(), "TCP: received data");
data_processed = true;
(data.len(), ())
} else {
(0, ())
}
});
if let Err(e) = result {
debug!(handle=%handle, "TCP recv error: {:?}", e);
}
}
// Если сокет закрывается, нужно дать ему закрыться, а не крутить вечно
if !socket.may_recv() && socket.state() == tcp::State::CloseWait {
socket.close();
}
}
fn handle_udp(&mut self, handle: SocketHandle, socket: &mut udp::Socket) {
if socket.can_recv() {
match socket.recv() {
Ok((data, endpoint)) => {
info!(handle=%handle, from=?endpoint, len=data.len(), "UDP: packet received");
// МОК: Обработка UDP датаграммы
}
Err(_) => {}
}
}
}
fn handle_icmp(&mut self, handle: SocketHandle, socket: &mut icmp::Socket) {
if socket.can_recv() {
match socket.recv() {
Ok((data, endpoint)) => {
debug!(handle=%handle, from=?endpoint, "ICMP: packet received");
// МОК: Ответ на пинг или обработка ошибок
}
Err(_) => {}
}
}
}
fn create_tcp_socket<'a>() -> tcp::Socket<'a> {
const BUF_SIZE: usize = 65535;
tcp::Socket::new(
tcp::SocketBuffer::new(vec![0; BUF_SIZE]),
tcp::SocketBuffer::new(vec![0; BUF_SIZE]),
)
}
fn create_udp_socket<'a>() -> udp::Socket<'a> {
const BUF_SIZE: usize = 65535;
udp::Socket::new(
udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY; 16], vec![0; BUF_SIZE]),
udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY; 16], vec![0; BUF_SIZE]),
)
}
fn create_icmp_socket<'a>() -> icmp::Socket<'a> {
let icmp_rx_buffer =
icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; 4], vec![0; 1024]);
let icmp_tx_buffer =
icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; 4], vec![0; 1024]);
icmp::Socket::new(icmp_rx_buffer, icmp_tx_buffer)
}
pub fn refill_sockets(&mut self, socket_set: &mut SocketSet) {
self.prune_sockets(socket_set);
const TARGET_FREE_TCP: usize = 16;
const TARGET_FREE_UDP: usize = 8;
let tcp_listeners = socket_set
.iter()
.filter(|(_, s)| {
if let Some(tcp) = tcp::Socket::downcast(s) {
tcp.state() == tcp::State::Listen
} else {
false
}
})
.count();
let udp_active = socket_set
.iter()
.filter(|(_, s)| udp::Socket::downcast(s).is_some())
.count();
let has_icmp = socket_set
.iter()
.any(|(_, s)| icmp::Socket::downcast(s).is_some());
if tcp_listeners < TARGET_FREE_TCP {
let diff = TARGET_FREE_TCP - tcp_listeners;
debug!("Refilling TCP pool: adding {} listeners", diff);
for _ in 0..diff {
let mut s = Self::create_tcp_socket();
let endpoint = IpListenEndpoint {
addr: None,
port: 8080,
};
s.listen(endpoint).unwrap();
socket_set.add(s);
}
}
if udp_active < TARGET_FREE_UDP {
let diff = TARGET_FREE_UDP - udp_active;
debug!("Refilling UDP pool: adding {} sockets", diff);
for _ in 0..diff {
let s = Self::create_udp_socket();
socket_set.add(s);
}
}
if !has_icmp {
debug!("Adding ICMP socket for echo requests");
let s = Self::create_icmp_socket();
socket_set.add(s);
}
}
fn prune_sockets(&mut self, socket_set: &mut SocketSet) {
let now = StdInstant::now();
let udp_timeout = Duration::from_secs(60); // 1 минута для UDP
let mut to_remove = Vec::new();
for (handle, socket) in socket_set.iter() {
if let Some(tcp) = tcp::Socket::downcast(socket) {
if tcp.state() == tcp::State::Closed {
to_remove.push(handle);
continue;
}
}
if let Some(udp) = udp::Socket::downcast(socket) {
if udp.endpoint().port == 0 {
continue;
}
let last = self.last_activity.get(&handle).unwrap_or(&START_TIME);
if now.duration_since(*last) > udp_timeout {
debug!(handle=%handle, "UDP socket timeout reached");
to_remove.push(handle);
}
}
if let Some(_icmp) = icmp::Socket::downcast(socket) {
continue;
}
}
for handle in to_remove {
socket_set.remove(handle);
self.last_activity.remove(&handle);
}
}
}
@@ -5,8 +5,8 @@ use std::{
mem,
ops::{Deref, DerefMut},
sync::{
atomic::{AtomicBool, Ordering},
Arc, LazyLock, Mutex,
atomic::{AtomicBool, Ordering},
},
};
+144
View File
@@ -0,0 +1,144 @@
use smoltcp::time::Instant;
use smoltcp::{
iface::{Config, Interface, SocketSet},
phy::DeviceCapabilities,
};
use std::{
mem,
sync::{Arc, LazyLock, atomic::AtomicBool},
time::Instant as StdInstant,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
use tokio::time::{Duration, sleep};
use tun::{DeviceReader, DeviceWriter};
use tracing::{debug, warn};
use crate::tun::connection_manager::ConnectionManager;
use crate::tun::device::{TokenBuffer, VirtTunDevice};
use crate::tun::tun::Tun;
pub static START_TIME: LazyLock<StdInstant> = LazyLock::new(StdInstant::now);
pub struct Engine {
interface: Interface,
socket_set: SocketSet<'static>,
manager: ConnectionManager,
device: VirtTunDevice,
bridge_rx: UnboundedReceiver<TokenBuffer>,
bridge_tx: UnboundedSender<TokenBuffer>,
avail: Arc<AtomicBool>,
}
impl Engine {
pub fn new(config: Config, caps: DeviceCapabilities, ip: String) -> Self {
let now = Engine::current_time();
let (mut device, bridge_rx, bridge_tx, avail) = VirtTunDevice::new(caps);
let interface = Interface::new(config, &mut device, now);
let socket_set = SocketSet::new(vec![]);
Self {
interface,
socket_set,
device,
bridge_tx,
bridge_rx,
avail,
manager: ConnectionManager::new(ip),
}
}
pub async fn run(&mut self, tun: Tun) {
//Bridge from tun to stack and stack to tun
let (writer, reader) = tun.split().expect("Failed to split TUN");
// Забираем bridge_rx, так как он нам нужен только в одной задаче
let from_engine = mem::replace(&mut self.bridge_rx, mpsc::unbounded_channel().1);
Self::spawn_tun_to_engine(reader, self.bridge_tx.clone(), self.avail.clone());
Self::spawn_engine_to_tun(writer, from_engine);
loop {
self.manager.refill_sockets(&mut self.socket_set);
self.poll();
self.manager.process_sockets(&mut self.socket_set);
if self.avail.load(std::sync::atomic::Ordering::Acquire) {
tokio::task::yield_now().await;
continue;
}
self.poll_delay().await;
}
}
pub fn poll(&mut self) {
let now = Self::current_time();
// Передаем девайс и сокеты в интерфейс
self.interface
.poll(now, &mut self.device, &mut self.socket_set);
}
pub async fn poll_delay(&mut self) {
let timestamp = Self::current_time();
let delay = self.interface.poll_delay(timestamp, &self.socket_set);
let sleep_duration = match delay {
Some(d) => Duration::from_micros(d.micros()),
None => Duration::from_millis(10),
};
sleep(sleep_duration).await;
}
fn spawn_tun_to_engine(
mut reader: DeviceReader,
to_engine: UnboundedSender<TokenBuffer>,
is_avail: Arc<AtomicBool>,
) {
tokio::spawn(async move {
debug!("TUN-to-Engine bridge task started");
let mut buf = [0u8; 4096];
while let Ok(n) = reader.read(&mut buf).await {
if n == 0 {
break;
}
let mut token = TokenBuffer::with_capacity(n);
token.extend_from_slice(&buf[..n]);
if to_engine.send(token).is_ok() {
is_avail.store(true, std::sync::atomic::Ordering::Release);
} else {
break;
}
}
warn!("TUN-to-Engine bridge task stopped");
});
}
fn spawn_engine_to_tun(
mut writer: DeviceWriter,
mut from_engine: UnboundedReceiver<TokenBuffer>,
) {
tokio::spawn(async move {
debug!("Engine-to-TUN bridge task started");
while let Some(token) = from_engine.recv().await {
if writer.write_all(&token).await.is_err() {
break;
}
}
warn!("Engine-to-TUN bridge task stopped");
});
}
fn current_time() -> Instant {
let duration = StdInstant::now().duration_since(*START_TIME);
Instant::from_micros(duration.as_micros() as i64)
}
pub fn add_address(&mut self, address: smoltcp::wire::IpCidr) {
self.interface.update_ip_addrs(|addrs| {
addrs
.push(address)
.expect("Failed to add IP: address list is full");
});
self.interface
.routes_mut()
.add_default_ipv4_route(smoltcp::wire::Ipv4Address::new(10, 0, 0, 1))
.ok();
}
}
-53
View File
@@ -1,53 +0,0 @@
//! IP packet encapsulation
use std::net::IpAddr;
use smoltcp::wire::{IpProtocol, IpVersion, Ipv4Packet, Ipv6Packet};
#[derive(Debug)]
pub enum IpPacket<T: AsRef<[u8]>> {
Ipv4(Ipv4Packet<T>),
Ipv6(Ipv6Packet<T>),
}
impl<T: AsRef<[u8]> + Copy> IpPacket<T> {
pub fn new_checked(packet: T) -> smoltcp::wire::Result<Option<Self>> {
let buffer = packet.as_ref();
match IpVersion::of_packet(buffer)? {
IpVersion::Ipv4 => Ok(Some(Self::Ipv4(Ipv4Packet::new_checked(packet)?))),
IpVersion::Ipv6 => Ok(Some(Self::Ipv6(Ipv6Packet::new_checked(packet)?))),
}
}
pub fn src_addr(&self) -> IpAddr {
match *self {
Self::Ipv4(ref packet) => IpAddr::from(packet.src_addr()),
Self::Ipv6(ref packet) => IpAddr::from(packet.src_addr()),
}
}
pub fn dst_addr(&self) -> IpAddr {
match *self {
Self::Ipv4(ref packet) => IpAddr::from(packet.dst_addr()),
Self::Ipv6(ref packet) => IpAddr::from(packet.dst_addr()),
}
}
pub fn protocol(&self) -> IpProtocol {
match *self {
Self::Ipv4(ref packet) => packet.next_header(),
Self::Ipv6(ref packet) => packet.next_header(),
}
}
}
impl<'a, T: AsRef<[u8]> + ?Sized> IpPacket<&'a T> {
/// Return a pointer to the payload.
#[inline]
pub fn payload(&self) -> &'a [u8] {
match *self {
IpPacket::Ipv4(ref packet) => packet.payload(),
IpPacket::Ipv6(ref packet) => packet.payload(),
}
}
}
+3 -6
View File
@@ -1,7 +1,4 @@
mod ip_packet;
pub mod tcp;
pub mod connection_manager;
pub mod device;
pub mod engine;
pub mod tun;
pub mod udp;
pub mod virt_device;
pub use tun::Tun;
pub use tun::TunBuilder;
-669
View File
@@ -1,669 +0,0 @@
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
}
+29 -283
View File
@@ -1,299 +1,45 @@
use std::net::SocketAddr;
use std::os::unix::io::RawFd;
use std::{io, mem, net::IpAddr, time::Duration};
use std::io;
use tracing::error;
use tun::{AsyncDevice, Configuration, DeviceReader, DeviceWriter, create_as_async};
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::{
AbstractDevice, AsyncDevice, Configuration as TunConfiguration, Error as TunError, Layer,
create_as_async,
};
use crate::tun::ip_packet::IpPacket;
use crate::tun::virt_device::TokenBuffer;
pub struct TunBuilder {
tun_config: TunConfiguration,
udp_expiry_duration: Option<Duration>,
udp_capacity: Option<usize>,
}
/// TunConfiguration contains a HANDLE, which is a *mut c_void on Windows.
unsafe impl Send for TunBuilder {}
impl TunBuilder {
/// Create a Tun service builder
pub fn new() -> Self {
Self {
tun_config: TunConfiguration::default(),
udp_expiry_duration: None,
udp_capacity: None,
}
}
pub fn address(&mut self, addr: IpNet) {
self.tun_config.address(addr.addr()).netmask(addr.netmask());
}
pub fn destination(&mut self, addr: IpNet) {
self.tun_config.destination(addr.addr());
}
pub fn name(&mut self, name: &str) {
self.tun_config.tun_name(name);
}
#[cfg(unix)]
pub fn file_descriptor(&mut self, fd: RawFd) {
self.tun_config.raw_fd(fd);
}
pub fn udp_expiry_duration(&mut self, udp_expiry_duration: Duration) {
self.udp_expiry_duration = Some(udp_expiry_duration);
}
pub fn udp_capacity(&mut self, udp_capacity: usize) {
self.udp_capacity = Some(udp_capacity);
}
/// Build Tun server
pub async fn build(mut self) -> io::Result<Tun> {
self.tun_config.layer(Layer::L3).up();
let device = match create_as_async(&self.tun_config) {
Ok(d) => d,
Err(TunError::Io(err)) => return Err(err),
Err(err) => return Err(io::Error::other(err)),
};
Ok(Tun { device })
}
}
/// Tun service
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 {
pub async fn run(mut self) -> io::Result<()> {
info!(
"tun device {}",
self.device
.tun_name()
.or_else(|r| Ok::<_, ()>(r.to_string()))
.unwrap(),
);
pub fn build() -> Configuration {
Configuration::default()
}
let address = match self.device.address() {
Ok(a) => a,
Err(err) => {
error!("[TUN] failed to get device address, error: {}", err);
return Err(io::Error::other(err));
}
};
let netmask = match self.device.netmask() {
Ok(n) => n,
Err(err) => {
error!("[TUN] failed to get device netmask, error: {}", err);
return Err(io::Error::other(err));
}
};
let address_net = match IpNet::with_netmask(address, netmask) {
Ok(n) => n,
Err(err) => {
error!(
"[TUN] invalid address {}, netmask {}, error: {}",
address, netmask, err
);
return Err(io::Error::other(err));
}
};
trace!(
"[TUN] tun device network: {} (address: {}, netmask: {})",
address_net, address, netmask
);
let address_broadcast = address_net.broadcast();
let create_packet_buffer = || {
const PACKET_BUFFER_SIZE: usize = 65536;
let mut packet_buffer = TokenBuffer::with_capacity(PACKET_BUFFER_SIZE);
unsafe {
packet_buffer.set_len(PACKET_BUFFER_SIZE);
}
packet_buffer
};
let mut packet_buffer = create_packet_buffer();
loop {
tokio::select! {
n = self.device.read(&mut packet_buffer) => {
let n = n?;
let mut packet_buffer = mem::replace(&mut packet_buffer, create_packet_buffer());
unsafe {
packet_buffer.set_len(n);
}
trace!("[TUN] received IP packet {:?}", ByteStr::new(&packet_buffer));
if let Err(err) = self.handle_tun_frame(&address_broadcast, packet_buffer).await {
error!("[TUN] handle IP frame failed, error: {}", err);
}
}
pub fn new(config: &Configuration) -> io::Result<Self> {
match create_as_async(config) {
Ok(device) => Ok(Self { device }),
Err(e) => {
error!("Failed to create TUN device: {}", e);
Err(io::Error::new(io::ErrorKind::Other, e))
}
}
}
async fn handle_tun_frame(
&mut self,
device_broadcast_addr: &IpAddr,
frame: TokenBuffer,
) -> smoltcp::wire::Result<()> {
let packet = match IpPacket::new_checked(frame.as_ref())? {
Some(packet) => packet,
None => {
warn!("unrecognized IP packet {:?}", ByteStr::new(&frame));
return Ok(());
}
};
pub fn create<F>(f: F) -> io::Result<Self>
where
F: FnOnce(&mut Configuration),
{
let mut config = Configuration::default();
f(&mut config);
Self::new(&config)
}
trace!("[TUN] {:?}", packet);
pub fn from_android_fd(fd: i32) -> io::Result<Self> {
let mut config = Configuration::default();
config.raw_fd(fd); // Передаем дескриптор, который нам дал Android VpnService
config.up(); // Убеждаемся, что он поднят
let src_ip_addr = packet.src_addr();
let dst_ip_addr = packet.dst_addr();
let src_non_unicast = src_ip_addr == *device_broadcast_addr
|| match src_ip_addr {
IpAddr::V4(v4) => v4.is_broadcast() || v4.is_multicast() || v4.is_unspecified(),
IpAddr::V6(v6) => v6.is_multicast() || v6.is_unspecified(),
};
let dst_non_unicast = dst_ip_addr == *device_broadcast_addr
|| match dst_ip_addr {
IpAddr::V4(v4) => v4.is_broadcast() || v4.is_multicast() || v4.is_unspecified(),
IpAddr::V6(v6) => v6.is_multicast() || v6.is_unspecified(),
};
Self::new(&config)
}
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
);
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(())
pub fn split(self) -> io::Result<(DeviceWriter, DeviceReader)> {
let (writer, reader) = self.device.split()?;
Ok((writer, reader))
}
}
-171
View File
@@ -1,171 +0,0 @@
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(())
}
}