core engine rewrite and client side changes

This commit is contained in:
2026-03-21 22:24:21 +07:00
parent af5ab22712
commit 46e8b4ae73
13 changed files with 573 additions and 376 deletions
+3 -3
View File
@@ -1,6 +1,6 @@
use crate::protocol::codec::frame::FrameType;
use crate::proxy::connection::muxer::{MuxMessage, Muxer};
use crate::proxy::connection::BUF_SIZE;
use crate::proxy::connection::{TCP_BUF_SIZE, UDP_BUF_SIZE};
use bytes::{Bytes, BytesMut};
use netrunner_logger::{debug, error};
use tokio::net::UdpSocket;
@@ -16,7 +16,7 @@ pub async fn run_proxy_bridge<R, W>(
R: tokio::io::AsyncReadExt + Unpin,
W: tokio::io::AsyncWriteExt + Unpin,
{
let mut buf = BytesMut::with_capacity(BUF_SIZE);
let mut buf = BytesMut::with_capacity(TCP_BUF_SIZE);
loop {
tokio::select! {
@@ -81,7 +81,7 @@ pub async fn run_udp_bridge(
muxer: Muxer,
mut v_rx: mpsc::Receiver<Bytes>,
) {
let mut buf = [0u8; 65536];
let mut buf = BytesMut::with_capacity(UDP_BUF_SIZE);
loop {
tokio::select! {
+27 -29
View File
@@ -13,7 +13,7 @@ use crate::{
engine::TunnelEngine,
handler::StreamHandler,
muxer::{MuxMessage, Muxer},
BUF_SIZE, CHANNEL_SIZE,
MESSAGE_CHANNEL_SIZE, TCP_BUF_SIZE,
},
tlseng::profile::BrowserProfile,
};
@@ -53,7 +53,7 @@ impl Connection {
Self {
inbound,
outbound,
read_buf: BytesMut::with_capacity(BUF_SIZE),
read_buf: BytesMut::with_capacity(TCP_BUF_SIZE),
codec: Codec::new(init),
}
}
@@ -62,7 +62,7 @@ impl Connection {
Self {
inbound,
outbound,
read_buf: BytesMut::with_capacity(BUF_SIZE),
read_buf: BytesMut::with_capacity(TCP_BUF_SIZE),
codec: Codec::new(false),
}
}
@@ -138,8 +138,11 @@ impl ClientHandler {
}
}
let (mux_tx, mux_rx) = mpsc::channel(CHANNEL_SIZE);
let muxer = Muxer::new(mux_tx, true);
// --- ИЗМЕНЕНИЕ ДЛЯ НОВОГО MUXER ---
let (control_tx, control_rx) = mpsc::channel(MESSAGE_CHANNEL_SIZE);
let (data_tx, data_rx) = mpsc::channel(MESSAGE_CHANNEL_SIZE);
let muxer = Muxer::new(control_tx, data_tx, true);
let handler =
std::sync::Arc::new(StreamHandler::new(muxer.clone(), ConnectionRole::Client));
@@ -149,7 +152,8 @@ impl ClientHandler {
outbound: conn.outbound,
codec: conn.codec,
read_buf: conn.read_buf,
mux_rx,
control_rx, // Передаем оба ресивера в Engine
data_rx, // Передаем оба ресивера в Engine
handler,
token: token.clone(),
};
@@ -203,17 +207,17 @@ impl TunnelHandler for ClientHandler {
target,
} => {
let stream_id = self.muxer.next_id();
let (v_tx, mut v_rx) = mpsc::channel::<bytes::Bytes>(BUF_SIZE);
let (v_tx, mut v_rx) = mpsc::channel::<bytes::Bytes>(TCP_BUF_SIZE);
self.muxer.register_stream(stream_id, v_tx);
// Используем send_control для отправки FrameType::Connect
self.muxer
.send_to_netwrok(MuxMessage {
.send_control(
stream_id,
frame_type: FrameType::Connect,
data: bytes::Bytes::from(target.to_string()),
})
.await
.map_err(|e| e.to_string())?;
FrameType::Connect,
bytes::Bytes::from(target.to_string()),
)
.await?;
let first_payload =
tokio::time::timeout(std::time::Duration::from_secs(10), v_rx.recv())
@@ -251,6 +255,7 @@ impl TunnelHandler for ClientHandler {
}
}
}
pub struct ServerHandler {
pub conn: Connection,
pub token: CancellationToken,
@@ -258,8 +263,6 @@ pub struct ServerHandler {
impl ServerHandler {
async fn handle_fallback(outbound: &mut OwnedWriteHalf) {
// Формируем правдоподобный HTTP-ответ.
// Заголовок "Server: nginx" помогает мимикрировать под обычный веб-сервер.
let fallback_response = "HTTP/1.1 302 Found\r\n\
Server: nginx/1.18.0 (Ubuntu)\r\n\
Location: https://www.ubuntu.com/\r\n\
@@ -267,11 +270,8 @@ impl ServerHandler {
Connection: close\r\n\
\r\n";
// Пытаемся отправить ответ сканеру/цензору, игнорируя ошибки (если он уже отключился)
let _ = outbound.write_all(fallback_response.as_bytes()).await;
let _ = outbound.flush().await;
// Корректно закрываем соединение (отправляем FIN/RST)
let _ = outbound.shutdown().await;
}
}
@@ -280,10 +280,12 @@ impl ServerHandler {
impl TunnelHandler for ServerHandler {
async fn run(mut self) -> Result<(), String> {
info!("Acting as TLS Server");
let (mux_tx, mux_rx) = mpsc::channel(CHANNEL_SIZE);
let muxer = Muxer::new(mux_tx, false);
// Тайм-аут для защиты от "сканеров-молчунов" (Active Probing)
// --- ИЗМЕНЕНИЕ ДЛЯ НОВОГО MUXER ---
let (control_tx, control_rx) = mpsc::channel(MESSAGE_CHANNEL_SIZE);
let (data_tx, data_rx) = mpsc::channel(MESSAGE_CHANNEL_SIZE);
let muxer = Muxer::new(control_tx, data_tx, false);
let handshake_timeout = std::time::Duration::from_secs(5);
let hello = loop {
@@ -294,7 +296,6 @@ impl TunnelHandler for ServerHandler {
{
Ok(b) => break b,
Err(e) if e.action == ErrorAction::Wait => {
// Используем timeout: если за 5 сек не пришел полный handshake - рвем/фолбечим
let read_res = tokio::time::timeout(
handshake_timeout,
self.conn.inbound.read_buf(&mut self.conn.read_buf),
@@ -306,35 +307,31 @@ impl TunnelHandler for ServerHandler {
return Err("Client closed connection before handshake".into());
}
Ok(Ok(_)) => {
// Прочитали кусок, идем на следующий круг проверять handshake
continue;
}
Ok(Err(e)) => {
return Err(format!("Socket read error: {}", e));
}
Err(_) => {
// СРАБОТАЛ ТАЙМ-АУТ
netrunner_logger::warn!(
"Handshake timeout (Scanner detected). Triggering fallback."
);
ServerHandler::handle_fallback(&mut self.conn.outbound).await;
return Ok(()); // Выходим без ошибки, чтобы не крашить сервер
return Ok(());
}
}
}
Err(e) => {
// ПРИШЕЛ МУСОР ИЛИ НЕВЕРНЫЙ ФОРМАТ (HTTP сканер)
netrunner_logger::warn!(
error = ?e,
"Invalid handshake format (Scanner detected). Triggering fallback."
);
ServerHandler::handle_fallback(&mut self.conn.outbound).await;
return Ok(()); // Выходим без ошибки
return Ok(());
}
}
};
// --- ЕСЛИ МЫ ЗДЕСЬ, ХЕНДШЕЙК ПРОШЕЛ УСПЕШНО ---
self.conn
.outbound
.write_all(&hello)
@@ -348,7 +345,8 @@ impl TunnelHandler for ServerHandler {
outbound: self.conn.outbound,
codec: self.conn.codec,
read_buf: self.conn.read_buf,
mux_rx,
control_rx, // Передаем оба ресивера в Engine
data_rx, // Передаем оба ресивера в Engine
handler,
token: self.token,
}
+160 -90
View File
@@ -5,7 +5,7 @@ use netrunner_logger::{debug, error, info};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::tcp::{OwnedReadHalf, OwnedWriteHalf},
sync::mpsc::Receiver,
sync::{mpsc::Receiver, Mutex},
};
use tokio_util::sync::CancellationToken;
@@ -22,99 +22,161 @@ pub struct TunnelEngine {
pub outbound: OwnedWriteHalf,
pub codec: Codec,
pub read_buf: BytesMut,
pub mux_rx: Receiver<MuxMessage>,
pub control_rx: Receiver<MuxMessage>,
pub data_rx: Receiver<MuxMessage>,
pub handler: Arc<StreamHandler>,
pub token: CancellationToken,
}
impl TunnelEngine {
pub async fn run(self) -> Result<(), String> {
let mut inbound = self.inbound;
let mut outbound = self.outbound;
let mut codec = self.codec;
let mut read_buf = self.read_buf;
let mut mux_rx = self.mux_rx;
let inbound = self.inbound;
let outbound = self.outbound;
// Оборачиваем кодек в потокобезопасный мьютекс
let codec = Arc::new(Mutex::new(self.codec));
let read_buf = self.read_buf;
let control_rx = self.control_rx;
let data_rx = self.data_rx;
let handler = self.handler;
let token = self.token;
let mut heartbeat = tokio::time::interval(std::time::Duration::from_secs(15));
loop {
tokio::select! {
_ = token.cancelled() => {
info!("TunnelEngine: Shutdown signal received. Closing...");
return Ok(());
}
res = Self::process_inbound(&mut inbound, &mut codec, &mut read_buf, &handler) => {
res?
}
// Клонируем Arc для независимых тасок
let codec_reader = codec.clone();
let codec_writer = codec.clone();
_ = heartbeat.tick() => {
let msg = MuxMessage { stream_id: 0, frame_type: FrameType::Heartbeat, data: Bytes::new() };
Self::handle_outbound(&mut outbound, &mut codec, msg).await?;
}
let token_reader = token.clone();
let token_writer = token.clone();
msg_opt = mux_rx.recv() => {
if let Some(msg) = msg_opt {
Self::handle_outbound(&mut outbound, &mut codec, msg).await?;
} else {
// ==========================================
// 1. READER TASK (Только чтение и расшифровка)
// ==========================================
let reader_handle = tokio::spawn(async move {
let mut read_buf = read_buf;
let mut inbound = inbound;
loop {
tokio::select! {
_ = token_reader.cancelled() => {
info!("Reader Task: Shutdown signal received.");
break;
}
res = inbound.read_buf(&mut read_buf) => {
let n = res.map_err(|e| e.to_string())?;
// --- ИСПРАВЛЕННАЯ ЛОГИКА EOF ---
if n == 0 {
if read_buf.is_empty() {
info!("Connection closed by peer (Clean EOF)");
} else {
error!("Connection abruptly closed by peer (Incomplete frame: {} bytes left)", read_buf.len());
}
// ВЫХОДИМ В ЛЮБОМ СЛУЧАЕ, СОКЕТ МЕРТВ!
return Err::<(), String>("EOF".into());
}
// Вектор для фреймов. Мы расшифруем всё быстро,
// сложим сюда и отпустим лок Кодека.
let mut frames = Vec::new();
{
// Блокируем кодек только на время математики (расшифровки)
let mut c = codec_reader.lock().await;
loop {
match c.inbound(&mut read_buf) {
Ok(Some(frame)) => frames.push(frame),
Ok(None) => break,
Err(e) => {
if e.action == ErrorAction::Wait {
break;
}
if e.action == ErrorAction::Drop {
error!("CRITICAL: Crypto tampering or sync lost. Hard dropping tunnel!");
return Err("Crypto drop".into());
}
error!(error = ?e, "Codec inbound failed");
return Err(format!("Codec error: {:?}", e));
}
}
}
} // Лок кодека отпущен!
// Обрабатываем фреймы (I/O) без лока, не мешая Writer Task
for frame in frames {
handler.handle(frame).await;
}
}
}
}
}
Ok(())
}
Ok::<(), String>(())
});
async fn process_inbound(
inbound: &mut OwnedReadHalf,
codec: &mut Codec,
read_buf: &mut BytesMut,
handler: &Arc<StreamHandler>,
) -> Result<(), String> {
let n = inbound
.read_buf(read_buf)
.await
.map_err(|e| e.to_string())?;
// ==========================================
// 2. WRITER TASK (Только шифрование и отправка)
// ==========================================
let writer_handle = tokio::spawn(async move {
let mut outbound = outbound;
let mut control_rx = control_rx;
let mut data_rx = data_rx;
let mut heartbeat = tokio::time::interval(std::time::Duration::from_secs(15));
if n == 0 && read_buf.is_empty() {
netrunner_logger::info!("Connection closed by peer (EOF detected)");
return Err("EOF".into());
}
loop {
tokio::select! {
biased; // Приоритет сверху вниз
loop {
match codec.inbound(read_buf) {
Ok(Some(frame)) => {
handler.handle(frame).await;
}
Ok(None) => break,
Err(e) => {
if e.action == ErrorAction::Wait {
_ = token_writer.cancelled() => {
info!("Writer Task: Shutdown signal received.");
break;
}
if e.action == ErrorAction::Drop {
// ТСПУ подмешал мусор ИЛИ ключи не совпали.
// Мгновенный выход с ошибкой приведет к обрыву TCP-соединения.
netrunner_logger::error!(
"CRITICAL: Crypto tampering or sync lost. Hard dropping tunnel!"
);
return Err("Crypto drop".into());
// FAST TRACK: Управляющие команды
msg_opt = control_rx.recv() => {
if let Some(msg) = msg_opt {
Self::handle_outbound(&mut outbound, &codec_writer, msg).await?;
} else {
break;
}
}
error!(error = ?e, "Codec inbound failed");
return Err(format!("Codec error: {:?}", e));
// Пинг
_ = heartbeat.tick() => {
let msg = MuxMessage { stream_id: 0, frame_type: FrameType::Heartbeat, data: Bytes::new() };
Self::handle_outbound(&mut outbound, &codec_writer, msg).await?;
}
// SLOW TRACK: Данные
msg_opt = data_rx.recv() => {
if let Some(msg) = msg_opt {
Self::handle_outbound(&mut outbound, &codec_writer, msg).await?;
} else {
break;
}
}
}
}
Ok::<(), String>(())
});
// ==========================================
// Ожидаем завершения обеих тасок.
// Если одна падает с ошибкой, убиваем туннель.
// ==========================================
let res = tokio::select! {
res = reader_handle => res.unwrap_or_else(|e| Err(format!("Reader panic: {}", e))),
res = writer_handle => res.unwrap_or_else(|e| Err(format!("Writer panic: {}", e))),
};
if let Err(e) = &res {
error!("TunnelEngine critical failure: {}", e);
}
Ok(())
res
}
async fn handle_outbound(
outbound: &mut OwnedWriteHalf,
codec: &mut Codec,
codec: &Arc<Mutex<Codec>>,
msg: MuxMessage,
) -> Result<(), String> {
const MAX_CHUNK_SIZE: usize = 16000;
@@ -123,36 +185,44 @@ impl TunnelEngine {
let stream_id = msg.stream_id;
let frame_type = msg.frame_type;
if data.is_empty() {
match codec.encrypt_data(stream_id, frame_type.clone(), Bytes::new()) {
Ok(pkt) => {
outbound.write_all(&pkt).await.map_err(|e| e.to_string())?;
// Вектор зашифрованных пакетов. Собираем их быстро под локом.
let mut packets = Vec::new();
{
// Берем лок кодека только для шифрования
let mut c = codec.lock().await;
if data.is_empty() {
match c.encrypt_data(stream_id, frame_type.clone(), Bytes::new()) {
Ok(pkt) => packets.push(pkt),
Err(e) => {
error!(stream_id, error = ?e, "Encryption failed for empty message");
return Err(format!("Encryption error: {:?}", e));
}
}
Err(e) => {
error!(stream_id, error = ?e, "Encryption failed for empty message");
return Err(format!("Encryption error: {:?}", e));
} else {
while !data.is_empty() {
let chunk_size = std::cmp::min(data.len(), MAX_CHUNK_SIZE);
let chunk = data.split_to(chunk_size);
match c.encrypt_data(stream_id, frame_type.clone(), chunk) {
Ok(pkt) => packets.push(pkt),
Err(e) => {
error!(stream_id, error = ?e, "Encryption failed for chunked message");
return Err(format!("Encryption error: {:?}", e));
}
}
}
}
return Ok(());
}
} // Лок кодека отпущен!
while !data.is_empty() {
let chunk_size = std::cmp::min(data.len(), MAX_CHUNK_SIZE);
let chunk = data.split_to(chunk_size);
match codec.encrypt_data(stream_id, frame_type.clone(), chunk) {
Ok(pkt) => {
outbound.write_all(&pkt).await.map_err(|e| {
error!(stream_id, error = %e, "Failed to write encrypted data to network");
e.to_string()
})?;
}
Err(e) => {
error!(stream_id, error = ?e, "Encryption failed for chunked message");
return Err(format!("Encryption error: {:?}", e));
}
}
// Выполняем I/O операцию (которая может зависнуть при плохой сети) БЕЗ лока кодека.
// Это позволяет Reader Task продолжать читать сеть!
for pkt in packets {
outbound.write_all(&pkt).await.map_err(|e| {
error!(stream_id, error = %e, "Failed to write encrypted data to network");
e.to_string()
})?;
}
debug!(stream_id, "Outbound packet sent successfully");
+3 -2
View File
@@ -4,5 +4,6 @@ pub mod engine;
pub mod handler;
pub mod muxer;
pub const BUF_SIZE: usize = 65536;
pub const CHANNEL_SIZE: usize = 16;
pub const TCP_BUF_SIZE: usize = 1024 * 512;
pub const UDP_BUF_SIZE: usize = 1024 * 64;
pub const MESSAGE_CHANNEL_SIZE: usize = 1024 * 16;
+28 -35
View File
@@ -1,10 +1,9 @@
use crate::protocol::codec::frame::FrameType;
use bytes::Bytes;
use dashmap::DashMap; // Добавь в Cargo.toml: dashmap = "6.0"
use dashmap::DashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc::error::SendError;
use tokio::sync::mpsc::Sender;
use tokio::sync::mpsc::{error::SendError, Sender};
pub struct IdGenerator {
counter: AtomicU32,
@@ -31,15 +30,21 @@ pub struct MuxMessage {
#[derive(Clone)]
pub struct Muxer {
to_network: Sender<MuxMessage>,
control_tx: Sender<MuxMessage>,
data_tx: Sender<MuxMessage>,
streams: Arc<DashMap<u32, Sender<Bytes>>>,
id_gen: Arc<IdGenerator>,
}
impl Muxer {
pub fn new(to_network: Sender<MuxMessage>, is_client: bool) -> Self {
pub fn new(
control_tx: Sender<MuxMessage>,
data_tx: Sender<MuxMessage>,
is_client: bool,
) -> Self {
Self {
to_network,
control_tx,
data_tx,
streams: Arc::new(DashMap::new()),
id_gen: Arc::new(IdGenerator::new(is_client)),
}
@@ -49,30 +54,19 @@ impl Muxer {
self.id_gen.next()
}
// Обычные данные летят в канал с низким приоритетом
pub async fn send_to_netwrok(&self, message: MuxMessage) -> Result<(), SendError<MuxMessage>> {
self.to_network.send(message).await
}
pub fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
self.streams.insert(stream_id, tx);
netrunner_logger::debug!(
stream_id,
total_active = self.streams.len(),
"MUXER: [REGISTER] Stream added"
);
}
pub fn remove_stream(&self, stream_id: u32) {
self.streams.remove(&stream_id);
self.data_tx.send(message).await
}
// Управляющие фреймы летят в VIP-канал
pub async fn send_control(
&self,
stream_id: u32,
f_type: FrameType,
data: Bytes,
) -> Result<(), String> {
self.to_network
self.control_tx
.send(MuxMessage {
stream_id,
frame_type: f_type,
@@ -82,32 +76,31 @@ impl Muxer {
.map_err(|e| e.to_string())
}
pub fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
self.streams.insert(stream_id, tx);
}
pub fn remove_stream(&self, stream_id: u32) {
self.streams.remove(&stream_id);
}
pub async fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
// DashMap позволяет получить доступ к элементу без явного RwLock
let tx = self.streams.get(&stream_id).map(|r| r.value().clone());
if let Some(tx) = tx {
if data.is_empty() {
netrunner_logger::debug!(stream_id, "MUXER: [EOF] Forwarding EOF to local handler");
} else {
netrunner_logger::trace!(
stream_id,
len = data.len(),
"MUXER: [DISPATCH] Sending data"
);
}
if let Err(_e) = tx.send(data).await {
netrunner_logger::debug!(
netrunner_logger::warn!(
stream_id,
"MUXER: [WARN] Local channel closed, dropping packet"
"MUXER: [INBOUND_ERR] Local channel closed, dropping packet"
);
self.remove_stream(stream_id);
}
} else {
// Уменьшил уровень лога до TRACE, потому что это штатная ситуация при больших буферах!
netrunner_logger::trace!(
stream_id,
"MUXER: [IGNORE] Packet for already closed stream"
len = data.len(),
"MUXER: [IGNORE] Received data for already closed stream (draining pipe)"
);
}
}