sessions and muxer legs

This commit is contained in:
2026-03-27 15:57:04 +07:00
parent fffc3794b9
commit f4c6f2efdd
13 changed files with 974 additions and 540 deletions
+28 -29
View File
@@ -1,8 +1,10 @@
use std::sync::Arc;
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};
use netrunner_logger::{debug, error, info};
use tokio::net::UdpSocket;
use tokio::sync::mpsc;
@@ -10,7 +12,7 @@ pub(crate) async fn run_tcp_bridge<R, W>(
stream_id: u32,
mut reader: R,
mut writer: W,
muxer: Muxer,
muxer: Arc<Muxer>,
mut v_rx: mpsc::Receiver<Bytes>,
) where
R: tokio::io::AsyncReadExt + Unpin,
@@ -71,61 +73,58 @@ pub(crate) async fn run_tcp_bridge<R, W>(
muxer.remove_stream(stream_id);
}
pub(crate) async fn run_udp_bridge(
stream_id: u32,
socket: UdpSocket,
muxer: Muxer,
muxer: Arc<Muxer>,
mut v_rx: mpsc::Receiver<Bytes>,
) {
let mut buf = BytesMut::with_capacity(NetworkConfig::global().udp_buffer_size);
let config = NetworkConfig::global();
let mut buf = vec![0u8; config.udp_buffer_size];
info!(
"🌉 [UDP {}] Bridge active. Media traffic leg optimization enabled.",
stream_id
);
loop {
tokio::select! {
// Данные ИЗ ИНТЕРНЕТА -> В ТУННЕЛЬ
res = socket.recv(&mut buf) => {
match res {
Ok(0) => break,
Ok(n) => {
let msg = MuxMessage {
stream_id,
frame_type: FrameType::UdpData,
data: Bytes::copy_from_slice(&buf[..n]),
};
if muxer.send_to_network(msg).await.is_err() { break; }
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).await {
error!("❌ [UDP {}] Failed to send to tunnel: {}", stream_id, e);
break;
}
}
Ok(_) => break,
Err(e) => {
error!(stream_id, error = %e, "UDP socket read error");
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, error = %e, "UDP socket write error");
error!("❌ [UDP {}] Internet write error: {}", stream_id, e);
break;
}
}
None => {
debug!(stream_id, "Virtual channel closed (UDP)");
break;
}
None => break,
}
}
}
}
let _ = muxer
.send_to_network(MuxMessage {
stream_id,
frame_type: FrameType::Close,
data: Bytes::new(),
})
.await;
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
info!("🔌 [UDP {}] Bridge closed", stream_id);
muxer.remove_stream(stream_id);
}
+240 -117
View File
@@ -1,14 +1,18 @@
use std::{net::Ipv4Addr, sync::Arc};
use crate::{
net::{
connection::{engine::TunnelEngine, handler::StreamHandler, muxer::Muxer},
network::NetworkConfig,
},
nrxp::{Codec, ErrorAction, FrameType},
rawcast::{LocalProtocol, RawCastEvent, RawCastFrame},
nrxp::{Codec, ErrorAction, Frame, FrameType},
rawcast::{LocalProtocol, RawCastAdapter, RawCastFrame},
tlseng::BrowserProfile,
};
use bytes::{Bytes, BytesMut};
use dashmap::DashMap;
use netrunner_logger::{error, info, warn};
use rand::Rng;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{
@@ -18,6 +22,38 @@ use tokio::{
sync::mpsc,
};
pub struct SessionManager {
sessions: DashMap<String, Arc<Muxer>>,
}
impl SessionManager {
pub fn new() -> Self {
Self {
sessions: DashMap::new(),
}
}
// Генерация переехала сюда
pub fn generate_id() -> String {
let mut rng = rand::rng();
format!("{:016x}{:016x}", rng.next_u64(), rng.next_u64())
}
pub fn get_or_create(&self, session_id: &str) -> Arc<Muxer> {
self.sessions
.entry(session_id.to_string())
.or_insert_with(|| Arc::new(Muxer::new(false)))
.clone()
}
// Удаление мертвых сессий, чтобы не текла память
pub fn remove(&self, session_id: &str) {
if self.sessions.remove(session_id).is_some() {
netrunner_logger::info!("🧹 Session {} completely closed and cleaned up", session_id);
}
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum ConnectionRole {
Client,
@@ -57,14 +93,19 @@ impl Connection {
}
}
type StreamContext = (Ipv4Addr, u16, LocalProtocol);
pub struct ClientHandler;
impl ClientHandler {
/// Вспомогательная функция: устанавливает одно физическое TLS-соединение
/// и возвращает готовый Muxer для работы с ним.
async fn establish_leg(remote_proxy_addr: &str, leg_name: &str) -> Result<Muxer, String> {
async fn establish_leg(
remote_proxy_addr: &str,
leg_id: u32,
muxer: Arc<Muxer>,
session_id: &str,
) -> Result<(), String> {
let leg_name = if leg_id == 0 { "TCP-Leg" } else { "UDP-Leg" };
info!(
"Establishing dedicated {} tunnel to {}...",
"Establishing dedicated {} to {}...",
leg_name, remote_proxy_addr
);
@@ -72,12 +113,11 @@ impl ClientHandler {
.await
.map_err(|e| format!("Failed to connect: {}", e))?;
if let Err(e) = stream.set_nodelay(true) {
warn!("Failed to set TCP_NODELAY on {} leg: {}", leg_name, e);
}
let _ = stream.set_nodelay(true);
let (inbound, outbound) = stream.into_split();
let mut conn = Connection::new_raw(inbound, outbound);
// TLS Handshake
let ch = conn
.codec
.make_client_handshake(&BrowserProfile::CHROME_131, "ubuntu.com")
@@ -98,19 +138,39 @@ impl ClientHandler {
.await
.map_err(|e| e.to_string())?;
if n == 0 {
return Err(format!("EOF during handshake on {} leg", leg_name));
return Err(format!("EOF on {}", leg_name));
}
}
Err(e) => return Err(format!("TLS error on {} leg: {:?}", leg_name, e)),
Err(e) => return Err(format!("TLS error on {}: {:?}", leg_name, e)),
}
}
info!("{} TLS Handshake complete. Starting Engine.", leg_name);
info!("{} TLS Handshake complete.", leg_name);
let handshake_payload = Bytes::from(format!("{}:{}", session_id, leg_id));
// 2. Шифруем через Codec! (Внутри сгенерится auth_tag)
let encrypted_handshake = conn
.codec
.encrypt_data(0, FrameType::Handshake, handshake_payload)
.map_err(|e| format!("Failed to encrypt Handshake: {:?}", e))?;
// 3. Отправляем в сокет
conn.outbound
.write_all(&encrypted_handshake)
.await
.map_err(|e| format!("Failed to send Handshake: {}", e))?;
info!(
"{} Sent Encrypted Handshake (Leg: {}). Starting engine...",
leg_name, leg_id
);
let (control_tx, control_rx) = mpsc::channel(NetworkConfig::global().muxer_capacity);
let (data_tx, data_rx) = mpsc::channel(NetworkConfig::global().muxer_capacity);
let muxer = Muxer::new(control_tx, data_tx, true);
// ОБНОВЛЕННЫЙ ВЫЗОВ: передаем leg_id
muxer.add_leg(leg_id, control_tx, data_tx);
let handler =
std::sync::Arc::new(StreamHandler::new(muxer.clone(), ConnectionRole::Client));
@@ -124,114 +184,123 @@ impl ClientHandler {
handler,
};
tokio::spawn(async move { engine.run().await });
Ok(muxer)
engine.run().await;
Err(format!("{} Engine stopped", leg_name))
}
/// Главная функция запуска клиента
pub async fn connect(
remote_proxy_addr: &str,
mut rx_from_engine: mpsc::Receiver<RawCastFrame>,
tx_to_engine: mpsc::Sender<RawCastFrame>,
) -> Result<(), String> {
// 1. Устанавливаем ДВА независимых соединения
// Мы можем сделать это параллельно через tokio::try_join! для скорости
let (tcp_muxer, udp_muxer) = tokio::try_join!(
Self::establish_leg(remote_proxy_addr, "TCP"),
Self::establish_leg(remote_proxy_addr, "UDP"),
)?;
let muxer = Arc::new(Muxer::new(true));
let registry: Arc<DashMap<u32, StreamContext>> = Arc::new(DashMap::new());
info!("Dual-tunnel architecture established successfully!");
let mut rng = rand::rng();
let session_id = format!("{:016x}{:016x}", rng.next_u64(), rng.next_u64());
info!("🔑 Generated Master Session ID: {}", session_id);
// 2. ЗАПУСКАЕМ МОСТ-РОУТЕР
for id in 0..2 {
let addr = remote_proxy_addr.to_string();
let m = muxer.clone();
let sid = session_id.clone();
tokio::spawn(async move {
loop {
if let Err(e) = Self::establish_leg(&addr, id, m.clone(), &sid).await {
error!("Leg {} failed: {}. Reconnecting...", id, e);
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
}
}
});
}
info!("🚀 Netrunner Dual-Leg Tunnel Active.");
let muxer_inner = muxer.clone();
tokio::spawn(async move {
while let Some(raw_frame) = rx_from_engine.recv().await {
let stream_id = raw_frame.socket_id as u32;
let is_udp = raw_frame.protocol == LocalProtocol::Udp;
// 1. Сохраняем метаданные для локального реестра (до того, как raw_frame исчезнет)
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;
// РОУТИНГ: Выбираем нужный физический канал
let muxer = if is_udp {
udp_muxer.clone()
} else {
tcp_muxer.clone()
// 2. Вся магия здесь: адаптер сам решит, это Connect или Data, и склеит IP:Port
let nrxp_frame = match RawCastAdapter::to_nrxp(raw_frame) {
Ok(frame) => frame,
Err(e) => {
error!("⚠️ [Adapter] Failed to cast frame: {}", e);
continue;
}
};
match raw_frame.event {
RawCastEvent::Connect => {
let stream_id = nrxp_frame.header.stream_id;
let f_type = nrxp_frame.header.frame_type;
let payload = nrxp_frame.payload;
// 3. Маршрутизация на основе ГОТОВОГО FrameType
match f_type {
FrameType::Connect | FrameType::UdpConnect => {
info!(
"🔗 [Muxer] Forwarding {} request to {} (stream {})",
if is_udp { "UDP" } else { "TCP" },
String::from_utf8_lossy(&payload), // Адаптер уже положил сюда "IP:Port"
stream_id
);
registry.insert(stream_id, (dst_ip, dst_port, protocol));
let (v_tx, mut v_rx) =
mpsc::channel(NetworkConfig::global().tcp_buffer_size);
muxer.register_stream(stream_id, v_tx);
muxer_inner.register_stream(stream_id, v_tx);
// Читаем домен, который мы прокинули в предыдущем шаге
let target = if !raw_frame.payload.is_empty() {
String::from_utf8_lossy(&raw_frame.payload).to_string()
} else {
format!("{}:{}", raw_frame.dst_ip, raw_frame.dst_port)
};
let frame_type = if is_udp {
FrameType::UdpConnect
} else {
FrameType::Connect
};
if let Err(e) = muxer
.send_control(stream_id, frame_type, Bytes::from(target))
.await
{
error!("Failed to send connect control frame: {}", e);
continue;
}
let tx_engine_clone = tx_to_engine.clone();
let mut muxer_clone = muxer.clone();
let dst_ip = raw_frame.dst_ip;
let dst_port = raw_frame.dst_port;
let protocol = raw_frame.protocol;
let socket_id = raw_frame.socket_id;
let tx_to_tun = tx_to_engine.clone();
let reg = registry.clone();
tokio::spawn(async move {
while let Some(payload) = v_rx.recv().await {
let data_frame = RawCastFrame::data(
protocol,
socket_id,
dst_ip,
dst_port,
payload.to_vec(),
);
if tx_engine_clone.send(data_frame).await.is_err() {
break;
while let Some(back_payload) = v_rx.recv().await {
if let Some(ctx) = reg.get(&stream_id) {
let (ip, port, proto) = *ctx;
let out_f_type = if proto == LocalProtocol::Udp {
FrameType::UdpData
} else {
FrameType::Data
};
let mock_nrxp = Frame::new(stream_id, out_f_type, back_payload);
// Здесь from_nrxp уже используется у тебя правильно!
if let Ok(raw) = RawCastAdapter::from_nrxp(
mock_nrxp,
ip,
port,
proto == LocalProtocol::Udp,
) {
let _ = tx_to_tun.send(raw).await;
}
}
}
muxer_clone.remove_stream(stream_id);
});
// Отправляем управляющий фрейм с адресом
let _ = muxer_inner.send_control(stream_id, f_type, payload).await;
}
RawCastEvent::Data => {
let frame_type = if is_udp {
FrameType::UdpData
} else {
FrameType::Data
};
if let Err(e) = muxer
.send_control(stream_id, frame_type, raw_frame.payload)
.await
{
error!("Failed to send data frame: {}", e);
}
FrameType::Data | FrameType::UdpData => {
let _ = muxer_inner.send_data_safe(stream_id, payload, is_udp).await;
}
RawCastEvent::Close => {
let _ = muxer
FrameType::Close => {
let _ = muxer_inner
.send_control(stream_id, FrameType::Close, Bytes::new())
.await;
muxer.remove_stream(stream_id);
muxer_inner.remove_stream(stream_id);
}
_ => {
warn!("Unhandled FrameType from adapter: {:?}", f_type);
}
}
}
info!("ClientHandler bridge task terminated.");
});
Ok(())
@@ -239,9 +308,17 @@ impl ClientHandler {
}
pub struct ServerHandler {
pub(crate) conn: Connection,
pub(crate) session_manager: Arc<SessionManager>,
}
impl ServerHandler {
pub fn new(connection: Connection) -> Self {
Self {
conn: connection,
session_manager: Arc::new(SessionManager::new()),
}
}
async fn handle_stealth_fallback(
mut client_inbound: OwnedReadHalf,
mut client_outbound: OwnedWriteHalf,
@@ -288,20 +365,17 @@ impl ServerHandler {
}
}
}
use crate::parser::Parser; // Не забудь импортировать трейт Parser!
#[async_trait::async_trait]
impl TunnelHandler for ServerHandler {
async fn run(mut self) -> Result<(), String> {
info!("Acting as TLS Server with Stealth Fallback");
let (control_tx, control_rx) = mpsc::channel(NetworkConfig::global().muxer_capacity);
let (data_tx, data_rx) = mpsc::channel(NetworkConfig::global().muxer_capacity);
let muxer = Muxer::new(control_tx, data_tx, false);
// 1. TLS Хендшейк (без изменений)
let handshake_timeout = std::time::Duration::from_secs(1);
let hello = loop {
let buf_snapshot = self.conn.read_buf.clone().freeze();
match self
.conn
.codec
@@ -309,19 +383,15 @@ impl TunnelHandler for ServerHandler {
{
Ok(b) => break b,
Err(e) if e.action == ErrorAction::Wait => {
let read_res = tokio::time::timeout(
let res = tokio::time::timeout(
handshake_timeout,
self.conn.inbound.read_buf(&mut self.conn.read_buf),
)
.await;
match read_res {
Ok(Ok(n)) if n == 0 => return Err("Client closed".into()),
match res {
Ok(Ok(0)) => return Err("Client closed".into()),
Ok(Ok(_)) => continue,
Ok(Err(e)) => return Err(e.to_string()),
Err(_) => {
warn!("Handshake timeout. Going stealth.");
_ => {
ServerHandler::handle_stealth_fallback(
self.conn.inbound,
self.conn.outbound,
@@ -332,14 +402,7 @@ impl TunnelHandler for ServerHandler {
}
}
}
Err(e) => {
warn!("Auth/Format failed: {:?}. Going stealth.", e);
info!(
"DEBUG: Restoring {} bytes from snapshot for fallback",
buf_snapshot.len()
);
Err(_) => {
ServerHandler::handle_stealth_fallback(
self.conn.inbound,
self.conn.outbound,
@@ -357,9 +420,61 @@ impl TunnelHandler for ServerHandler {
.await
.map_err(|e| e.to_string())?;
let handler = std::sync::Arc::new(StreamHandler::new(muxer, ConnectionRole::Server));
// 1. ЧИТАЕМ ЗАШИФРОВАННЫЙ HANDSHAKE ЧЕРЕЗ CODEC
let (session_id, leg_id) = loop {
let n = tokio::time::timeout(
std::time::Duration::from_secs(3),
self.conn.inbound.read_buf(&mut self.conn.read_buf),
)
.await
.map_err(|_| "Timeout waiting for Handshake".to_string())?
.map_err(|e| e.to_string())?;
TunnelEngine {
if n == 0 {
return Err("Client closed connection before Handshake".into());
}
// Кодек сам расшифрует и проверит auth_tag
match self.conn.codec.inbound(&mut self.conn.read_buf) {
Ok(Some(frame)) => {
if frame.header.frame_type == FrameType::Handshake {
let payload_str = String::from_utf8_lossy(&frame.payload).to_string();
let parts: Vec<&str> = payload_str.split(':').collect();
if parts.len() == 2 {
let sid = parts[0].to_string();
let lid: u32 = parts[1].parse().unwrap_or(0);
info!(
"🤝 Secure Handshake verified! Session: {}, Leg: {}",
sid, lid
);
break (sid, lid);
} else {
return Err("Invalid Handshake format".into());
}
} else {
return Err("Invalid first frame (Expected Handshake)".into());
}
}
Ok(None) => continue, // Ждем еще байт
Err(e) => return Err(format!("Decryption/Auth failed during Handshake: {:?}", e)),
}
};
// 2. РЕГИСТРИРУЕМ НОГУ
let muxer = self.session_manager.get_or_create(&session_id);
let (control_tx, control_rx) = mpsc::channel(NetworkConfig::global().muxer_capacity);
let (data_tx, data_rx) = mpsc::channel(NetworkConfig::global().muxer_capacity);
// ОБНОВЛЕННЫЙ ВЫЗОВ: регистрируем конкретный ID ноги
muxer.add_leg(leg_id, control_tx, data_tx);
let handler =
std::sync::Arc::new(StreamHandler::new(muxer.clone(), ConnectionRole::Server));
// 3. ЗАПУСКАЕМ ДВИЖОК
let engine = TunnelEngine {
inbound: self.conn.inbound,
outbound: self.conn.outbound,
codec: self.conn.codec,
@@ -367,8 +482,16 @@ impl TunnelHandler for ServerHandler {
control_rx,
data_rx,
handler,
};
let _ = engine.run().await;
// 4. ОЧИСТКА ПОСЛЕ ОТКЛЮЧЕНИЯ
muxer.remove_leg(leg_id);
if muxer.active_legs_count() == 0 {
self.session_manager.remove(&session_id);
}
.run()
.await
Ok(())
}
}
+76 -35
View File
@@ -1,5 +1,7 @@
use std::sync::Arc;
use bytes::Bytes;
use netrunner_logger::{debug, error, info};
use netrunner_logger::{debug, error, info, warn};
use crate::{
net::{
@@ -14,18 +16,23 @@ use crate::{
};
pub(crate) struct StreamHandler {
muxer: Muxer,
muxer: Arc<Muxer>,
role: ConnectionRole,
}
impl StreamHandler {
pub(crate) fn new(muxer: Muxer, role: ConnectionRole) -> Self {
pub(crate) fn new(muxer: Arc<Muxer>, role: ConnectionRole) -> Self {
Self { muxer, role }
}
pub(crate) async fn handle(&self, frame: Frame) {
let stream_id = frame.header.stream_id;
info!(
stream_id,
"📥 [Tunnel] Received frame: {:?}", frame.header.frame_type
);
match frame.header.frame_type {
FrameType::Connect => self.on_connect(stream_id, frame.payload).await,
FrameType::UdpConnect => self.on_udp_connect(stream_id, frame.payload).await,
@@ -35,10 +42,11 @@ impl StreamHandler {
_ => debug!(stream_id, "Unhandled frame type"),
}
}
async fn on_connect(&self, stream_id: u32, payload: Bytes) {
let target_str = String::from_utf8_lossy(&payload).to_string();
if self.role == ConnectionRole::Server {
let target_str = String::from_utf8_lossy(&payload).to_string();
info!(stream_id, target = %target_str, "🌐 [TCP] Request to establish remote connection");
let muxer = self.muxer.clone();
let (v_tx, v_rx) =
@@ -47,41 +55,46 @@ impl StreamHandler {
tokio::spawn(async move {
let start = std::time::Instant::now();
info!(stream_id, target = %target_str, "Attempting remote TCP connection");
info!(
stream_id,
"⏳ [TCP Worker] Attempting connection to {}", target_str
);
let connect_timeout = tokio::time::timeout(
std::time::Duration::from_secs(5),
let connect_res = tokio::time::timeout(
std::time::Duration::from_secs(7),
tokio::net::TcpStream::connect(&target_str),
)
.await;
match connect_timeout {
match connect_res {
Ok(Ok(stream)) => {
let elapsed = start.elapsed();
info!(stream_id, target = %target_str, latency_ms = elapsed.as_millis(), "Remote TCP connection established");
// Больше никаких ответов (Reply), просто прокидываем байты!
info!(stream_id, target = %target_str, "✅ [TCP] Established in {:?}. Starting bridge.", elapsed);
let (r, w) = stream.into_split();
run_tcp_bridge(stream_id, r, w, muxer, v_rx).await;
info!(stream_id, "🔚 [TCP Worker] Bridge task finished");
}
Ok(Err(e)) => {
error!(stream_id, target = %target_str, error = %e, "TCP connection failed");
error!(stream_id, target = %target_str, error = %e, "❌ [TCP] Connection failed");
Self::close_stream(&muxer, stream_id).await;
}
Err(_) => {
error!(stream_id, target = %target_str, "Connection timed out (DNS/TCP)");
error!(stream_id, target = %target_str, "⏰ [TCP] Connection timeout (DNS or IP unreachable)");
Self::close_stream(&muxer, stream_id).await;
}
}
});
} else {
info!(stream_id, "📲 [TCP] Dispatching payload to local stack");
self.muxer.dispatch_to_local(stream_id, payload).await;
}
}
async fn on_udp_connect(&self, stream_id: u32, payload: Bytes) {
let target_str = String::from_utf8_lossy(&payload).to_string();
if self.role == ConnectionRole::Server {
let target_str = String::from_utf8_lossy(&payload).to_string();
info!(stream_id, target = %target_str, "🚀 [UDP] Request to establish remote bridge");
let muxer = self.muxer.clone();
let (v_tx, v_rx) =
@@ -89,47 +102,75 @@ impl StreamHandler {
muxer.register_stream(stream_id, v_tx);
tokio::spawn(async move {
info!(stream_id, target = %target_str, "Attempting remote UDP connection");
info!(
stream_id,
"⏳ [UDP Worker] Binding socket for {}", target_str
);
match tokio::net::UdpSocket::bind("0.0.0.0:0").await {
Ok(socket) => {
if let Err(e) = socket.connect(&target_str).await {
error!(stream_id, target = %target_str, error = %e, "UDP connect failed");
Self::close_stream(&muxer, stream_id).await;
return;
}
let socket_res = tokio::net::UdpSocket::bind("[::]:0").await;
let socket = match socket_res {
Ok(s) => s,
Err(e) => {
warn!(
stream_id,
"⚠️ [UDP] IPv6 bind failed, falling back to IPv4: {}", e
);
tokio::net::UdpSocket::bind("0.0.0.0:0")
.await
.expect("UDP bind fail")
}
};
// Успех - просто начинаем слушать UDP и слать в туннель
match socket.connect(&target_str).await {
Ok(_) => {
let local = socket.local_addr().unwrap();
let remote = socket.peer_addr().unwrap();
info!(
"✅ [UDP {}] Socket ready. Local: {}, Remote: {}. Starting bridge.",
stream_id, local, remote
);
run_udp_bridge(stream_id, socket, muxer, v_rx).await;
info!(stream_id, "🔚 [UDP Worker] Bridge task finished");
}
Err(e) => {
error!(stream_id, target = %target_str, error = %e, "UDP bind failed");
error!(stream_id, target = %target_str, error = %e, "❌ [UDP] Target connect failed");
Self::close_stream(&muxer, stream_id).await;
}
}
});
} else {
info!(
stream_id,
"📲 [UDP] Dispatching connection payload to local stack"
);
self.muxer.dispatch_to_local(stream_id, payload).await;
}
}
async fn on_udp_data(&self, stream_id: u32, payload: Bytes) {
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;
}
// Вспомогательная функция вместо send_error_reply
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;
}
async fn on_close(&self, stream_id: u32) {
info!(stream_id, "🏁 [Close] Received close signal for stream");
self.muxer.remove_stream(stream_id);
}
async fn close_stream(muxer: &Muxer, stream_id: u32) {
info!(
stream_id,
"📡 [Control] Sending CLOSE signal to remote peer"
);
let _ = muxer
.send_control(stream_id, FrameType::Close, Bytes::new())
.await;
muxer.remove_stream(stream_id);
}
async fn on_data(&self, stream_id: u32, payload: Bytes) {
self.muxer.dispatch_to_local(stream_id, payload).await;
}
async fn on_close(&self, stream_id: u32) {
self.muxer.remove_stream(stream_id);
}
}
+113 -57
View File
@@ -2,10 +2,18 @@ use bytes::Bytes;
use dashmap::DashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc::{error::SendError, Sender};
use tokio::sync::mpsc::Sender;
use crate::net::network::NetworkConfig;
use crate::nrxp::FrameType;
// Вспомогательная структура для отдельного TLS-соединения
#[derive(Clone)]
struct MuxLeg {
control_tx: Sender<MuxMessage>,
data_tx: Sender<MuxMessage>,
}
struct IdGenerator {
counter: AtomicU32,
}
@@ -31,67 +39,129 @@ pub struct MuxMessage {
#[derive(Clone)]
pub struct Muxer {
control_tx: Sender<MuxMessage>,
data_tx: Sender<MuxMessage>,
legs: Arc<DashMap<u32, MuxLeg>>,
streams: Arc<DashMap<u32, Sender<Bytes>>>,
id_gen: Arc<IdGenerator>,
leg_selector: Arc<AtomicU32>,
}
impl Muxer {
pub fn new(
control_tx: Sender<MuxMessage>,
data_tx: Sender<MuxMessage>,
is_client: bool,
) -> Self {
pub fn new(is_client: bool) -> Self {
Self {
control_tx,
data_tx,
legs: Arc::new(DashMap::new()),
streams: Arc::new(DashMap::new()),
id_gen: Arc::new(IdGenerator::new(is_client)),
leg_selector: Arc::new(AtomicU32::new(0)),
// leg_id_gen удален!
}
}
pub fn next_id(&self) -> u32 {
self.id_gen.next()
// ТЕПЕРЬ leg_id ПЕРЕДАЕТСЯ СНАРУЖИ
pub fn add_leg(
&self,
leg_id: u32,
control_tx: Sender<MuxMessage>,
data_tx: Sender<MuxMessage>,
) {
self.legs.insert(
leg_id,
MuxLeg {
control_tx,
data_tx,
},
);
netrunner_logger::info!(leg_id, "MUXER: Leg registered");
}
pub async fn send_to_network(&self, message: MuxMessage) -> Result<(), SendError<MuxMessage>> {
self.data_tx.send(message).await
pub fn remove_leg(&self, leg_id: u32) {
self.legs.remove(&leg_id);
netrunner_logger::info!(leg_id, "MUXER: Leg removed");
}
pub async fn send_data_safe(&self, stream_id: u32, mut data: Bytes) -> Result<(), String> {
// Лимит полезной нагрузки, чтобы вместе с заголовком и паддингом
// пакет оставался в пределах ~1400-1450 байт.
pub fn active_legs_count(&self) -> usize {
self.legs.len()
}
// Возвращаем кортеж (leg_id, MuxLeg), чтобы знать, кого удалять в случае ошибки
fn select_leg(&self, frame_type: &FrameType) -> Option<(u32, MuxLeg)> {
if self.legs.is_empty() {
return None;
}
match frame_type {
FrameType::UdpData | FrameType::UdpConnect => {
// Пытаемся взять UDP ногу (1), если её нет — берем TCP (0) как фоллбэк
if let Some(leg) = self.legs.get(&1) {
Some((1, leg.clone()))
} else if let Some(leg) = self.legs.get(&0) {
Some((0, leg.clone()))
} else {
None
}
}
_ => {
// Пытаемся взять TCP ногу (0), фоллбэк на UDP (1)
if let Some(leg) = self.legs.get(&0) {
Some((0, leg.clone()))
} else if let Some(leg) = self.legs.get(&1) {
Some((1, leg.clone()))
} else {
None
}
}
}
}
pub async fn send_to_network(&self, message: MuxMessage) -> Result<(), String> {
let (leg_id, leg) = self
.select_leg(&message.frame_type)
.ok_or_else(|| "MUXER: No active legs available".to_string())?;
let target_tx = match message.frame_type {
FrameType::Connect
| FrameType::Close
| FrameType::UdpConnect
| FrameType::Handshake => &leg.control_tx,
_ => &leg.data_tx,
};
// Если канал мертв (TunnelEngine упал), удаляем ногу из роутинга!
if let Err(e) = target_tx.send(message).await {
self.remove_leg(leg_id);
return Err(format!("MUXER: Leg {} died during send: {}", leg_id, e));
}
Ok(())
}
pub async fn send_data_safe(
&self,
stream_id: u32,
mut data: Bytes,
is_udp: bool,
) -> Result<(), String> {
const MAX_PAYLOAD_CHUNK: usize = 1300;
let frame_type = if is_udp {
FrameType::UdpData
} else {
FrameType::Data
};
if data.len() <= MAX_PAYLOAD_CHUNK {
return self
.send_to_network(MuxMessage {
stream_id,
frame_type: FrameType::Data,
data,
})
.await
.map_err(|e| e.to_string());
}
// Если данных много (например, те самые 4096 байт), режем их на куски
while !data.is_empty() {
let chunk_size = std::cmp::min(data.len(), MAX_PAYLOAD_CHUNK);
let chunk = data.split_to(chunk_size);
self.send_to_network(MuxMessage {
stream_id,
frame_type: FrameType::Data,
frame_type,
data: chunk,
})
.await
.map_err(|e| e.to_string())?;
.await?;
// Небольшая уступка планировщику, чтобы не забить канал мгновенно
tokio::task::yield_now().await;
if !data.is_empty() {
tokio::task::yield_now().await;
}
}
Ok(())
}
@@ -101,14 +171,12 @@ impl Muxer {
f_type: FrameType,
data: Bytes,
) -> Result<(), String> {
self.control_tx
.send(MuxMessage {
stream_id,
frame_type: f_type,
data,
})
.await
.map_err(|e| e.to_string())
self.send_to_network(MuxMessage {
stream_id,
frame_type: f_type,
data,
})
.await
}
pub fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
@@ -120,22 +188,10 @@ impl Muxer {
}
pub async fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
let tx = self.streams.get(&stream_id).map(|r| r.value().clone());
if let Some(tx) = tx {
if let Err(_e) = tx.send(data).await {
netrunner_logger::warn!(
stream_id,
"MUXER: [INBOUND_ERR] Local channel closed, dropping packet"
);
if let Some(tx) = self.streams.get(&stream_id) {
if let Err(_) = tx.send(data).await {
self.remove_stream(stream_id);
}
} else {
netrunner_logger::trace!(
stream_id,
len = data.len(),
"MUXER: [IGNORE] Received data for already closed stream (draining pipe)"
);
}
}
}
+4 -15
View File
@@ -12,22 +12,11 @@ pub struct Network {
host: String,
port: u16,
role: ConnectionRole,
remote_proxy_addr: Option<String>,
}
impl Network {
pub fn new(
host: String,
port: u16,
role: ConnectionRole,
remote_proxy_addr: Option<String>,
) -> Self {
Self {
host,
port,
role,
remote_proxy_addr,
}
pub fn new(host: String, port: u16, role: ConnectionRole) -> Self {
Self { host, port, role }
}
pub async fn run(&self, token: CancellationToken) {
@@ -60,7 +49,7 @@ impl Network {
// Создаем соединение (init = true для сервера)
let conn = Connection::new(stream, true);
let handler = ServerHandler { conn };
let handler = ServerHandler::new(conn);
tokio::spawn(async move {
if let Err(e) = handler.run().await {
@@ -108,7 +97,7 @@ impl NetworkConfig {
// 2. Каналы Muxer (Баланс между скоростью и задержкой)
let muxer_capacity = 512; // Глобальная очередь (выдержит много вкладок)
let tcp_stream_capacity = 16; // Хватит для скорости, но не даст пингу взлететь
let udp_stream_capacity = 64; // Простор для голосового трафика и игр
let udp_stream_capacity = 32; // Простор для голосового трафика и игр
Self {
mtu: system_mtu,