stable connection

This commit is contained in:
2026-04-03 18:53:13 +07:00
parent dfac8ebc23
commit 2db4ba1e9d
14 changed files with 125 additions and 204 deletions
+2 -2
View File
@@ -59,7 +59,7 @@ pub(crate) async fn run_tcp_bridge<R, W>(
data: buf.split().freeze(),
};
if muxer.send_to_network(msg).is_err() {
if muxer.send_to_network(msg).await.is_err() {
return Ok(false);
}
Ok(true)
@@ -139,7 +139,7 @@ pub(crate) async fn run_udp_bridge(
match res {
Ok(n) if n > 0 => {
let data = Bytes::copy_from_slice(&buf[..n]);
if let Err(e) = muxer.send_data_safe(stream_id, data, true) {
if let Err(e) = muxer.send_data_safe(stream_id, data, true).await {
error!(stream_id, "UDP Failed to send to tunnel: {}", e);
return Err(e);
}
+54 -73
View File
@@ -89,11 +89,9 @@ impl Connection {
}
}
// Измененный контекст: теперь храним оригинальный (локальный) socket_id
type StreamContext = (u32, Ipv4Addr, u16, LocalProtocol);
pub struct ClientHandler;
impl ClientHandler {
async fn establish_leg(
remote_proxy_addr: &str,
@@ -197,9 +195,7 @@ impl ClientHandler {
let muxer = Arc::new(Muxer::new(true, session_id.clone()));
// Реестр для отправки данных обратно в виртуальный сетевой стек
let registry: Arc<DashMap<u32, StreamContext>> = Arc::new(DashMap::new());
// Мапа для трансляции: локальный socket_id -> глобальный stream_id
let local_to_global: Arc<DashMap<u32, u32>> = Arc::new(DashMap::new());
info!("🚀 Netrunner Multi-Path Tunnel Initializing (10 Legs max).");
@@ -228,10 +224,7 @@ impl ClientHandler {
break;
}
// 🔥 ЗАПУСК HEALTH CHECK: Пингуем все леги
m_stats.perform_health_check().await;
// После пинга выводим красивую топологию
m_stats.print_topology_tree();
}
});
@@ -239,21 +232,18 @@ impl ClientHandler {
let muxer_inner = muxer.clone();
tokio::spawn(async move {
while let Some(raw_frame) = rx_from_engine.recv().await {
// Извлекаем нужные метаданные до передачи во фрейм
let dst_ip = raw_frame.dst_ip;
let dst_port = raw_frame.dst_port;
let protocol = raw_frame.protocol;
let is_udp = protocol == LocalProtocol::Udp;
if let Ok(nrxp_frame) = RawCastAdapter::to_nrxp(raw_frame) {
// Адаптер кладет socket_id от smoltcp в заголовок stream_id
let local_socket_id = nrxp_frame.header.stream_id;
let f_type = nrxp_frame.header.frame_type;
let payload = nrxp_frame.payload;
match f_type {
FrameType::Connect | FrameType::UdpConnect => {
// 🔥 ШАГ 2: Выделяем глобальный ID для нового потока
let global_stream_id = muxer_inner.next_stream_id();
info!(
@@ -264,7 +254,6 @@ impl ClientHandler {
global_stream_id
);
// Запоминаем связь
local_to_global.insert(local_socket_id, global_stream_id);
registry.insert(
global_stream_id,
@@ -274,7 +263,6 @@ impl ClientHandler {
let (v_tx, mut v_rx) =
mpsc::channel(NetworkConfig::global().tcp_stream_capacity);
// Регистрируем локальный поток под глобальным ID
muxer_inner.register_stream(global_stream_id, v_tx);
let tx_to_tun = tx_to_engine.clone();
@@ -282,16 +270,16 @@ impl ClientHandler {
tokio::spawn(async move {
while let Some(back_payload) = v_rx.recv().await {
if let Some(ctx) = reg.get(&global_stream_id) {
let (orig_local_id, ip, port, proto) = *ctx;
// ✅ ИСПРАВЛЕНИЕ 1: Копируем данные и отпускаем DashMap до .await
let ctx = reg.get(&global_stream_id).map(|r| *r);
if let Some((orig_local_id, ip, port, proto)) = ctx {
let out_f_type = if proto == LocalProtocol::Udp {
FrameType::UdpData
} else {
FrameType::Data
};
// Оборачиваем обратно в оригинальный локальный ID,
// чтобы smoltcp и tun-адаптер поняли, кому это принадлежит
let mock_nrxp =
Frame::new(orig_local_id, out_f_type, back_payload);
@@ -307,14 +295,19 @@ impl ClientHandler {
}
});
let _ = muxer_inner.send_control(global_stream_id, f_type, payload);
// ✅ Добавлен .await
let _ = muxer_inner
.send_control(global_stream_id, f_type, payload)
.await;
}
FrameType::Data | FrameType::UdpData => {
// Ищем глобальный ID по локальному
if let Some(global_ref) = local_to_global.get(&local_socket_id) {
let global_stream_id = *global_ref;
let _ =
muxer_inner.send_data_safe(global_stream_id, payload, is_udp);
// ✅ ИСПРАВЛЕНИЕ 2: Сброс блокировки DashMap
let global_stream_id =
local_to_global.get(&local_socket_id).map(|r| *r);
if let Some(id) = global_stream_id {
// ✅ Добавлен .await
let _ = muxer_inner.send_data_safe(id, payload, is_udp).await;
} else {
warn!(
"Data received for unknown local socket_id: {}",
@@ -323,14 +316,16 @@ impl ClientHandler {
}
}
FrameType::Close => {
// .remove() сам по себе возвращает владение данными и снимает лок мгновенно,
// поэтому здесь дедлока не было, но мы добавляем .await для вызова Муксера.
if let Some((_, global_stream_id)) =
local_to_global.remove(&local_socket_id)
{
let _ = muxer_inner.send_control(
global_stream_id,
FrameType::Close,
Bytes::new(),
);
// ✅ Добавлен .await
let _ = muxer_inner
.send_control(global_stream_id, FrameType::Close, Bytes::new())
.await;
muxer_inner.remove_stream(global_stream_id);
registry.remove(&global_stream_id);
}
@@ -346,7 +341,6 @@ impl ClientHandler {
Ok(())
}
}
pub struct ServerHandler {
pub(crate) conn: Connection,
pub(crate) session_manager: Arc<SessionManager>,
@@ -414,7 +408,6 @@ impl TunnelHandler for ServerHandler {
let handshake_timeout = std::time::Duration::from_secs(1);
// 1. ФАЗА TLS HANDSHAKE
let hello = loop {
let buf_snapshot = self.conn.read_buf.clone().freeze();
match self
@@ -423,59 +416,50 @@ impl TunnelHandler for ServerHandler {
.make_server_handshake(&mut self.conn.read_buf)
{
Ok(b) => break b,
Err(e) => {
// 🔥 ВЫЗЫВАЕМ СТРАТЕГИЮ (Она заодно и залогирует ошибку)
match e.execute_strategy() {
ErrorAction::Wait => {
let res = tokio::time::timeout(
handshake_timeout,
self.conn.inbound.read_buf(&mut self.conn.read_buf),
)
.await;
Err(e) => match e.execute_strategy() {
ErrorAction::Wait => {
let res = tokio::time::timeout(
handshake_timeout,
self.conn.inbound.read_buf(&mut self.conn.read_buf),
)
.await;
match res {
Ok(Ok(0)) => return Err("Client closed".into()),
Ok(Ok(_)) => continue,
_ => {
// По таймауту тоже редиректим, так как сканеры часто
// присылают пол-пакета и "висят"
Self::handle_stealth_fallback(
self.conn.inbound,
self.conn.outbound,
buf_snapshot,
)
.await;
return Ok(());
}
match res {
Ok(Ok(0)) => return Err("Client closed".into()),
Ok(Ok(_)) => continue,
_ => {
Self::handle_stealth_fallback(
self.conn.inbound,
self.conn.outbound,
buf_snapshot,
)
.await;
return Ok(());
}
}
ErrorAction::Redirect => {
// Пришел валидный TLS пакет, но от чужака -> притворяемся сайтом
Self::handle_stealth_fallback(
self.conn.inbound,
self.conn.outbound,
buf_snapshot,
)
.await;
return Ok(());
}
ErrorAction::Drop => {
// Полный мусор, отключаем немедленно (экономим сокеты ОС)
return Err("Dropped by security strategy (TLS Phase)".into());
}
}
}
ErrorAction::Redirect => {
Self::handle_stealth_fallback(
self.conn.inbound,
self.conn.outbound,
buf_snapshot,
)
.await;
return Ok(());
}
ErrorAction::Drop => {
return Err("Dropped by security strategy (TLS Phase)".into());
}
},
}
};
// Отвечаем ServerHello
self.conn
.outbound
.write_all(&hello)
.await
.map_err(|e| e.to_string())?;
// 2. ФАЗА NETRUNNER HANDSHAKE (Внутри туннеля)
let (session_id, leg_id) = loop {
let n = tokio::time::timeout(
std::time::Duration::from_secs(3),
@@ -513,17 +497,15 @@ impl TunnelHandler for ServerHandler {
let buf_snapshot = self.conn.read_buf.clone().freeze();
match e.execute_strategy() {
ErrorAction::Redirect => {
// Если мы решаем прятать VPN даже на уровне неудачной криптографии
Self::handle_stealth_fallback(
self.conn.inbound,
self.conn.outbound,
buf_snapshot, // Отправляем зашифрованный мусор на целевой сервак (он сам его дропнет)
buf_snapshot,
)
.await;
return Ok(());
}
_ => {
// Wait или Drop на этапе Handshake лучше просто дропать
return Err("Dropped by security strategy (Auth Phase)".into());
}
}
@@ -531,7 +513,6 @@ impl TunnelHandler for ServerHandler {
}
};
// ... Инициализация Muxer и TunnelEngine (остается без изменений)
let muxer = self.session_manager.get_or_create(&session_id);
let (control_tx, control_rx) = mpsc::channel(NetworkConfig::global().muxer_capacity);
-1
View File
@@ -185,7 +185,6 @@ impl TunnelEngine {
}
}
} else {
// TCP - потоковый, режем на чанки по 16KB для совместимости с TLS Record
while !data.is_empty() {
let chunk_size =
std::cmp::min(data.len(), NetworkConfig::global().tcp_chunk_size);
+3 -9
View File
@@ -14,8 +14,6 @@ use crate::net::{
};
use crate::nrxp::{Frame, FrameType};
// --- ОТКРЫВАШКА ДЛЯ СЕРВЕРА (Remote) ---
// Этот парень реально идет в интернет через сокеты ОС
pub struct RemoteOpener {
pub muxer: Arc<Muxer>,
}
@@ -63,11 +61,9 @@ impl RemoteOpener {
}
}
// --- ЕДИНЫЙ ХЭНДЛЕР ---
pub(crate) struct StreamHandler {
muxer: Arc<Muxer>,
opener: Option<Arc<RemoteOpener>>, // На сервере Some, на клиенте None
opener: Option<Arc<RemoteOpener>>,
}
impl StreamHandler {
@@ -89,8 +85,7 @@ impl StreamHandler {
}
FrameType::Data | FrameType::UdpData => {
// Унифицировано: просто кидаем в локальный канал
self.muxer.dispatch_to_local(stream_id, frame.payload);
self.muxer.dispatch_to_local(stream_id, frame.payload).await;
}
FrameType::Close => {
@@ -105,7 +100,6 @@ impl StreamHandler {
let target = String::from_utf8_lossy(&payload).to_string();
if let Some(opener) = &self.opener {
// Создаем виртуальную трубу для этого стрима
let capacity = if is_udp {
NetworkConfig::global().udp_stream_capacity
} else {
@@ -125,7 +119,7 @@ impl StreamHandler {
stream_id,
"⚠️ [Tunnel] Rejected incoming connection to {} (Client mode)", target
);
// Если мы клиент, мы не принимаем входящие соединения из туннеля
let _ = self
.muxer
.send_control(stream_id, FrameType::Close, Bytes::new());
+48 -79
View File
@@ -1,14 +1,13 @@
use bytes::Bytes;
use dashmap::DashMap;
use netrunner_logger::{debug, info, warn};
use tokio::sync::mpsc::error::TrySendError;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc::error::TrySendError;
use tokio::sync::mpsc::Sender;
use crate::nrxp::FrameType;
#[derive(Default, Debug)]
pub struct LegStats {
pub tx_bytes: AtomicU64,
@@ -59,9 +58,7 @@ pub struct Muxer {
id_gen: Arc<IdGenerator>,
session_id: Arc<String>,
}
impl Muxer {
pub fn new(is_client: bool, session_id: String) -> Self {
Self {
legs: Arc::new(DashMap::new()),
@@ -77,7 +74,6 @@ impl Muxer {
control_tx: Sender<MuxMessage>,
data_tx: Sender<MuxMessage>,
) {
if self.legs.len() >= 10 {
warn!(
leg_id,
@@ -106,50 +102,47 @@ impl Muxer {
self.legs.len()
}
fn select_leg(&self, frame_type: &FrameType, stream_id: u32) -> Option<(u32, MuxLeg)> {
fn select_leg(&self, frame_type: &FrameType, stream_id: u32) -> Option<(u32, MuxLeg)> {
if self.legs.is_empty() {
return None;
}
let is_udp = matches!(frame_type, FrameType::UdpData | FrameType::UdpConnect);
let mut candidates: Vec<(u32, MuxLeg)> = self.legs
let mut candidates: Vec<(u32, MuxLeg)> = self
.legs
.iter()
.map(|kv| (*kv.key(), kv.value().clone()))
.filter(|(id, _)| {
if is_udp { id % 2 != 0 } else { id % 2 == 0 }
})
.filter(|(id, _)| if is_udp { id % 2 != 0 } else { id % 2 == 0 })
.collect();
if candidates.is_empty() {
candidates = self.legs.iter().map(|kv| (*kv.key(), kv.value().clone())).collect();
candidates = self
.legs
.iter()
.map(|kv| (*kv.key(), kv.value().clone()))
.collect();
}
candidates.sort_by_key(|(_, leg)| {
let rtt = leg.stats.rtt_ms.load(Ordering::Relaxed);
if rtt == 0 { 9999 } else { rtt }
if rtt == 0 {
9999
} else {
rtt
}
});
let pool_size = std::cmp::min(candidates.len(), 2);
let target = &candidates[stream_id as usize % pool_size];
Some(target.clone())
}
pub fn send_to_network(&self, message: MuxMessage) -> Result<(), String> {
pub async fn send_to_network(&self, message: MuxMessage) -> Result<(), String> {
let stream_id = message.stream_id;
let size = message.data.len() as u64;
let (leg_id, leg) = self
.select_leg(&message.frame_type, stream_id)
.ok_or_else(|| "MUXER: No active legs available".to_string())?;
@@ -158,42 +151,30 @@ impl Muxer {
FrameType::Connect
| FrameType::Close
| FrameType::UdpConnect
| FrameType::Handshake => &leg.control_tx,
_ => &leg.data_tx,
| FrameType::Handshake => leg.control_tx.clone(),
_ => leg.data_tx.clone(),
};
match target_tx.try_send(message) {
match target_tx.send(message).await {
Ok(_) => {
leg.stats.tx_bytes.fetch_add(size, Ordering::Relaxed);
if let Some(stream_ref) = self.streams.get(&stream_id) {
stream_ref
.value()
.1
.tx_bytes
.fetch_add(size, Ordering::Relaxed);
stream_ref.1.tx_bytes.fetch_add(size, Ordering::Relaxed);
}
Ok(())
}
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
warn!(
leg_id,
"MUXER: Network queue full! Dropping outbound packet."
);
Ok(())
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
Err(_) => {
self.remove_leg(leg_id);
Err(format!("MUXER: Leg {} died during send", leg_id))
}
}
}
pub fn send_data_safe(
pub async fn send_data_safe(
&self,
stream_id: u32,
data: Bytes, // Больше не mut
data: Bytes,
is_udp: bool,
) -> Result<(), String> {
let frame_type = if is_udp {
@@ -202,15 +183,15 @@ impl Muxer {
FrameType::Data
};
// Отправляем как есть, целиком!
self.send_to_network(MuxMessage {
stream_id,
frame_type,
data,
})
.await
}
pub fn send_control(
pub async fn send_control(
&self,
stream_id: u32,
f_type: FrameType,
@@ -221,6 +202,7 @@ impl Muxer {
frame_type: f_type,
data,
})
.await
}
pub fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
@@ -232,33 +214,25 @@ impl Muxer {
self.streams.remove(&stream_id);
}
pub fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
let stream_opt = self
.streams
.get(&stream_id)
.map(|s| (s.0.clone(), s.1.clone()));
pub async fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
let stream_opt = self
.streams
.get(&stream_id)
.map(|s| (s.0.clone(), s.1.clone()));
if let Some((tx, stats)) = stream_opt {
let size = data.len() as u64;
if let Some((tx, stats)) = stream_opt {
let size = data.len() as u64;
match tx.try_send(data) {
Ok(_) => {
stats.rx_bytes.fetch_add(size, Ordering::Relaxed);
}
Err(TrySendError::Full(_)) => {
netrunner_logger::warn!(
stream_id,
"Muxer -> Local: Buffer full, dropping packet. Check your bridge/TUN speed."
);
}
Err(TrySendError::Closed(_)) => {
self.remove_stream(stream_id);
match tx.send(data).await {
Ok(_) => {
stats.rx_bytes.fetch_add(size, Ordering::Relaxed);
}
Err(_) => {
self.remove_stream(stream_id);
}
}
}
}
}
pub fn record_leg_rx(&self, leg_id: u32, bytes: u64) {
if let Some(leg) = self.legs.get(&leg_id) {
@@ -271,16 +245,16 @@ pub fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
}
pub async fn perform_health_check(&self) {
let legs: Vec<(u32, Sender<MuxMessage>)> = self.legs.iter()
let legs: Vec<(u32, Sender<MuxMessage>)> = self
.legs
.iter()
.map(|k| (*k.key(), k.value().control_tx.clone()))
.collect();
for (leg_id, tx) in legs {
let probe_stream_id = self.id_gen.next();
let (probe_tx, mut probe_rx) = tokio::sync::mpsc::channel(2);
self.register_stream(probe_stream_id, probe_tx);
let msg = MuxMessage {
@@ -290,11 +264,10 @@ pub fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
};
let start = std::time::Instant::now();
if tx.try_send(msg).is_ok() {
match tokio::time::timeout(std::time::Duration::from_secs(2), probe_rx.recv()).await {
match tokio::time::timeout(std::time::Duration::from_secs(2), probe_rx.recv()).await
{
Ok(Some(_)) => {
let rtt = start.elapsed().as_millis() as u32;
if let Some(leg) = self.legs.get(&leg_id) {
@@ -303,7 +276,6 @@ pub fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
}
}
_ => {
if let Some(leg) = self.legs.get(&leg_id) {
leg.stats.rtt_ms.store(5000, Ordering::Relaxed);
warn!(leg_id, "❌ Leg Health Check Timeout (marked as slow)");
@@ -312,12 +284,10 @@ pub fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
}
}
self.remove_stream(probe_stream_id);
}
}
fn format_size(bytes: u64) -> String {
const KB: u64 = 1024;
const MB: u64 = KB * 1024;
@@ -333,7 +303,6 @@ pub fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
}
}
pub fn print_topology_tree(&self) {
println!(
"\n🌐 Netrunner Tunnel Topology [Session: {}]",