nineap/netrunner-proxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

device to smoltcp [AI]

Browse Source
This commit is contained in:
nineAp
2026-06-25 14:00:00 +07:00
parent 30a843d8a6
commit b6056b2a66
7 changed files with 333 additions and 470 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/src/net/connection.rs
+103 -135
View File
@@ -1,4 +1,4 @@
use bytes::{Buf, Bytes};
use bytes::Bytes;
use netrunner_core::{
net::{GLOBAL_MIN_RTT, NetworkConfig, UDP_IDLE_TIMEOUT},
rawcast::{LocalProtocol, RawCastFrame},
@@ -11,10 +11,11 @@ use smoltcp::{
Ipv6Address,
},
};
use std::{collections::VecDeque, sync::{Arc, atomic::{AtomicBool, Ordering}}};
use std::{sync::{Arc, atomic::{AtomicBool, Ordering}}};
use tokio::sync::{OwnedSemaphorePermit, mpsc, oneshot};
use netrunner_logger::{debug, info, instrument};
pub struct ConnectionCore<T> {
pub handle: SocketHandle,
pub tx: mpsc::Sender<T>,
@@ -23,7 +24,9 @@ pub struct ConnectionCore<T> {
}
impl<T> ConnectionCore<T> {
pub fn new(handle: SocketHandle) -> (Self, mpsc::Receiver<T>, mpsc::Sender<Bytes>, Arc<AtomicBool>) {
pub fn new(
handle: SocketHandle,
) -> (Self, mpsc::Receiver<T>, mpsc::Sender<Bytes>, Arc<AtomicBool>) {
let cap = NetworkConfig::global().channel_capacity;
let (tx_to_net, rx_from_smol) = mpsc::channel::<T>(cap);
let (tx_to_smol, rx_from_net) = mpsc::channel::<Bytes>(cap);
@@ -48,13 +51,12 @@ pub enum ConnectionState {
Closed,
}
pub struct TcpConnection {
core: ConnectionCore<Bytes>,
state: ConnectionState,
pending_data: VecDeque<Bytes>,
pending_bytes: usize,
/// A single in-flight chunk that partially fit into smoltcp's tx_buf.
/// When Some, we retry flushing it before reading the next chunk from core.rx.
pending_chunk: Option<Bytes>,
handshake_rx: Option<oneshot::Receiver<()>>,
chunk_buf: Vec<u8>,
server_eof: bool,
@@ -62,8 +64,7 @@ pub struct TcpConnection {
total_up_bytes: u64,
total_down_bytes: u64,
rx_congested: bool,
tx_congested: bool,
tx_congested: bool,
last_rtt_push_ms: i64,
last_pushed_rtt_ms: u32,
}
@@ -72,21 +73,15 @@ impl TcpConnection {
const RTT_PUSH_INTERVAL_MS: i64 = 50;
const RTT_CHANGE_RATIO: f64 = 0.10;
const BUF_FLOOR: usize = 64 * 1024;
/// Верхняя граница (2 МБ) — тоже как в smoltcp.
const BUF_CEIL: usize = 2 * 1024 * 1024;
/// Fallback, когда BBR ещё не дал оценку BDP.
const BUF_FALLBACK: usize = 256 * 1024;
pub fn new(
handle: SocketHandle,
permit: tokio::sync::OwnedSemaphorePermit,
permit: OwnedSemaphorePermit,
) -> (
Self,
mpsc::Receiver<Bytes>,
mpsc::Sender<Bytes>,
oneshot::Sender<()>,
Arc<AtomicBool>
Arc<AtomicBool>,
) {
let (core, rx_from_smol, tx_to_smol, is_saturated) = ConnectionCore::new(handle);
let (handshake_tx, handshake_rx) = oneshot::channel();
@@ -95,14 +90,12 @@ impl TcpConnection {
core,
state: ConnectionState::Handshaking,
permit: Some(permit),
pending_data: VecDeque::new(),
pending_bytes: 0,
pending_chunk: None,
handshake_rx: Some(handshake_rx),
chunk_buf: vec![0u8; NetworkConfig::global().tcp_chunk_size],
server_eof: false,
total_up_bytes: 0,
total_down_bytes: 0,
rx_congested: false,
tx_congested: false,
last_rtt_push_ms: i64::MIN,
last_pushed_rtt_ms: 0,
@@ -139,9 +132,9 @@ impl TcpConnection {
if matches!(socket.state(), tcp::State::Closed | tcp::State::TimeWait) {
info!(
%self.core.handle,
UP = %self.total_up_bytes,
DOWN = %self.total_down_bytes,
%self.core.handle,
UP = %self.total_up_bytes,
DOWN = %self.total_down_bytes,
"🏁 TCP Socket finished and closed"
);
self.state = ConnectionState::Closed;
@@ -149,9 +142,7 @@ impl TcpConnection {
}
}
ConnectionState::Closed => {
return false;
}
ConnectionState::Closed => return false,
ConnectionState::Established => {
info!(
@@ -167,13 +158,12 @@ impl TcpConnection {
}
fn maybe_update_tunnel_rtt(
&mut self,
socket: &mut tcp::Socket,
timestamp: smoltcp::time::Instant,
&mut self,
socket: &mut tcp::Socket,
timestamp: smoltcp::time::Instant,
) {
let now_ms = timestamp.total_millis();
// 1. Троттлинг по времени (кроме самого первого вызова).
if self.last_rtt_push_ms != i64::MIN
&& now_ms - self.last_rtt_push_ms < Self::RTT_PUSH_INTERVAL_MS
{
@@ -181,13 +171,10 @@ impl TcpConnection {
}
let current_rtt = GLOBAL_MIN_RTT.load(Ordering::Relaxed);
// 2. Игнорируем заведомо некорректное нулевое значение.
if current_rtt == 0 {
return;
}
// 3. Обновляем только при значимом изменении (или при первом запуске).
let first_push = self.last_rtt_push_ms == i64::MIN;
let changed_enough = if self.last_pushed_rtt_ms == 0 {
true
@@ -209,133 +196,122 @@ impl TcpConnection {
}
}
fn optimal_buffer_size(&self, socket: &tcp::Socket) -> usize {
let bdp = socket.estimated_bdp();
if bdp == 0 {
// BBR ещё не оценил полосу (старт соединения) — безопасный дефолт.
Self::BUF_FALLBACK
} else {
(bdp.saturating_mul(2)).clamp(Self::BUF_FLOOR, Self::BUF_CEIL)
}
}
fn poll_and_process(&mut self, socket: &mut tcp::Socket, timestamp: smoltcp::time::Instant) {
self.maybe_update_tunnel_rtt(socket, timestamp);
// 1. Читаем из браузера -> в Туннель (Upload)
// ── Upload: browser → smoltcp rx_buf → channel → muxer ──────────
while socket.can_recv() && self.core.tx.capacity() > 0 {
if let Ok(n) = socket.peek_slice(&mut self.chunk_buf, timestamp) {
if n == 0 { break; }
if n == 0 {
break;
}
let chunk = Bytes::copy_from_slice(&self.chunk_buf[..n]);
match self.core.tx.try_send(chunk) {
Ok(_) => {
socket.recv_slice(&mut self.chunk_buf[..n]).unwrap();
self.total_up_bytes += n as u64; // Учет Upload
// Если была перегрузка, а теперь прошло — пишем радостный лог
self.total_up_bytes += n as u64;
if self.tx_congested {
netrunner_logger::debug!(%self.core.handle, "🟢 Upload channel cleared. Resuming read from browser.");
netrunner_logger::debug!(
%self.core.handle,
"🟢 Upload channel cleared. Resuming read from browser."
);
self.tx_congested = false;
}
}
Err(mpsc::error::TrySendError::Full(_)) => {
if !self.tx_congested {
netrunner_logger::debug!(%self.core.handle, "🟡 Upload Congestion: Channel to Muxer is full. Pausing read from browser.");
netrunner_logger::debug!(
%self.core.handle,
"🟡 Upload Congestion: Channel to Muxer is full. Pausing."
);
self.tx_congested = true;
}
continue;
break;
}
Err(mpsc::error::TrySendError::Closed(_)) => {
self.server_eof = true;
break;
}
}
} else { break; }
} else {
break;
}
}
// 2. Читаем из Туннеля -> в Браузер (Download)
// ── Download: muxer → channel → smoltcp tx_buf → browser ─────────
//
// We keep at most ONE in-flight chunk (pending_chunk). When smoltcp's
// tx_buf is full we stop reading from core.rx, which naturally
// backpressures the bounded channel and eventually the muxer.
if !self.server_eof {
// Считаем пороги ОДИН раз до цикла: BDP за время цикла не изменится,
// а дёргать BBR на каждой итерации незачем.
let optimal = self.optimal_buffer_size(socket);
let congestion_threshold = optimal + optimal / 2; // optimal * 1.5
let relief_threshold = optimal / 2; // optimal * 0.5 (гистерезис 3:1)
loop {
if self.pending_bytes > congestion_threshold {
if !self.rx_congested {
netrunner_logger::warn!(
%self.core.handle,
"🟡 Download Congestion: pending={} > {} (optimal={}). Pausing tunnel read.",
self.pending_bytes, congestion_threshold, optimal
);
self.rx_congested = true;
self.core.is_saturated.store(true, Ordering::Release);
}
// First flush any partially-sent chunk from a previous tick.
let chunk = match self.pending_chunk.take() {
Some(c) => c,
None => match self.core.rx.try_recv() {
Ok(data) => {
self.total_down_bytes += data.len() as u64;
data
}
Err(mpsc::error::TryRecvError::Empty) => break,
Err(mpsc::error::TryRecvError::Disconnected) => {
self.server_eof = true;
break;
}
},
};
if !socket.can_send() {
// tx_buf is full → hold the chunk, signal saturation.
self.pending_chunk = Some(chunk);
self.core.is_saturated.store(true, Ordering::Release);
break;
}
else if self.rx_congested && self.pending_bytes < relief_threshold {
netrunner_logger::debug!(
%self.core.handle,
"🟢 Download buffer relieved: pending={} < {}. Resuming tunnel read.",
self.pending_bytes, relief_threshold
);
self.rx_congested = false;
self.core.is_saturated.store(false, Ordering::Release);
}
match self.core.rx.try_recv() {
Ok(data) => {
self.pending_bytes += data.len();
self.total_down_bytes += data.len() as u64;
self.pending_data.push_back(data);
match socket.send_slice(&chunk) {
Ok(n) if n == chunk.len() => {
// Entire chunk accepted.
}
Err(mpsc::error::TryRecvError::Empty) => break,
Err(mpsc::error::TryRecvError::Disconnected) => {
self.server_eof = true;
Ok(n) => {
// Partial write — keep the remainder for the next tick.
self.pending_chunk = Some(chunk.slice(n..));
break;
}
Err(e) => {
netrunner_logger::debug!(
%self.core.handle,
"smoltcp send error: {:?}", e
);
self.pending_chunk = Some(chunk);
break;
}
}
}
// Clear saturation when tx_buf has room and no chunk is pending.
if self.pending_chunk.is_none() && socket.can_send() {
self.core.is_saturated.store(false, Ordering::Release);
}
}
// 3. Отправляем в smoltcp то, что накопилось (оставляем как было)
while socket.can_send() {
if let Some(mut chunk) = self.pending_data.pop_front() {
match socket.send_slice(&chunk) {
Ok(n) => {
self.pending_bytes -= n;
if n < chunk.len() {
chunk.advance(n);
self.pending_data.push_front(chunk);
break;
}
}
Err(e) => {
netrunner_logger::debug!(%self.core.handle, "Smoltcp send error: {:?}", e);
self.pending_data.push_front(chunk);
break;
}
}
} else { break; }
}
if self.server_eof && self.pending_data.is_empty() {
if self.server_eof && self.pending_chunk.is_none() {
let state = socket.state();
if state == tcp::State::Established || state == tcp::State::CloseWait {
netrunner_logger::debug!(%self.core.handle, "All data flushed, sending FIN to browser");
netrunner_logger::debug!(
%self.core.handle,
"All data flushed, sending FIN to browser"
);
socket.close();
}
}
}
pub fn app_pending_out_size(&self) -> usize {
self.pending_bytes
self.pending_chunk.as_ref().map(|c| c.len()).unwrap_or(0)
}
#[instrument(skip(rx_smol, handshake_tx, tx_tunnel), fields(
socket_id = socket_id,
socket_id = socket_id,
dst = %target
))]
pub fn spawn(
@@ -359,27 +335,23 @@ impl TcpConnection {
let _ = handshake_tx.send(());
while let Some(data) = rx_smol.recv().await {
let data_frame = RawCastFrame::data(
LocalProtocol::Tcp,
socket_id,
dst_ip,
dst_port,
data
);
let data_frame =
RawCastFrame::data(LocalProtocol::Tcp, socket_id, dst_ip, dst_port, data);
if tx_tunnel.send(data_frame).await.is_err() {
break;
}
}
let close_frame = RawCastFrame::close(LocalProtocol::Tcp, socket_id, dst_ip, dst_port);
let close_frame =
RawCastFrame::close(LocalProtocol::Tcp, socket_id, dst_ip, dst_port);
let _ = tx_tunnel.send(close_frame).await;
debug!("🏁 [TCP {}] Spawned task finished", socket_id);
});
}
}
// ─── UDP ────────────────────────────────────────────────────────────────────
pub type UdpPacketTarget = (Bytes, std::net::Ipv4Addr, u16);
pub struct UdpConnection {
@@ -391,7 +363,7 @@ pub struct UdpConnection {
impl UdpConnection {
pub fn new(
handle: SocketHandle,
client_addr: smoltcp::wire::IpAddress,
client_addr: IpAddress,
client_port: u16,
) -> (Self, mpsc::Receiver<UdpPacketTarget>, mpsc::Sender<Bytes>, Arc<AtomicBool>) {
let (core, rx_from_smol, tx_to_smol, is_saturated) = ConnectionCore::new(handle);
@@ -416,7 +388,7 @@ impl UdpConnection {
if socket.can_recv() {
while let Ok((data, metadata)) = socket.recv(timestamp) {
if let smoltcp::wire::IpAddress::Ipv4(ip) = metadata.endpoint.addr {
if let IpAddress::Ipv4(ip) = metadata.endpoint.addr {
self.last_client_endpoint = Some(metadata.endpoint);
let target_ip = std::net::Ipv4Addr::from(ip);
let target_port = metadata.endpoint.port;
@@ -447,7 +419,7 @@ impl UdpConnection {
}
#[instrument(skip(rx_smol, tx_tunnel), fields(
socket_id = socket_id,
socket_id = socket_id,
dst = %target
))]
pub fn spawn(
@@ -469,20 +441,22 @@ impl UdpConnection {
}
while let Some((data, ip, port)) = rx_smol.recv().await {
let data_frame =
RawCastFrame::data(LocalProtocol::Udp, socket_id, ip, port, data);
let data_frame = RawCastFrame::data(LocalProtocol::Udp, socket_id, ip, port, data);
if tx_tunnel.send(data_frame).await.is_err() {
break;
}
}
let close_frame = RawCastFrame::close(LocalProtocol::Udp, socket_id, dst_ip, dst_port);
let close_frame =
RawCastFrame::close(LocalProtocol::Udp, socket_id, dst_ip, dst_port);
let _ = tx_tunnel.send(close_frame).await;
info!("🛑 [UDP {}] Task stopped", socket_id);
});
}
}
// ─── ICMP ───────────────────────────────────────────────────────────────────
use smoltcp::socket::icmp;
pub struct IcmpResponder;
@@ -493,11 +467,8 @@ impl IcmpResponder {
return;
}
let result = socket.recv(timestamp);
if let Ok((data, src_addr)) = result {
if let Ok((data, src_addr)) = socket.recv(timestamp) {
let payload = data.to_vec();
match src_addr {
IpAddress::Ipv4(_) => Self::reply_v4(socket, payload, src_addr),
IpAddress::Ipv6(v6) => Self::reply_v6(socket, payload, v6),
@@ -511,7 +482,6 @@ impl IcmpResponder {
let mut reply_pkt = Icmpv4Packet::new_unchecked(&mut payload);
reply_pkt.set_msg_type(Icmpv4Message::EchoReply);
reply_pkt.fill_checksum();
let _ = socket.send_slice(&payload, src);
info!("🏓 [ICMPv4] Echo Reply -> {}", src);
}
@@ -523,10 +493,8 @@ impl IcmpResponder {
if pkt.msg_type() == Icmpv6Message::EchoRequest {
let mut reply_pkt = Icmpv6Packet::new_unchecked(&mut payload);
reply_pkt.set_msg_type(Icmpv6Message::EchoReply);
let gateway = Ipv6Address::new(0xfe80, 0, 0, 0, 0, 0, 0, 1);
reply_pkt.fill_checksum(&gateway, &src);
let _ = socket.send_slice(&payload, src.into());
info!("🏓 [ICMPv6] Echo Reply -> {}", src);
}
client/src/net/engine.rs
+169 -134
View File
@@ -3,6 +3,7 @@ use netrunner_core::net::ClientHandler;
use netrunner_core::net::NetworkConfig;
use netrunner_core::rawcast::{RawCastEvent, RawCastFrame};
use smoltcp::iface::PollResult;
use smoltcp::phy::ChannelDevice;
use smoltcp::time::Instant;
use smoltcp::wire::{IpAddress, IpCidr};
use smoltcp::{
@@ -10,13 +11,11 @@ use smoltcp::{
phy::DeviceCapabilities,
};
use std::net::Ipv4Addr;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::{
sync::{Arc, LazyLock, atomic::AtomicBool},
time::Instant as StdInstant,
};
use std::{sync::LazyLock, time::Instant as StdInstant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::mpsc::{self, Receiver, Sender, UnboundedReceiver, UnboundedSender};
use tokio::sync::mpsc;
use tokio::time::{Duration, sleep};
use tun::{DeviceReader, DeviceWriter};
@@ -25,22 +24,32 @@ use netrunner_logger::{debug, error, info, warn};
use crate::net::connection_manager::ConnectionManager;
use crate::net::dns::DnsHandler;
use crate::net::socket_factory::{SmolSocketFactory, SocketProvider};
use crate::tun::device::{TokenBuffer, VirtTunDevice};
use crate::tun::device::TrafficCounter;
use crate::tun::routing::setup_platform_routing;
use crate::tun::tun::Tun;
pub static START_TIME: LazyLock<StdInstant> = LazyLock::new(StdInstant::now);
/// How many inbound packets the device can hold before backpressure kicks in.
/// Each packet is ≤ MTU bytes, so at 1500 B × 64 = 96 KB max queue.
const DEVICE_RX_CAP: usize = 64;
/// How many outbound packets smoltcp can stage before we drain them.
const DEVICE_TX_CAP: usize = 64;
/// Bounded capacity for the TUN-reader → engine channel (packets).
const TUN_CHAN_CAP: usize = 128;
pub struct Engine {
interface: Interface,
socket_set: SocketSet<'static>,
manager: ConnectionManager,
device: VirtTunDevice,
to_smoltcp_tx: UnboundedSender<TokenBuffer>,
from_smoltcp_rx: Option<UnboundedReceiver<TokenBuffer>>,
avail: Arc<AtomicBool>,
rx_from_tunnel: Option<Receiver<RawCastFrame>>,
device: ChannelDevice,
/// Bounded channel from TUN reader task to engine loop.
tun_rx: Option<mpsc::Receiver<Vec<u8>>>,
/// Bounded channel from engine loop to TUN writer task.
tun_tx: mpsc::Sender<Vec<u8>>,
rx_from_tunnel: Option<mpsc::Receiver<RawCastFrame>>,
factory: Arc<dyn SocketProvider>,
stats: TrafficCounter,
}
impl Engine {
@@ -48,28 +57,32 @@ impl Engine {
config: Config,
caps: DeviceCapabilities,
dns_handler: DnsHandler,
tx_to_tunnel: Sender<RawCastFrame>,
rx_from_tunnel: Receiver<RawCastFrame>,
tx_to_tunnel: mpsc::Sender<RawCastFrame>,
rx_from_tunnel: mpsc::Receiver<RawCastFrame>,
factory: Arc<dyn SocketProvider>,
) -> Self {
let now = Engine::current_time();
let (mut device, to_smoltcp_tx, from_smoltcp_rx, avail) = VirtTunDevice::new(caps);
let mut device = ChannelDevice::new(caps, DEVICE_RX_CAP, DEVICE_TX_CAP);
let interface = Interface::new(config, &mut device, now);
let socket_set = ConnectionManager::setup_sockets(factory.as_ref(), 2);
let manager = ConnectionManager::new(dns_handler, tx_to_tunnel, factory.clone());
// Bounded TUN writer channel — smoltcp's TCP window limits how many
// outgoing packets can pile up, so a modest cap is enough.
let (tun_tx, _placeholder) = mpsc::channel(DEVICE_TX_CAP * 2);
Self {
interface,
socket_set,
device,
to_smoltcp_tx,
from_smoltcp_rx: Some(from_smoltcp_rx),
avail,
tun_rx: None,
tun_tx,
manager,
rx_from_tunnel: Some(rx_from_tunnel),
factory,
stats: TrafficCounter::new(),
}
}
@@ -77,129 +90,163 @@ impl Engine {
info!("Current routes: {:?}", self.interface.routes());
let (writer, reader) = tun.split().expect("Failed to split TUN");
let (tun_to_engine_tx, mut tun_to_engine_rx) = mpsc::unbounded_channel::<TokenBuffer>();
// Bounded: TUN reader blocks when engine is overloaded → kernel TUN
// buffer fills → natural backpressure to the OS.
let (tun_to_engine_tx, tun_to_engine_rx) = mpsc::channel::<Vec<u8>>(TUN_CHAN_CAP);
// Bounded: engine drops TX packets if TUN writer is slow (TCP retransmits).
let (engine_to_tun_tx, engine_to_tun_rx) = mpsc::channel::<Vec<u8>>(DEVICE_TX_CAP * 2);
Self::spawn_tun_reader(reader, tun_to_engine_tx, self.avail.clone());
self.tun_tx = engine_to_tun_tx;
self.tun_rx = Some(tun_to_engine_rx);
let from_smoltcp_rx = self.from_smoltcp_rx.take().expect("Engine started twice");
Self::spawn_tun_writer(writer, from_smoltcp_rx);
Self::spawn_tun_reader(reader, tun_to_engine_tx);
Self::spawn_tun_writer(writer, engine_to_tun_rx);
let mut last_stats_log = StdInstant::now();
let inbound_map = self.manager.tracker.inbound_tx.clone();
let mut rx_tunnel = self.rx_from_tunnel.take().unwrap();
let mut tun_rx = self.tun_rx.take().unwrap();
tokio::spawn(async move {
let mut local_cache: std::collections::HashMap<
u64,
(mpsc::Sender<Bytes>, Arc<AtomicBool>),
> = std::collections::HashMap::new();
while let Some(frame) = rx_tunnel.recv().await {
if frame.event == RawCastEvent::Close {
local_cache.remove(&frame.socket_id);
inbound_map.remove(&frame.socket_id);
} else if frame.event == RawCastEvent::Data {
let tx_opt: Option<(mpsc::Sender<Bytes>, Arc<AtomicBool>)> =
if let Some(cached) = local_cache.get(&frame.socket_id) {
Some(cached.clone())
} else if let Some(ref_tx) = inbound_map.get(&frame.socket_id) {
let val = ref_tx.value().clone();
local_cache.insert(frame.socket_id, val.clone());
Some(val)
} else {
None
};
if let Some((tx, is_saturated)) = tx_opt {
// Pressure-Aware drop (нулевая нагрузка на ОС)
if is_saturated.load(Ordering::Relaxed) {
continue;
}
match tx.try_send(frame.payload) {
Ok(_) => {}
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
netrunner_logger::trace!(
"⚠️ Local socket {} buffer suddenly full. Dropping frame.",
frame.socket_id
);
}
Err(_) => {}
}
}
}
}
});
// Local cache to avoid DashMap lookups on every frame.
let mut local_cache: std::collections::HashMap<
u64,
(mpsc::Sender<Bytes>, Arc<std::sync::atomic::AtomicBool>),
> = std::collections::HashMap::new();
loop {
let mut work_done = false;
let mut repeat_poll = true;
let now = Self::current_time();
let mut work_done = false;
// 1. Опрашиваем сокеты до тех пор, пока есть движение данных
while repeat_poll {
self.manager.process_sockets(&mut self.socket_set, now);
let poll_res = self.poll();
self.manager.cleanup(&mut self.socket_set);
if matches!(poll_res, PollResult::SocketStateChanged) {
work_done = true;
repeat_poll = true;
} else {
repeat_poll = false;
// ── 1. Dispatch tunnel → local sockets (download) ────────────
// Drain all available frames without blocking.
loop {
match rx_tunnel.try_recv() {
Ok(frame) => {
work_done = true;
if frame.event == RawCastEvent::Close {
local_cache.remove(&frame.socket_id);
inbound_map.remove(&frame.socket_id);
} else if frame.event == RawCastEvent::Data {
let tx_opt = if let Some(cached) = local_cache.get(&frame.socket_id) {
Some(cached.clone())
} else if let Some(ref_tx) = inbound_map.get(&frame.socket_id) {
let val = ref_tx.value().clone();
local_cache.insert(frame.socket_id, val.clone());
Some(val)
} else {
None
};
if let Some((tx, is_saturated)) = tx_opt {
if is_saturated.load(Ordering::Relaxed) {
continue;
}
match tx.try_send(frame.payload) {
Ok(_) => {}
Err(mpsc::error::TrySendError::Full(data)) => {
// Spawn a task to wait for space rather than drop.
let tx2 = tx.clone();
tokio::spawn(async move {
let _ = tx2.send(data).await;
});
}
Err(_) => {}
}
}
}
}
Err(_) => break,
}
}
// ── 2. Accept TUN packets → device (upload) ──────────────────
// Stop reading when the device's RX queue is full (backpressure).
let mut packets_read = 0;
while !self.device.rx_full() {
match tun_rx.try_recv() {
Ok(pkt) => {
self.stats.record_rx(pkt.len());
self.manager
.try_create_socket_from_packet(&pkt, &mut self.socket_set);
self.device.push_rx(pkt);
work_done = true;
packets_read += 1;
if packets_read >= 250 {
break; // Yield occasionally to prevent starvation.
}
}
Err(_) => break,
}
}
// ── 3. Run smoltcp ───────────────────────────────────────────
let mut repeat = true;
while repeat {
self.manager.process_sockets(&mut self.socket_set, now);
let res = self.poll();
self.manager.cleanup(&mut self.socket_set);
repeat = matches!(res, PollResult::SocketStateChanged);
if repeat {
work_done = true;
}
}
// ── 4. Drain smoltcp TX → TUN writer ─────────────────────────
while let Some(pkt) = self.device.pop_tx() {
self.stats.record_tx(pkt.len());
// Non-blocking: if TUN writer is overloaded, drop the packet.
// TCP will retransmit; UDP is best-effort.
if self.tun_tx.try_send(pkt).is_err() {
break;
}
}
// ── 5. Stats logging ─────────────────────────────────────────
if last_stats_log.elapsed() >= Duration::from_secs(5) {
let stats = self.stats.get_stats();
info!(
"TunDevice Traffic: RX: {:.2} MB ({} pkts) | TX: {:.2} MB ({} pkts) | Speed: ↓{:.2} MB/s, ↑{:.2} MB/s",
stats.rx_bytes as f64 / 1_048_576.0,
stats.rx_packets,
stats.tx_bytes as f64 / 1_048_576.0,
stats.tx_packets,
stats.rx_speed_mb_s,
stats.tx_speed_mb_s,
);
let manager_ref = &self.manager;
self.factory
.log_stats(&self.socket_set, &|handle| manager_ref.get_buf_info(handle));
last_stats_log = StdInstant::now();
}
// 2. БАТЧИНГ: Выгребаем все доступные пакеты из TUN без блокировки
let mut packets_processed = 0;
while let Ok(token) = tun_to_engine_rx.try_recv() {
self.manager
.try_create_socket_from_packet(&token, &mut self.socket_set);
if self.to_smoltcp_tx.send(token).is_ok() {
self.device.mark_rx_available();
}
work_done = true;
packets_processed += 1;
// Предотвращаем starvation: если пакетов слишком много, даем другим потокам подышать
if packets_processed > 250 {
break;
}
}
// 3. АДАПТИВНЫЙ ТАЙМИНГ
// ── 6. Adaptive timing ───────────────────────────────────────
if work_done {
// Если мы обработали пакеты или сокеты сдвинулись - НЕ СПИМ.
// Уступаем квант времени другим задачам и моментально возвращаемся в цикл.
tokio::task::yield_now().await;
} else {
// Мы отдыхаем только если сеть полностью простаивает
let delay = self
.interface
.poll_delay(Self::current_time(), &self.socket_set);
let sleep_time = delay
.map(|d| {
std::cmp::min(Duration::from_micros(d.micros()), Duration::from_millis(5))
std::cmp::min(
Duration::from_micros(d.micros()),
Duration::from_millis(5),
)
})
.unwrap_or(Duration::from_millis(5));
tokio::select! {
_ = sleep(sleep_time) => {}
msg = tun_to_engine_rx.recv() => {
if let Some(token) = msg {
self.manager.try_create_socket_from_packet(&token, &mut self.socket_set);
if self.to_smoltcp_tx.send(token).is_ok() {
self.device.mark_rx_available();
msg = tun_rx.recv() => {
match msg {
Some(pkt) => {
self.stats.record_rx(pkt.len());
self.manager.try_create_socket_from_packet(&pkt, &mut self.socket_set);
self.device.push_rx(pkt);
}
} else {
break;
None => break,
}
}
}
@@ -213,36 +260,23 @@ impl Engine {
.poll(now, &mut self.device, &mut self.socket_set)
}
fn spawn_tun_reader(
mut reader: DeviceReader,
to_engine: mpsc::UnboundedSender<TokenBuffer>,
is_avail: Arc<AtomicBool>,
) {
fn spawn_tun_reader(mut reader: DeviceReader, to_engine: mpsc::Sender<Vec<u8>>) {
tokio::spawn(async move {
debug!("TUN Reader task started");
let mut buf = vec![0u8; 65536];
loop {
// 1. Берем чистый буфер
let mut token = TokenBuffer::with_capacity(65536);
// 2. Используем временный срез для чтения
// Это гарантирует, что мы пишем в начало буфера
let mut temp_buf = [0u8; 65536];
match reader.read(&mut temp_buf).await {
match reader.read(&mut buf).await {
Ok(n) if n > 0 => {
// 3. Копируем в наш TokenBuffer (это быстро, 64кБ - это L1 кэш)
token.extend_from_slice(&temp_buf[..n]);
if to_engine.send(token).is_ok() {
is_avail.store(true, Ordering::Release);
} else {
let pkt = buf[..n].to_vec();
// .send().await blocks when engine channel is full →
// backpressure propagates to OS TUN device.
if to_engine.send(pkt).await.is_err() {
break;
}
}
Ok(_) => break, // EOF
Ok(_) => break,
Err(e) => {
error!("FATAL: TUN Reader task died with error: {}", e);
error!("FATAL: TUN Reader task died: {}", e);
break;
}
}
@@ -250,14 +284,11 @@ impl Engine {
});
}
fn spawn_tun_writer(
mut writer: DeviceWriter,
mut from_smoltcp: UnboundedReceiver<TokenBuffer>,
) {
fn spawn_tun_writer(mut writer: DeviceWriter, mut from_engine: mpsc::Receiver<Vec<u8>>) {
tokio::spawn(async move {
debug!("TUN Writer task started");
while let Some(token) = from_smoltcp.recv().await {
if writer.write_all(&token).await.is_err() {
while let Some(pkt) = from_engine.recv().await {
if writer.write_all(&pkt).await.is_err() {
break;
}
}
@@ -265,7 +296,7 @@ impl Engine {
});
}
fn current_time() -> Instant {
pub fn current_time() -> Instant {
let duration = StdInstant::now().duration_since(*START_TIME);
Instant::from_micros(duration.as_micros() as i64)
}
@@ -301,6 +332,8 @@ impl Engine {
}
}
// ─── EngineConfig & EngineBuilder (unchanged API surface) ──────────────────
#[derive(Clone, Debug)]
pub struct EngineConfig {
pub remote_address: String,
@@ -457,7 +490,10 @@ impl EngineBuilder {
);
#[cfg(target_os = "windows")]
let _ = crate::tun::routing::run_cmd_ext(
&format!("route add {} mask 255.255.255.255 {}", ipv4, phys_gw),
&format!(
"route add {} mask 255.255.255.255 {}",
ipv4, phys_gw
),
true,
);
}
@@ -480,7 +516,6 @@ impl EngineBuilder {
if self.config.transparent_mode {
engine.set_transparent_mode();
}
engine.set_default_gateway(self.config.default_gateway);
engine.activate();