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

backpressure, bufferbloat fix and legs reconnects

Browse Source
This commit is contained in:
Kirill
2026-04-14 16:35:59 +07:00
parent 6d3bc81a6e
commit a6ca710954
10 changed files with 200 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/src/net/connection.rs
+46 -44
View File
@@ -19,20 +19,15 @@ use netrunner_logger::{debug, info};
pub struct ConnectionCore<T> {
pub handle: SocketHandle,
pub tx: mpsc::UnboundedSender<T>, // 🔥 Полностью Unbounded
pub rx: mpsc::UnboundedReceiver<Bytes>, // 🔥 Полностью Unbounded
pub tx: mpsc::Sender<T>,
pub rx: mpsc::Receiver<Bytes>,
}
impl<T> ConnectionCore<T> {
pub fn new(
handle: SocketHandle,
) -> (
Self,
mpsc::UnboundedReceiver<T>,
mpsc::UnboundedSender<Bytes>,
) {
let (tx_to_net, rx_from_smol) = mpsc::unbounded_channel::<T>();
let (tx_to_smol, rx_from_net) = mpsc::unbounded_channel::<Bytes>();
pub fn new(handle: SocketHandle) -> (Self, mpsc::Receiver<T>, mpsc::Sender<Bytes>) {
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);
let core = Self {
handle,
@@ -56,6 +51,7 @@ pub struct TcpConnection {
core: ConnectionCore<Bytes>,
state: ConnectionState,
pending_data: VecDeque<Bytes>,
pending_bytes: usize,
handshake_rx: Option<oneshot::Receiver<()>>,
chunk_buf: Vec<u8>,
server_eof: bool,
@@ -66,8 +62,8 @@ impl TcpConnection {
handle: SocketHandle,
) -> (
Self,
mpsc::UnboundedReceiver<Bytes>,
mpsc::UnboundedSender<Bytes>,
mpsc::Receiver<Bytes>,
mpsc::Sender<Bytes>,
oneshot::Sender<()>,
) {
let (core, rx_from_smol, tx_to_smol) = ConnectionCore::new(handle);
@@ -77,6 +73,7 @@ impl TcpConnection {
core,
state: ConnectionState::Handshaking,
pending_data: VecDeque::new(),
pending_bytes: 0,
handshake_rx: Some(handshake_rx),
chunk_buf: vec![0u8; NetworkConfig::global().tcp_chunk_size],
server_eof: false,
@@ -139,7 +136,8 @@ impl TcpConnection {
let rtt = Duration::from_millis(current_rtt as u64);
socket.set_tunnel_rtt(rtt);
if self.app_pending_out_size() < 512 * 1024 {
// Читаем из браузера в Туннель
if self.core.tx.capacity() > 0 {
while socket.can_recv() {
if let Ok(n) = socket.peek_slice(&mut self.chunk_buf) {
if n == 0 {
@@ -147,11 +145,17 @@ impl TcpConnection {
}
let chunk = Bytes::copy_from_slice(&self.chunk_buf[..n]);
if self.core.tx.send(chunk).is_ok() {
socket.recv_slice(&mut self.chunk_buf[..n]).unwrap();
} else {
self.server_eof = true;
break;
match self.core.tx.try_send(chunk) {
Ok(_) => {
socket.recv_slice(&mut self.chunk_buf[..n]).unwrap();
}
Err(mpsc::error::TrySendError::Full(_)) => {
break;
}
Err(mpsc::error::TrySendError::Closed(_)) => {
self.server_eof = true;
break;
}
}
} else {
break;
@@ -159,10 +163,15 @@ impl TcpConnection {
}
}
// Читаем из Туннеля в браузер
if !self.server_eof {
loop {
// 🔥 РЕАЛЬНЫЙ BACKPRESSURE:
// Читаем из канала, ТОЛЬКО если в очереди < 4 МБ данных.
// Если больше - оставляем лежать в канале Tokio, пока smoltcp не освободится!
while self.pending_bytes < 4 * 1024 * 1024 {
match self.core.rx.try_recv() {
Ok(data) => {
self.pending_bytes += data.len();
self.pending_data.push_back(data);
}
Err(mpsc::error::TryRecvError::Empty) => break,
@@ -179,6 +188,7 @@ impl TcpConnection {
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);
@@ -205,22 +215,24 @@ impl TcpConnection {
}
}
pub fn app_pending_out_size(&self) -> usize {
self.pending_bytes
}
pub fn spawn(
socket_id: u64,
dst_ip: std::net::Ipv4Addr,
dst_port: u16,
target: String,
mut rx_smol: mpsc::UnboundedReceiver<Bytes>,
mut rx_smol: mpsc::Receiver<Bytes>,
handshake_tx: oneshot::Sender<()>,
// 🔥 ФИКС: Тип изменен на UnboundedSender
tx_tunnel: mpsc::UnboundedSender<RawCastFrame>,
tx_tunnel: mpsc::Sender<RawCastFrame>,
) {
tokio::spawn(async move {
let mut frame = RawCastFrame::connect(LocalProtocol::Tcp, socket_id, dst_ip, dst_port);
frame.payload = Bytes::from(target);
// send() на UnboundedSender не асинхронный и возвращает Result<(), T>
if tx_tunnel.send(frame).is_err() {
if tx_tunnel.send(frame).await.is_err() {
netrunner_logger::error!("❌ [TCP {}] Failed to send CONNECT to tunnel", socket_id);
return;
}
@@ -236,21 +248,17 @@ impl TcpConnection {
data.to_vec(),
);
if tx_tunnel.send(data_frame).is_err() {
if tx_tunnel.send(data_frame).await.is_err() {
break;
}
}
let close_frame = RawCastFrame::close(LocalProtocol::Tcp, socket_id, dst_ip, dst_port);
let _ = tx_tunnel.send(close_frame);
let _ = tx_tunnel.send(close_frame).await;
debug!("🏁 [TCP {}] Spawned task finished", socket_id);
});
}
pub fn app_pending_out_size(&self) -> usize {
self.pending_data.iter().map(|b| b.len()).sum()
}
}
pub type UdpPacketTarget = (Bytes, std::net::Ipv4Addr, u16);
@@ -266,11 +274,7 @@ impl UdpConnection {
handle: SocketHandle,
client_addr: smoltcp::wire::IpAddress,
client_port: u16,
) -> (
Self,
mpsc::UnboundedReceiver<UdpPacketTarget>,
mpsc::UnboundedSender<Bytes>,
) {
) -> (Self, mpsc::Receiver<UdpPacketTarget>, mpsc::Sender<Bytes>) {
let (core, rx_from_smol, tx_to_smol) = ConnectionCore::new(handle);
let conn = Self {
core,
@@ -299,7 +303,7 @@ impl UdpConnection {
let target_port = metadata.endpoint.port;
let payload = (Bytes::copy_from_slice(data), target_ip, target_port);
if self.core.tx.send(payload).is_ok() {
if self.core.tx.try_send(payload).is_ok() {
self.last_activity = std::time::Instant::now();
}
}
@@ -328,16 +332,15 @@ impl UdpConnection {
dst_ip: std::net::Ipv4Addr,
dst_port: u16,
target: String,
mut rx_smol: mpsc::UnboundedReceiver<UdpPacketTarget>,
// 🔥 ФИКС: Тип изменен на UnboundedSender
tx_tunnel: mpsc::UnboundedSender<RawCastFrame>,
mut rx_smol: mpsc::Receiver<UdpPacketTarget>,
tx_tunnel: mpsc::Sender<RawCastFrame>,
) {
tokio::spawn(async move {
debug!("📡 [UDP {}] Task started for {}", socket_id, target);
let mut frame = RawCastFrame::connect(LocalProtocol::Udp, socket_id, dst_ip, dst_port);
frame.payload = Bytes::from(target);
if tx_tunnel.send(frame).is_err() {
if tx_tunnel.send(frame).await.is_err() {
netrunner_logger::error!("❌ [UDP {}] Failed to send CONNECT to tunnel", socket_id);
return;
}
@@ -345,19 +348,18 @@ impl UdpConnection {
while let Some((data, ip, port)) = rx_smol.recv().await {
let data_frame =
RawCastFrame::data(LocalProtocol::Udp, socket_id, ip, port, data.to_vec());
if tx_tunnel.send(data_frame).is_err() {
if tx_tunnel.send(data_frame).await.is_err() {
break;
}
}
let close_frame = RawCastFrame::close(LocalProtocol::Udp, socket_id, dst_ip, dst_port);
let _ = tx_tunnel.send(close_frame);
let _ = tx_tunnel.send(close_frame).await;
info!("🛑 [UDP {}] Task stopped", socket_id);
});
}
}
// IcmpResponder оставляем как был (он не использует каналы)
use smoltcp::socket::icmp;
pub struct IcmpResponder;
client/src/net/connection_manager.rs
+12 -17
View File
@@ -13,7 +13,7 @@ use smoltcp::{
wire::{IpAddress, IpListenEndpoint, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket},
};
use std::sync::Arc;
use std::time::Instant; // Используем DashMap для потокобезопасного трекинга
use std::time::Instant;
use tokio::sync::mpsc;
@@ -24,7 +24,6 @@ use crate::net::{
socket_factory::SocketProvider,
};
// Ключ для идентификации уникального потока (4-tuple)
type FlowKey = (IpAddress, u16, IpAddress, u16);
struct Flow {
@@ -78,7 +77,7 @@ impl TargetResolver {
pub struct ConnectionManager {
tracker: SessionTracker,
resolver: TargetResolver,
tx_to_tunnel: mpsc::UnboundedSender<RawCastFrame>,
tx_to_tunnel: mpsc::Sender<RawCastFrame>, // 🔥 Bounded
factory: Arc<dyn SocketProvider>,
pending_connects: DashMap<FlowKey, Instant>,
}
@@ -86,7 +85,7 @@ pub struct ConnectionManager {
impl ConnectionManager {
pub fn new(
dns_handler: DnsHandler,
tx_to_tunnel: mpsc::UnboundedSender<RawCastFrame>, // 🔥 Unbounded
tx_to_tunnel: mpsc::Sender<RawCastFrame>,
factory: Arc<dyn SocketProvider>,
) -> Self {
Self {
@@ -109,9 +108,15 @@ impl ConnectionManager {
}
if let Some(tx) = self.tracker.get_inbound_tx(frame.socket_id) {
// 🔥 UnboundedSender::send не требует await и не возвращает Full
if let Err(_) = tx.send(frame.payload.clone()) {
self.tracker.close_tunnel_session(frame.socket_id);
// 🔥 Защита от переполнения: Если очередь Bounded канала полна,
// мы тихо удаляем пакет. Очередь освободится, и TCP/BBR сделает свое дело (надежная доставка).
// Нельзя закрывать соединение!
if let Err(e) = tx.try_send(frame.payload.clone()) {
if let mpsc::error::TrySendError::Closed(_) = e {
self.tracker.close_tunnel_session(frame.socket_id);
} else {
// Backpressure in action: drop the packet, TCP will handle it
}
}
Ok(())
} else {
@@ -174,7 +179,6 @@ impl ConnectionManager {
if let Ok(p) = TcpPacket::new_checked(ip.payload()) {
flow.src_p = p.src_port();
flow.dst_p = p.dst_port();
// Игнорируем пакеты, если это не SYN (новое соединение)
if p.syn() && !p.ack() {
self.intercept_tcp(flow, socket_set);
}
@@ -262,7 +266,6 @@ impl ConnectionManager {
if socket.is_open() {
let handle = socket_set.add(socket);
// 🔥 Принимаем Unbounded части из конструктора
let (conn, rx_smol, tx_smol) = UdpConnection::new(handle, f.src, f.src_p);
self.tracker.register_udp(handle, socket_id, conn, tx_smol);
@@ -278,11 +281,9 @@ impl ConnectionManager {
}
pub fn process_sockets(&mut self, socket_set: &mut SocketSet) {
// Собираем хендлы, чтобы избежать конфликтов заимствования (borrow checker)
let handles: Vec<SocketHandle> = socket_set.iter().map(|(h, _)| h).collect();
for handle in handles {
// 🔥 ИСПРАВЛЕНИЕ: get_mut возвращает &mut Socket, а не Option<&mut Socket>
let socket = socket_set.get_mut(handle);
match socket {
@@ -311,7 +312,6 @@ impl ConnectionManager {
let socket_id = self.tracker.next_id();
let (dst_ip, target) = self.resolver.resolve_destination(local.addr, local.port);
// 🔥 Принимаем Unbounded части
let (conn, rx_smol, tx_smol, handshake_tx) = TcpConnection::new(handle);
self.tracker.register_tcp(handle, socket_id, conn, tx_smol);
@@ -350,7 +350,6 @@ impl ConnectionManager {
}
pub fn cleanup(&mut self, socket_set: &mut SocketSet) {
// 🔥 ФИКС 3: Очистка зависших pending_connects (например, если SYN потерялся)
self.pending_connects
.retain(|_, timestamp| timestamp.elapsed() < TCP_HANDSHAKE_TIMEOUT);
@@ -362,7 +361,3 @@ impl ConnectionManager {
self.tracker.get_app_buffer_info(handle)
}
}
fn is_handshaking(state: tcp::State) -> bool {
matches!(state, tcp::State::SynSent | tcp::State::SynReceived)
}
client/src/net/engine.rs
+10 -13
View File
@@ -15,8 +15,7 @@ use std::{
time::Instant as StdInstant,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::mpsc::UnboundedReceiver;
use tokio::sync::mpsc::{self, UnboundedSender}; // 🔥 Добавили UnboundedSender
use tokio::sync::mpsc::{self, Receiver, Sender, UnboundedReceiver, UnboundedSender};
use tokio::time::{Duration, sleep};
use tun::{DeviceReader, DeviceWriter};
@@ -36,10 +35,10 @@ pub struct Engine {
socket_set: SocketSet<'static>,
manager: ConnectionManager,
device: VirtTunDevice,
to_smoltcp_tx: UnboundedSender<TokenBuffer>, // 🔥 Стало Unbounded
to_smoltcp_tx: UnboundedSender<TokenBuffer>,
from_smoltcp_rx: Option<UnboundedReceiver<TokenBuffer>>,
avail: Arc<AtomicBool>,
rx_from_tunnel: mpsc::UnboundedReceiver<RawCastFrame>, // 🔥 Стало Unbounded
rx_from_tunnel: Receiver<RawCastFrame>, // 🔥 Bounded
factory: Arc<dyn SocketProvider>,
}
@@ -48,13 +47,12 @@ impl Engine {
config: Config,
caps: DeviceCapabilities,
dns_handler: DnsHandler,
tx_to_tunnel: mpsc::UnboundedSender<RawCastFrame>, // 🔥 Стало Unbounded
rx_from_tunnel: mpsc::UnboundedReceiver<RawCastFrame>, // 🔥 Стало Unbounded
tx_to_tunnel: Sender<RawCastFrame>, // 🔥 Bounded
rx_from_tunnel: Receiver<RawCastFrame>, // 🔥 Bounded
factory: Arc<dyn SocketProvider>,
) -> Self {
let now = Engine::current_time();
// VirtTunDevice::new возвращает UnboundedSender и UnboundedReceiver
let (mut device, to_smoltcp_tx, from_smoltcp_rx, avail) = VirtTunDevice::new(caps);
let interface = Interface::new(config, &mut device, now);
@@ -79,7 +77,6 @@ impl Engine {
info!("Current routes: {:?}", self.interface.routes());
let (writer, reader) = tun.split().expect("Failed to split TUN");
// TUN Reader теперь тоже Unbounded
let (tun_to_engine_tx, mut tun_to_engine_rx) = mpsc::unbounded_channel::<TokenBuffer>();
Self::spawn_tun_reader(reader, tun_to_engine_tx, self.avail.clone());
@@ -118,7 +115,6 @@ impl Engine {
msg = self.rx_from_tunnel.recv() => {
if let Some(frame) = msg {
let _ = self.manager.try_inject_inbound(frame);
// Выгребаем пачку
let mut count = 0;
while let Ok(frame) = self.rx_from_tunnel.try_recv() {
let _ = self.manager.try_inject_inbound(frame);
@@ -131,7 +127,6 @@ impl Engine {
msg = tun_to_engine_rx.recv() => {
if let Some(token) = msg {
self.manager.try_create_socket_from_packet(&token, &mut self.socket_set);
// Отправка в Unbounded всегда мгновенная
if self.to_smoltcp_tx.send(token).is_ok() {
self.device.mark_rx_available();
}
@@ -149,7 +144,7 @@ impl Engine {
fn spawn_tun_reader(
mut reader: DeviceReader,
to_engine: mpsc::UnboundedSender<TokenBuffer>, // 🔥 Unbounded
to_engine: mpsc::UnboundedSender<TokenBuffer>,
is_avail: Arc<AtomicBool>,
) {
tokio::spawn(async move {
@@ -303,8 +298,10 @@ impl EngineBuilder {
caps.max_transmission_unit = self.config.mtu;
caps.medium = smoltcp::phy::Medium::Ip;
let (tx_to_tunnel, rx_for_client_handler) = mpsc::unbounded_channel::<RawCastFrame>();
let (tx_for_client_handler, rx_from_tunnel) = mpsc::unbounded_channel::<RawCastFrame>();
// 🔥 Используем Bounded каналы для Tunnel <-> Engine
let cap = NetworkConfig::global().channel_capacity;
let (tx_to_tunnel, rx_for_client_handler) = mpsc::channel::<RawCastFrame>(cap);
let (tx_for_client_handler, rx_from_tunnel) = mpsc::channel::<RawCastFrame>(cap);
info!("Establishing secure tunnel to proxy server...");
ClientHandler::connect(
client/src/net/session_tracker.rs
+4 -5
View File
@@ -17,8 +17,7 @@ pub struct SessionTracker {
last_activity: HashMap<SocketHandle, StdInstant>,
active_tcp: HashMap<SocketHandle, TcpConnection>,
active_udp: HashMap<SocketHandle, UdpConnection>,
// 🔥 Теперь храним UnboundedSender
inbound_tx: HashMap<u64, mpsc::UnboundedSender<Bytes>>,
inbound_tx: HashMap<u64, mpsc::Sender<Bytes>>, // 🔥 Bounded
handle_to_id: HashMap<SocketHandle, u64>,
id_to_handle: HashMap<u64, SocketHandle>,
@@ -57,7 +56,7 @@ impl SessionTracker {
handle: SocketHandle,
id: u64,
conn: TcpConnection,
tx: mpsc::UnboundedSender<Bytes>, // 🔥 Unbounded
tx: mpsc::Sender<Bytes>,
) {
self.pending_tcp.remove(&handle);
self.handle_to_id.insert(handle, id);
@@ -72,7 +71,7 @@ impl SessionTracker {
handle: SocketHandle,
id: u64,
conn: UdpConnection,
tx: mpsc::UnboundedSender<Bytes>, // 🔥 Unbounded
tx: mpsc::Sender<Bytes>,
) {
self.handle_to_id.insert(handle, id);
self.id_to_handle.insert(id, handle);
@@ -123,7 +122,7 @@ impl SessionTracker {
self.last_activity.insert(handle, StdInstant::now());
}
pub fn get_inbound_tx(&self, id: u64) -> Option<&mpsc::UnboundedSender<Bytes>> {
pub fn get_inbound_tx(&self, id: u64) -> Option<&mpsc::Sender<Bytes>> {
self.inbound_tx.get(&id)
}