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

fixes and upgrades. async remove, buf changes

Browse Source
This commit is contained in:
Kirill
2026-04-03 11:43:03 +07:00
parent 337f191cca
commit b90083539c
9 changed files with 225 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/src/net/connection/bridge.rs
+127 -69
View File
@@ -1,12 +1,31 @@
use std::sync::Arc;
use std::time::Duration;
use crate::net::connection::muxer::{MuxMessage, Muxer};
use crate::net::network::NetworkConfig;
use crate::nrxp::FrameType;
use bytes::{Bytes, BytesMut};
use netrunner_logger::{debug, error, info};
use netrunner_logger::{debug, error, info, warn};
use tokio::net::UdpSocket;
use tokio::sync::mpsc;
use tokio::time::timeout;
/// Тайм-аут бездействия. Если данных нет 5 минут — закрываем мост.
const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
/// RAII Guard: гарантирует удаление стрима из Muxer при любом выходе из функции,
/// даже если произошла паника (panic) или ранняя ошибка.
struct StreamGuard {
stream_id: u32,
muxer: Arc<Muxer>,
}
impl Drop for StreamGuard {
fn drop(&mut self) {
debug!(self.stream_id, "StreamGuard: Cleaning up resources");
self.muxer.remove_stream(self.stream_id);
}
}
pub(crate) async fn run_tcp_bridge<R, W>(
stream_id: u32,
@@ -18,111 +37,150 @@ pub(crate) async fn run_tcp_bridge<R, W>(
R: tokio::io::AsyncReadExt + Unpin,
W: tokio::io::AsyncWriteExt + Unpin,
{
// Инициализируем защитника ресурсов
let _guard = StreamGuard {
stream_id,
muxer: muxer.clone(),
};
let mut buf = BytesMut::with_capacity(NetworkConfig::global().tcp_buffer_size);
loop {
tokio::select! {
res = reader.read_buf(&mut buf) => {
match res {
Ok(0) => {
debug!(stream_id, "Socket closed (EOF)");
break;
// Оборачиваем весь select в таймаут, чтобы убивать "мертвые" задачи
let select_res = timeout(IDLE_TIMEOUT, async {
tokio::select! {
// Данные из реального TCP сокета -> в Muxer
res = reader.read_buf(&mut buf) => {
match res {
Ok(0) => {
debug!(stream_id, "TCP Socket reached EOF");
return Ok(false); // Выход из цикла
}
Ok(_) => {
let msg = MuxMessage {
stream_id,
frame_type: FrameType::Data,
data: buf.split().freeze(),
};
// Если туннель закрыт — выходим
if muxer.send_to_network(msg).is_err() {
return Ok(false);
}
Ok(true)
}
Err(e) => {
error!(stream_id, error = %e, "TCP Socket read error");
Err(e.to_string())
}
}
Ok(_) => {
let msg = MuxMessage {
stream_id,
frame_type: FrameType::Data,
data: buf.split().freeze(),
};
if muxer.send_to_network(msg).is_err() { break; }
}
Err(e) => {
error!(stream_id, error = %e, "Socket read error");
break;
}
// Данные из туннеля -> в реальный TCP сокет
maybe_data = v_rx.recv() => {
match maybe_data {
Some(data) => {
if data.is_empty() { return Ok(true); }
if let Err(e) = writer.write_all(&data).await {
error!(stream_id, error = %e, "TCP Socket write error");
return Err(e.to_string());
}
Ok(true)
}
None => {
debug!(stream_id, "Virtual channel closed (Muxer removed stream)");
Ok(false)
}
}
}
}
})
.await;
maybe_data = v_rx.recv() => {
match maybe_data {
Some(data) => {
if data.is_empty() { continue; }
if let Err(e) = writer.write_all(&data).await {
error!(stream_id, error = %e, "Socket write error");
break;
}
}
None => {
debug!(stream_id, "Virtual channel closed");
break;
}
}
match select_res {
Ok(Ok(true)) => continue,
Ok(Ok(false)) => break, // Штатное закрытие
Ok(Err(e)) => {
debug!(stream_id, "Bridge closing due to error: {}", e);
break;
}
Err(_) => {
warn!(stream_id, "TCP Bridge IDLE timeout reached. Evicting task.");
break;
}
}
}
// Попытка вежливо сказать удаленной стороне, что мы закрылись
let _ = muxer.send_to_network(MuxMessage {
stream_id,
frame_type: FrameType::Close,
data: Bytes::new(),
});
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
muxer.remove_stream(stream_id);
// Небольшая пауза, чтобы пакет Close успел уйти в сеть перед тем, как Guard удалит стрим
tokio::time::sleep(Duration::from_millis(50)).await;
}
pub(crate) async fn run_udp_bridge(
stream_id: u32,
socket: UdpSocket,
muxer: Arc<Muxer>,
mut v_rx: mpsc::Receiver<Bytes>,
) {
let _guard = StreamGuard {
stream_id,
muxer: muxer.clone(),
};
let config = NetworkConfig::global();
let mut buf = vec![0u8; config.udp_buffer_size];
info!(
"🌉 [UDP {}] Bridge active. Media traffic leg optimization enabled.",
stream_id
);
info!(stream_id, "🌉 UDP Bridge active");
loop {
tokio::select! {
// Данные ИЗ ИНТЕРНЕТА -> В ТУННЕЛЬ
res = socket.recv(&mut buf) => {
match res {
Ok(n) if n > 0 => {
let data = Bytes::copy_from_slice(&buf[..n]);
// КРИТИЧЕСКИЙ МОМЕНТ: передаем true (is_udp),
// чтобы Muxer выбрал правильную ногу для реалтайм трафика
if let Err(e) = muxer.send_data_safe(stream_id, data, true) {
error!("❌ [UDP {}] Failed to send to tunnel: {}", stream_id, e);
break;
let select_res = timeout(IDLE_TIMEOUT, async {
tokio::select! {
// Из интернета в туннель
res = socket.recv(&mut buf) => {
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) {
error!(stream_id, "UDP Failed to send to tunnel: {}", e);
return Err(e);
}
Ok(true)
}
Ok(_) => Ok(false),
Err(e) => {
error!(stream_id, "UDP Socket read error: {}", e);
Err(e.to_string())
}
}
Ok(_) => break,
Err(e) => {
error!("❌ [UDP {}] Internet read error: {}", stream_id, e);
break;
}
// Из туннеля в интернет
maybe_data = v_rx.recv() => {
match maybe_data {
Some(data) => {
if let Err(e) = socket.send(&data).await {
error!(stream_id, "UDP Internet write error: {}", e);
return Err(e.to_string());
}
Ok(true)
}
None => Ok(false),
}
}
}
})
.await;
// Данные ИЗ ТУННЕЛЯ -> В ИНТЕРНЕТ
maybe_data = v_rx.recv() => {
match maybe_data {
Some(data) => {
if let Err(e) = socket.send(&data).await {
error!("❌ [UDP {}] Internet write error: {}", stream_id, e);
break;
}
}
None => break,
}
}
match select_res {
Ok(Ok(true)) => continue,
_ => break, // Любая ошибка или таймаут закрывают UDP мост
}
}
info!("🔌 [UDP {}] Bridge closed", stream_id);
muxer.remove_stream(stream_id);
debug!(stream_id, "🔌 UDP Bridge closed");
}
core/src/net/connection/connection.rs
+12 -15
View File
@@ -1,4 +1,4 @@
use std::{net::Ipv4Addr, sync::Arc};
use std::{net::Ipv4Addr, sync::Arc, time::Duration};
use crate::{
net::{
@@ -225,14 +225,14 @@ impl ClientHandler {
}
// === ДЕМОН СТАТИСТИКИ (КЛИЕНТ) ===
let m_stats = muxer.clone();
let m_weak = Arc::downgrade(&muxer);
tokio::spawn(async move {
loop {
tokio::time::sleep(std::time::Duration::from_secs(15)).await;
if m_stats.active_legs_count() > 0 {
m_stats.perform_health_check().await;
m_stats.print_topology_tree();
while let Some(m_stats) = m_weak.upgrade() {
tokio::time::sleep(Duration::from_secs(15)).await;
if m_stats.active_legs_count() == 0 {
break;
}
m_stats.print_topology_tree();
}
});
@@ -483,17 +483,14 @@ impl TunnelHandler for ServerHandler {
let handler =
std::sync::Arc::new(StreamHandler::new(muxer.clone(), ConnectionRole::Server));
// === ДЕМОН СТАТИСТИКИ (СЕРВЕР) ===
let m_stats = muxer.clone();
let m_weak = Arc::downgrade(&muxer);
tokio::spawn(async move {
loop {
tokio::time::sleep(std::time::Duration::from_secs(15)).await;
if m_stats.active_legs_count() > 0 {
m_stats.perform_health_check().await;
m_stats.print_topology_tree();
} else {
while let Some(m_stats) = m_weak.upgrade() {
tokio::time::sleep(Duration::from_secs(15)).await;
if m_stats.active_legs_count() == 0 {
break;
}
m_stats.print_topology_tree();
}
});
core/src/net/connection/handler.rs
+4 -4
View File
@@ -87,7 +87,7 @@ impl StreamHandler {
});
} else {
info!(stream_id, "📲 [TCP] Dispatching payload to local stack");
self.muxer.dispatch_to_local(stream_id, payload).await;
self.muxer.dispatch_to_local(stream_id, payload);
}
}
@@ -144,19 +144,19 @@ impl StreamHandler {
stream_id,
"📲 [UDP] Dispatching connection payload to local stack"
);
self.muxer.dispatch_to_local(stream_id, payload).await;
self.muxer.dispatch_to_local(stream_id, payload);
}
}
async fn on_data(&self, stream_id: u32, payload: Bytes) {
// Здесь info может быть избыточным при большой нагрузке, но для отладки полезно
debug!(stream_id, "📦 [TCP Data] Size: {} bytes", payload.len());
self.muxer.dispatch_to_local(stream_id, payload).await;
self.muxer.dispatch_to_local(stream_id, payload);
}
async fn on_udp_data(&self, stream_id: u32, payload: Bytes) {
debug!(stream_id, "📦 [UDP Data] Size: {} bytes", payload.len());
self.muxer.dispatch_to_local(stream_id, payload).await;
self.muxer.dispatch_to_local(stream_id, payload);
}
async fn on_close(&self, stream_id: u32) {
core/src/net/connection/muxer.rs
+73 -47
View File
@@ -1,6 +1,7 @@
use bytes::Bytes;
use dashmap::DashMap;
use netrunner_logger::{info, warn};
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::Sender;
@@ -13,7 +14,7 @@ use crate::nrxp::FrameType;
pub struct LegStats {
pub tx_bytes: AtomicU64,
pub rx_bytes: AtomicU64,
pub rtt_ms: AtomicU32, // Задержка (пинг) конкретно этой ноги
pub rtt_ms: AtomicU32,
}
#[derive(Default, Debug)]
@@ -107,35 +108,42 @@ impl Muxer {
self.legs.len()
}
// ДИНАМИЧЕСКИЙ ВЫБОР НОГИ ПО STREAM_ID (Session Affinity)
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 mut active_legs: Vec<u32> = self.legs.iter().map(|kv| *kv.key()).collect();
active_legs.sort_unstable(); // Сортируем для предсказуемой привязки
let is_udp = matches!(frame_type, FrameType::UdpData | FrameType::UdpConnect);
// Фильтруем: UDP идет по нечетным ногам, TCP по четным (если они есть)
let preferred_legs: Vec<u32> = active_legs
// Собираем доступные ноги
let mut candidates: Vec<(u32, MuxLeg)> = self.legs
.iter()
.copied()
.filter(|id| if is_udp { id % 2 != 0 } else { id % 2 == 0 })
.map(|kv| (*kv.key(), kv.value().clone()))
.filter(|(id, _)| {
// Фильтр по протоколу, если ноги разделены
if is_udp { id % 2 != 0 } else { id % 2 == 0 }
})
.collect();
// БАЛАНСИРОВКА: Привязываем stream_id к конкретной ноге.
// Все пакеты одного стрима пойдут строго по одному маршруту.
let target_id = if !preferred_legs.is_empty() {
preferred_legs[(stream_id as usize) % preferred_legs.len()]
} else {
active_legs[(stream_id as usize) % active_legs.len()]
};
// Если по протоколу ничего не нашли, берем любые активные
if candidates.is_empty() {
candidates = self.legs.iter().map(|kv| (*kv.key(), kv.value().clone())).collect();
}
self.legs
.get(&target_id)
.map(|leg| (target_id, leg.clone()))
// --- МАГИЯ ЗДЕСЬ ---
// Сортируем кандидатов по RTT (от меньшего к большему)
// Ноги с RTT = 0 (еще не проверенные) или таймаутом считаем медленными
candidates.sort_by_key(|(_, leg)| {
let rtt = leg.stats.rtt_ms.load(Ordering::Relaxed);
if rtt == 0 { 9999 } else { rtt }
});
// Выбираем из ТОП-2 лучших ног по RTT, используя stream_id для аффинити
// Это защитит от ситуации, когда все стримы прыгнут на одну ногу с RTT 10мс
let pool_size = std::cmp::min(candidates.len(), 2);
let target = &candidates[stream_id as usize % pool_size];
Some(target.clone())
}
// СИНХРОННАЯ ФУНКЦИЯ (с отбрасыванием пакетов при перегрузке)
@@ -232,23 +240,33 @@ impl Muxer {
self.streams.remove(&stream_id);
}
// ЛОКАЛЬНАЯ ДИСПЕТЧЕРИЗАЦИЯ (Остается async, так как тут мы используем wait для TUN)
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()));
pub fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
// Получаем клон Sender и Stats из DashMap
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;
if tx.send(data).await.is_err() {
self.remove_stream(stream_id);
} else {
// Используем try_send вместо send().await
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);
}
}
}
}
pub fn record_leg_rx(&self, leg_id: u32, bytes: u64) {
if let Some(leg) = self.legs.get(&leg_id) {
@@ -256,41 +274,49 @@ impl Muxer {
}
}
// УМНЫЙ HEALTH CHECK: Пингуем каждую ногу отдельно для честного RTT
pub async fn perform_health_check(&self) {
let legs: Vec<(u32, Sender<MuxMessage>)> = self
.legs
.iter()
// Создаем снапшот ног, чтобы не держать lock DashMap слишком долго
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);
// Регистрируем временный стрим для ответа на PING
self.register_stream(probe_stream_id, probe_tx);
let start = std::time::Instant::now();
let msg = MuxMessage {
stream_id: probe_stream_id,
frame_type: FrameType::Handshake,
data: Bytes::from("PING"),
};
// Отправляем PING прямо в конкретную ногу в обход балансировщика
let start = std::time::Instant::now();
// Отправляем напрямую в ногу
if tx.try_send(msg).is_ok() {
if let Ok(Some(_)) =
tokio::time::timeout(std::time::Duration::from_secs(2), probe_rx.recv()).await
{
let rtt = start.elapsed().as_millis() as u32;
// Сохраняем RTT в стату ноги
if let Some(leg) = self.legs.get(&leg_id) {
leg.stats.rtt_ms.store(rtt, Ordering::Relaxed);
// Ждем ответа с жестким таймаутом
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) {
leg.stats.rtt_ms.store(rtt, Ordering::Relaxed);
debug!(leg_id, rtt, "✅ Leg Health Check OK");
}
}
_ => {
// Если таймаут — задираем RTT до небес, чтобы select_leg ее не выбирал
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)");
}
}
} else {
warn!(leg_id, "❌ Health check timed out for Leg");
}
}
// Обязательная очистка, чтобы не было утечки памяти в DashMap streams
self.remove_stream(probe_stream_id);
}
}