global buf size mtu based config and protocol rename
This commit is contained in:
@@ -0,0 +1,132 @@
|
||||
use crate::nrxp::codec::frame::FrameType;
|
||||
use crate::net::connection::muxer::{MuxMessage, Muxer};
|
||||
use crate::net::network::NetworkConfig;
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use netrunner_logger::{debug, error};
|
||||
use tokio::net::UdpSocket;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
pub async fn run_tcp_bridge<R, W>(
|
||||
stream_id: u32,
|
||||
mut reader: R,
|
||||
mut writer: W,
|
||||
muxer: Muxer,
|
||||
mut v_rx: mpsc::Receiver<Bytes>,
|
||||
) where
|
||||
R: tokio::io::AsyncReadExt + Unpin,
|
||||
W: tokio::io::AsyncWriteExt + Unpin,
|
||||
{
|
||||
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;
|
||||
}
|
||||
Ok(_) => {
|
||||
let msg = MuxMessage {
|
||||
stream_id,
|
||||
frame_type: FrameType::Data,
|
||||
data: buf.split().freeze(),
|
||||
};
|
||||
if muxer.send_to_network(msg).await.is_err() { break; }
|
||||
}
|
||||
Err(e) => {
|
||||
error!(stream_id, error = %e, "Socket read error");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
maybe_data = v_rx.recv() => {
|
||||
match maybe_data {
|
||||
Some(data) => {
|
||||
if data.is_empty() { break; }
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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;
|
||||
|
||||
muxer.remove_stream(stream_id);
|
||||
}
|
||||
|
||||
pub async fn run_udp_bridge(
|
||||
stream_id: u32,
|
||||
socket: UdpSocket,
|
||||
muxer: Muxer,
|
||||
mut v_rx: mpsc::Receiver<Bytes>,
|
||||
) {
|
||||
let mut buf = BytesMut::with_capacity(NetworkConfig::global().udp_buffer_size);
|
||||
|
||||
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; }
|
||||
}
|
||||
Err(e) => {
|
||||
error!(stream_id, error = %e, "UDP socket read error");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
maybe_data = v_rx.recv() => {
|
||||
match maybe_data {
|
||||
Some(data) => {
|
||||
if data.is_empty() { break; }
|
||||
if let Err(e) = socket.send(&data).await {
|
||||
error!(stream_id, error = %e, "UDP socket write error");
|
||||
break;
|
||||
}
|
||||
}
|
||||
None => {
|
||||
debug!(stream_id, "Virtual channel closed (UDP)");
|
||||
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;
|
||||
|
||||
muxer.remove_stream(stream_id);
|
||||
}
|
||||
@@ -0,0 +1,388 @@
|
||||
use crate::{
|
||||
nrxp::{
|
||||
codec::{
|
||||
codec::Codec,
|
||||
frame::FrameType,
|
||||
socks::{SocksReply, SocksRequest},
|
||||
},
|
||||
errors::ErrorAction,
|
||||
parser::parser::Parser,
|
||||
},
|
||||
net::{
|
||||
connection::{
|
||||
bridge::run_tcp_bridge, engine::TunnelEngine, handler::StreamHandler, muxer::Muxer,
|
||||
},
|
||||
network::NetworkConfig,
|
||||
},
|
||||
tlseng::profile::BrowserProfile,
|
||||
};
|
||||
use bytes::BytesMut;
|
||||
use netrunner_logger::{info, warn};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::{
|
||||
tcp::{OwnedReadHalf, OwnedWriteHalf},
|
||||
TcpStream,
|
||||
},
|
||||
sync::mpsc,
|
||||
};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub enum ConnectionRole {
|
||||
Client,
|
||||
Server,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
pub trait TunnelHandler {
|
||||
async fn run(self) -> Result<(), String>;
|
||||
}
|
||||
|
||||
pub struct Connection {
|
||||
pub inbound: OwnedReadHalf,
|
||||
pub outbound: OwnedWriteHalf,
|
||||
pub read_buf: BytesMut,
|
||||
pub codec: Codec,
|
||||
}
|
||||
|
||||
impl Connection {
|
||||
pub fn new(stream: TcpStream, init: bool) -> Self {
|
||||
let (inbound, outbound) = stream.into_split();
|
||||
Self {
|
||||
inbound,
|
||||
outbound,
|
||||
read_buf: BytesMut::with_capacity(NetworkConfig::global().tcp_buffer_size),
|
||||
codec: Codec::new(init),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_raw(inbound: OwnedReadHalf, outbound: OwnedWriteHalf) -> Self {
|
||||
Self {
|
||||
inbound,
|
||||
outbound,
|
||||
read_buf: BytesMut::with_capacity(NetworkConfig::global().tcp_buffer_size),
|
||||
codec: Codec::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn read_socks_request(&mut self) -> Result<SocksRequest, String> {
|
||||
loop {
|
||||
match SocksRequest::parse(&mut self.read_buf) {
|
||||
Ok(Some(req)) => return Ok(req),
|
||||
Ok(None) => {}
|
||||
Err(e) => return Err(format!("Socks parse error: {}", e)),
|
||||
}
|
||||
let n = self
|
||||
.inbound
|
||||
.read_buf(&mut self.read_buf)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
if n == 0 {
|
||||
return Err("Client closed connection".into());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn send_socks_reply(&mut self, reply: SocksReply) -> Result<(), String> {
|
||||
let mut buf = BytesMut::with_capacity(24);
|
||||
reply.write_to(&mut buf);
|
||||
self.outbound
|
||||
.write_all(&buf)
|
||||
.await
|
||||
.map_err(|e| e.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ClientHandler {
|
||||
pub conn: Connection,
|
||||
pub muxer: Muxer,
|
||||
}
|
||||
|
||||
impl ClientHandler {
|
||||
pub async fn connect(remote_proxy_addr: &str) -> Result<Muxer, String> {
|
||||
let stream = TcpStream::connect(remote_proxy_addr)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
let (inbound, outbound) = stream.into_split();
|
||||
|
||||
let mut conn = Connection::new_raw(inbound, outbound);
|
||||
|
||||
let ch = conn
|
||||
.codec
|
||||
.make_client_handshake(&BrowserProfile::CHROME_131, "ubuntu.com")
|
||||
.map_err(|e| format!("{:?}", e))?;
|
||||
conn.outbound
|
||||
.write_all(&ch)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
loop {
|
||||
match conn.codec.process_handshake(&mut conn.read_buf) {
|
||||
Ok(_) => break,
|
||||
Err(e) if e.action == ErrorAction::Wait => {
|
||||
let n = conn
|
||||
.inbound
|
||||
.read_buf(&mut conn.read_buf)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
if n == 0 {
|
||||
return Err("EOF during handshake".into());
|
||||
}
|
||||
}
|
||||
Err(e) => return Err(format!("TLS error: {:?}", e)),
|
||||
}
|
||||
}
|
||||
|
||||
let (control_tx, control_rx) = mpsc::channel(NetworkConfig::global().channel_capacity);
|
||||
let (data_tx, data_rx) = mpsc::channel(NetworkConfig::global().channel_capacity);
|
||||
|
||||
let muxer = Muxer::new(control_tx, data_tx, true);
|
||||
|
||||
let handler =
|
||||
std::sync::Arc::new(StreamHandler::new(muxer.clone(), ConnectionRole::Client));
|
||||
|
||||
let engine = TunnelEngine {
|
||||
inbound: conn.inbound,
|
||||
outbound: conn.outbound,
|
||||
codec: conn.codec,
|
||||
read_buf: conn.read_buf,
|
||||
control_rx,
|
||||
data_rx,
|
||||
handler,
|
||||
};
|
||||
|
||||
tokio::spawn(async move { engine.run().await });
|
||||
|
||||
Ok(muxer)
|
||||
}
|
||||
|
||||
async fn handle_udp_associate(&mut self) -> Result<(), String> {
|
||||
let reply = SocksReply::ConnectResult {
|
||||
reply_code: 0x00,
|
||||
atyp: 0x01,
|
||||
addr: [0, 0, 0, 0],
|
||||
port: 0,
|
||||
};
|
||||
self.conn.send_socks_reply(reply).await?;
|
||||
|
||||
let mut buf = [0u8; 1024];
|
||||
loop {
|
||||
if self
|
||||
.conn
|
||||
.inbound
|
||||
.read(&mut buf)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?
|
||||
== 0
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl TunnelHandler for ClientHandler {
|
||||
async fn run(mut self) -> Result<(), String> {
|
||||
info!("Starting SOCKS multiplexed handling");
|
||||
|
||||
self.conn.read_socks_request().await?;
|
||||
self.conn
|
||||
.send_socks_reply(SocksReply::HandshakeSelect { method: 0x00 })
|
||||
.await?;
|
||||
|
||||
let req = self.conn.read_socks_request().await?;
|
||||
|
||||
match req {
|
||||
SocksRequest::Connect {
|
||||
command: 0x01,
|
||||
target,
|
||||
} => {
|
||||
let stream_id = self.muxer.next_id();
|
||||
let (v_tx, mut v_rx) =
|
||||
mpsc::channel::<bytes::Bytes>(NetworkConfig::global().tcp_buffer_size);
|
||||
self.muxer.register_stream(stream_id, v_tx);
|
||||
|
||||
self.muxer
|
||||
.send_control(
|
||||
stream_id,
|
||||
FrameType::Connect,
|
||||
bytes::Bytes::from(target.to_string()),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let first_payload =
|
||||
tokio::time::timeout(std::time::Duration::from_secs(10), v_rx.recv())
|
||||
.await
|
||||
.map_err(|_| "Timeout waiting for proxy response")?
|
||||
.ok_or("No data from proxy")?;
|
||||
|
||||
if first_payload.len() >= 2 && first_payload[1] == 0x00 {
|
||||
self.conn
|
||||
.outbound
|
||||
.write_all(&first_payload)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
} else {
|
||||
self.conn.outbound.write_all(&first_payload).await.ok();
|
||||
return Err("Proxy rejected connection".into());
|
||||
}
|
||||
|
||||
let browser_in = self.conn.inbound;
|
||||
let browser_out = self.conn.outbound;
|
||||
let muxer = self.muxer;
|
||||
|
||||
tokio::spawn(async move {
|
||||
run_tcp_bridge(stream_id, browser_in, browser_out, muxer, v_rx).await;
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
SocksRequest::Connect { command: 0x03, .. } => {
|
||||
info!("Handling UDP Associate request");
|
||||
self.handle_udp_associate().await
|
||||
}
|
||||
|
||||
_ => Err("Unsupported SOCKS command".into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ServerHandler {
|
||||
pub conn: Connection,
|
||||
pub token: CancellationToken,
|
||||
}
|
||||
impl ServerHandler {
|
||||
async fn handle_stealth_fallback(
|
||||
mut client_inbound: OwnedReadHalf,
|
||||
mut client_outbound: OwnedWriteHalf,
|
||||
initial_data: bytes::Bytes,
|
||||
) {
|
||||
let target_host = "ubuntu.com:443";
|
||||
info!(target = %target_host, "Stealth fallback: bridging to Target");
|
||||
|
||||
let target_stream = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(3),
|
||||
TcpStream::connect(target_host),
|
||||
)
|
||||
.await;
|
||||
|
||||
match target_stream {
|
||||
Ok(Ok(target_server)) => {
|
||||
let (mut server_read, mut server_write) = target_server.into_split();
|
||||
|
||||
if !initial_data.is_empty() {
|
||||
if let Err(e) = server_write.write_all(&initial_data).await {
|
||||
warn!("Failed to push initial data to fallback: {}", e);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
let res = tokio::io::copy_bidirectional(
|
||||
&mut tokio::io::join(&mut client_inbound, &mut client_outbound),
|
||||
&mut tokio::io::join(&mut server_read, &mut server_write),
|
||||
)
|
||||
.await;
|
||||
|
||||
match res {
|
||||
Ok((from_client, from_server)) => {
|
||||
info!(
|
||||
"Fallback closed. Sent: {} bytes, Recv: {} bytes",
|
||||
from_client, from_server
|
||||
);
|
||||
}
|
||||
Err(e) => warn!("Fallback bridge error: {}", e),
|
||||
}
|
||||
}
|
||||
Ok(Err(e)) => warn!("Fallback connect error: {}", e),
|
||||
Err(_) => warn!("Fallback connection timed out (Target unreachable)"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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().channel_capacity);
|
||||
let (data_tx, data_rx) = mpsc::channel(NetworkConfig::global().channel_capacity);
|
||||
let muxer = Muxer::new(control_tx, data_tx, false);
|
||||
|
||||
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
|
||||
.make_server_handshake(&mut self.conn.read_buf)
|
||||
{
|
||||
Ok(b) => break b,
|
||||
Err(e) if e.action == ErrorAction::Wait => {
|
||||
let read_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()),
|
||||
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,
|
||||
buf_snapshot,
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Auth/Format failed: {:?}. Going stealth.", e);
|
||||
|
||||
info!(
|
||||
"DEBUG: Restoring {} bytes from snapshot for fallback",
|
||||
buf_snapshot.len()
|
||||
);
|
||||
|
||||
ServerHandler::handle_stealth_fallback(
|
||||
self.conn.inbound,
|
||||
self.conn.outbound,
|
||||
buf_snapshot,
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
self.conn
|
||||
.outbound
|
||||
.write_all(&hello)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
let handler = std::sync::Arc::new(StreamHandler::new(muxer, ConnectionRole::Server));
|
||||
|
||||
TunnelEngine {
|
||||
inbound: self.conn.inbound,
|
||||
outbound: self.conn.outbound,
|
||||
codec: self.conn.codec,
|
||||
read_buf: self.conn.read_buf,
|
||||
control_rx,
|
||||
data_rx,
|
||||
handler,
|
||||
}
|
||||
.run()
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use netrunner_logger::{debug, error, info};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::tcp::{OwnedReadHalf, OwnedWriteHalf},
|
||||
sync::{mpsc::Receiver, Mutex},
|
||||
};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::{
|
||||
net::connection::{handler::StreamHandler, muxer::MuxMessage},
|
||||
nrxp::{
|
||||
codec::{codec::Codec, frame::FrameType},
|
||||
errors::ErrorAction,
|
||||
},
|
||||
};
|
||||
|
||||
pub struct TunnelEngine {
|
||||
pub inbound: OwnedReadHalf,
|
||||
pub outbound: OwnedWriteHalf,
|
||||
pub codec: Codec,
|
||||
pub read_buf: BytesMut,
|
||||
pub control_rx: Receiver<MuxMessage>,
|
||||
pub data_rx: Receiver<MuxMessage>,
|
||||
pub handler: Arc<StreamHandler>,
|
||||
}
|
||||
|
||||
impl TunnelEngine {
|
||||
pub async fn run(self) -> Result<(), String> {
|
||||
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 = CancellationToken::new();
|
||||
|
||||
let codec_reader = codec.clone();
|
||||
let codec_writer = codec.clone();
|
||||
|
||||
let token_reader = token.clone();
|
||||
let token_writer = token.clone();
|
||||
|
||||
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())?;
|
||||
|
||||
|
||||
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));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
for frame in frames {
|
||||
handler.handle(frame).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok::<(), String>(())
|
||||
});
|
||||
|
||||
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));
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
|
||||
_ = token_writer.cancelled() => {
|
||||
info!("Writer Task: Shutdown signal received.");
|
||||
break;
|
||||
}
|
||||
|
||||
|
||||
msg_opt = control_rx.recv() => {
|
||||
if let Some(msg) = msg_opt {
|
||||
Self::handle_outbound(&mut outbound, &codec_writer, msg).await?;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
_ = heartbeat.tick() => {
|
||||
let msg = MuxMessage { stream_id: 0, frame_type: FrameType::Heartbeat, data: Bytes::new() };
|
||||
Self::handle_outbound(&mut outbound, &codec_writer, msg).await?;
|
||||
}
|
||||
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
res
|
||||
}
|
||||
|
||||
async fn handle_outbound(
|
||||
outbound: &mut OwnedWriteHalf,
|
||||
codec: &Arc<Mutex<Codec>>,
|
||||
msg: MuxMessage,
|
||||
) -> Result<(), String> {
|
||||
const MAX_CHUNK_SIZE: usize = 1024 * 4;
|
||||
|
||||
let mut data = msg.data;
|
||||
let stream_id = msg.stream_id;
|
||||
let frame_type = msg.frame_type;
|
||||
|
||||
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));
|
||||
}
|
||||
}
|
||||
} 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));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use netrunner_logger::{debug, error, info};
|
||||
|
||||
use crate::{
|
||||
net::{
|
||||
connection::{
|
||||
bridge::{run_tcp_bridge, run_udp_bridge},
|
||||
connection::ConnectionRole,
|
||||
muxer::Muxer,
|
||||
},
|
||||
network::NetworkConfig,
|
||||
},
|
||||
nrxp::codec::{
|
||||
frame::{Frame, FrameType},
|
||||
socks::SocksReply,
|
||||
},
|
||||
};
|
||||
|
||||
pub struct StreamHandler {
|
||||
muxer: Muxer,
|
||||
role: ConnectionRole,
|
||||
}
|
||||
|
||||
impl StreamHandler {
|
||||
pub fn new(muxer: Muxer, role: ConnectionRole) -> Self {
|
||||
Self { muxer, role }
|
||||
}
|
||||
|
||||
pub async fn handle(&self, frame: Frame) {
|
||||
let stream_id = frame.header.stream_id;
|
||||
|
||||
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,
|
||||
FrameType::Data => self.on_data(stream_id, frame.payload).await,
|
||||
FrameType::UdpData => self.on_udp_data(stream_id, frame.payload).await,
|
||||
FrameType::Close => self.on_close(stream_id).await,
|
||||
_ => debug!(stream_id, "Unhandled frame type"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_connect(&self, stream_id: u32, payload: Bytes) {
|
||||
if self.role == ConnectionRole::Server {
|
||||
let target_str = String::from_utf8_lossy(&payload).to_string();
|
||||
let muxer = self.muxer.clone();
|
||||
|
||||
let (v_tx, v_rx) = tokio::sync::mpsc::channel(NetworkConfig::global().channel_capacity);
|
||||
muxer.register_stream(stream_id, v_tx);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let start = std::time::Instant::now();
|
||||
info!(stream_id, target = %target_str, "Attempting remote TCP connection");
|
||||
|
||||
let connect_timeout = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
tokio::net::TcpStream::connect(&target_str),
|
||||
)
|
||||
.await;
|
||||
|
||||
match connect_timeout {
|
||||
Ok(Ok(stream)) => {
|
||||
let elapsed = start.elapsed();
|
||||
info!(stream_id, target = %target_str, latency_ms = elapsed.as_millis(), "Remote TCP connection established");
|
||||
|
||||
let mut reply_buf = BytesMut::with_capacity(10);
|
||||
let reply = SocksReply::ConnectResult {
|
||||
reply_code: 0x00,
|
||||
atyp: 0x01,
|
||||
addr: [0, 0, 0, 0],
|
||||
port: 0,
|
||||
};
|
||||
reply.write_to(&mut reply_buf);
|
||||
|
||||
let _ = muxer
|
||||
.send_control(stream_id, FrameType::Connect, reply_buf.freeze())
|
||||
.await;
|
||||
|
||||
let (r, w) = stream.into_split();
|
||||
run_tcp_bridge(stream_id, r, w, muxer, v_rx).await;
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
error!(stream_id, target = %target_str, error = %e, "TCP connection failed");
|
||||
Self::send_error_reply(&muxer, stream_id, 0x01, FrameType::Connect).await;
|
||||
}
|
||||
Err(_) => {
|
||||
error!(stream_id, target = %target_str, "Connection timed out (DNS/TCP)");
|
||||
Self::send_error_reply(&muxer, stream_id, 0x04, FrameType::Connect).await;
|
||||
}
|
||||
}
|
||||
});
|
||||
} else {
|
||||
self.muxer.dispatch_to_local(stream_id, payload).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_udp_connect(&self, stream_id: u32, payload: Bytes) {
|
||||
if self.role == ConnectionRole::Server {
|
||||
let target_str = String::from_utf8_lossy(&payload).to_string();
|
||||
let muxer = self.muxer.clone();
|
||||
|
||||
let (v_tx, v_rx) = tokio::sync::mpsc::channel(NetworkConfig::global().channel_capacity);
|
||||
muxer.register_stream(stream_id, v_tx);
|
||||
|
||||
tokio::spawn(async move {
|
||||
info!(stream_id, target = %target_str, "Attempting remote UDP connection");
|
||||
|
||||
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::send_error_reply(&muxer, stream_id, 0x01, FrameType::UdpConnect)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
let mut reply_buf = BytesMut::with_capacity(10);
|
||||
let reply = SocksReply::ConnectResult {
|
||||
reply_code: 0x00,
|
||||
atyp: 0x01,
|
||||
addr: [0, 0, 0, 0],
|
||||
port: 0,
|
||||
};
|
||||
reply.write_to(&mut reply_buf);
|
||||
|
||||
let _ = muxer
|
||||
.send_control(stream_id, FrameType::UdpConnect, reply_buf.freeze())
|
||||
.await;
|
||||
|
||||
run_udp_bridge(stream_id, socket, muxer, v_rx).await;
|
||||
}
|
||||
Err(e) => {
|
||||
error!(stream_id, target = %target_str, error = %e, "UDP bind failed");
|
||||
Self::send_error_reply(&muxer, stream_id, 0x01, FrameType::UdpConnect)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
});
|
||||
} else {
|
||||
self.muxer.dispatch_to_local(stream_id, payload).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_udp_data(&self, stream_id: u32, payload: Bytes) {
|
||||
self.muxer.dispatch_to_local(stream_id, payload).await;
|
||||
}
|
||||
|
||||
async fn send_error_reply(muxer: &Muxer, stream_id: u32, code: u8, frame_type: FrameType) {
|
||||
muxer.remove_stream(stream_id);
|
||||
let mut reply_buf = BytesMut::with_capacity(10);
|
||||
let reply = SocksReply::ConnectResult {
|
||||
reply_code: code,
|
||||
atyp: 0x01,
|
||||
addr: [0, 0, 0, 0],
|
||||
port: 0,
|
||||
};
|
||||
reply.write_to(&mut reply_buf);
|
||||
let _ = muxer
|
||||
.send_control(stream_id, frame_type, reply_buf.freeze())
|
||||
.await;
|
||||
}
|
||||
|
||||
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.dispatch_to_local(stream_id, Bytes::new()).await;
|
||||
self.muxer.remove_stream(stream_id);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
pub mod bridge;
|
||||
pub mod connection;
|
||||
pub mod engine;
|
||||
pub mod handler;
|
||||
pub mod muxer;
|
||||
|
||||
pub const MESSAGE_CHANNEL_SIZE: usize = 4;
|
||||
@@ -0,0 +1,140 @@
|
||||
use crate::nrxp::codec::frame::FrameType;
|
||||
use bytes::Bytes;
|
||||
use dashmap::DashMap;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc::{error::SendError, Sender};
|
||||
|
||||
pub struct IdGenerator {
|
||||
counter: AtomicU32,
|
||||
}
|
||||
|
||||
impl IdGenerator {
|
||||
pub fn new(is_client: bool) -> Self {
|
||||
let start = if is_client { 1 } else { 2 };
|
||||
Self {
|
||||
counter: AtomicU32::new(start),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn next(&self) -> u32 {
|
||||
self.counter.fetch_add(2, Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct MuxMessage {
|
||||
pub stream_id: u32,
|
||||
pub frame_type: FrameType,
|
||||
pub data: Bytes,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Muxer {
|
||||
control_tx: Sender<MuxMessage>,
|
||||
data_tx: Sender<MuxMessage>,
|
||||
streams: Arc<DashMap<u32, Sender<Bytes>>>,
|
||||
id_gen: Arc<IdGenerator>,
|
||||
}
|
||||
|
||||
impl Muxer {
|
||||
pub fn new(
|
||||
control_tx: Sender<MuxMessage>,
|
||||
data_tx: Sender<MuxMessage>,
|
||||
is_client: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
control_tx,
|
||||
data_tx,
|
||||
streams: Arc::new(DashMap::new()),
|
||||
id_gen: Arc::new(IdGenerator::new(is_client)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn next_id(&self) -> u32 {
|
||||
self.id_gen.next()
|
||||
}
|
||||
|
||||
pub async fn send_to_network(&self, message: MuxMessage) -> Result<(), SendError<MuxMessage>> {
|
||||
self.data_tx.send(message).await
|
||||
}
|
||||
|
||||
pub async fn send_data_safe(&self, stream_id: u32, mut data: Bytes) -> Result<(), String> {
|
||||
// Лимит полезной нагрузки, чтобы вместе с заголовком и паддингом
|
||||
// пакет оставался в пределах ~1400-1450 байт.
|
||||
const MAX_PAYLOAD_CHUNK: usize = 1300;
|
||||
|
||||
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,
|
||||
data: chunk,
|
||||
})
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
// Небольшая уступка планировщику, чтобы не забить канал мгновенно
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn send_control(
|
||||
&self,
|
||||
stream_id: u32,
|
||||
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())
|
||||
}
|
||||
|
||||
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) {
|
||||
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"
|
||||
);
|
||||
self.remove_stream(stream_id);
|
||||
}
|
||||
} else {
|
||||
netrunner_logger::trace!(
|
||||
stream_id,
|
||||
len = data.len(),
|
||||
"MUXER: [IGNORE] Received data for already closed stream (draining pipe)"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod connection;
|
||||
pub mod network;
|
||||
@@ -0,0 +1,152 @@
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use crate::{
|
||||
nrxp::codec::{frame::FRAME_HEADER_SIZE, MAX_PADDING_SIZE},
|
||||
net::connection::connection::{
|
||||
ClientHandler, Connection, ConnectionRole, ServerHandler, TunnelHandler,
|
||||
},
|
||||
};
|
||||
use netrunner_logger::{error, info};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
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 async fn run(&self, token: CancellationToken) {
|
||||
let addr = format!("{}:{}", self.host, self.port);
|
||||
NetworkConfig::init_global(1350);
|
||||
match self.role {
|
||||
ConnectionRole::Client => {
|
||||
info!("Starting Client mode");
|
||||
let server_addr = self
|
||||
.remote_proxy_addr
|
||||
.as_ref()
|
||||
.ok_or("No proxy addr")
|
||||
.unwrap();
|
||||
let muxer = match ClientHandler::connect(server_addr).await {
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
error!(error = %e, "Global tunnel failed.");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let listener = TcpListener::bind(&addr).await.expect("SOCKS bind failed");
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = token.cancelled() => break,
|
||||
res = listener.accept() => {
|
||||
if let Ok((stream, _client_addr)) = res {
|
||||
let conn = Connection::new(stream, false);
|
||||
let handler = ClientHandler{ conn, muxer: muxer.clone() };
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = handler.run().await {
|
||||
error!(error = %e, "Client handler error");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
ConnectionRole::Server => {
|
||||
let listener = TcpListener::bind(&addr).await.expect("Server bind failed");
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = token.cancelled() => break,
|
||||
res = listener.accept() => {
|
||||
if let Ok((stream, client_addr)) = res {
|
||||
let conn = Connection::new(stream, true);
|
||||
let handler = ServerHandler { conn, token: token.clone() };
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = handler.run().await {
|
||||
error!(client = %client_addr, error = %e, "Server handler error");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub static GLOBAL_NET_CONFIG: OnceLock<NetworkConfig> = OnceLock::new();
|
||||
|
||||
pub struct NetworkConfig {
|
||||
pub mtu: usize,
|
||||
pub max_wire_frame_size: usize,
|
||||
pub safe_payload_size: usize,
|
||||
pub tcp_buffer_size: usize,
|
||||
pub udp_buffer_size: usize,
|
||||
pub channel_capacity: usize,
|
||||
}
|
||||
|
||||
impl NetworkConfig {
|
||||
pub fn new(system_mtu: usize) -> Self {
|
||||
let transport_overhead = 28; // IPv4 + UDP
|
||||
|
||||
let max_wire_frame = system_mtu.saturating_sub(transport_overhead);
|
||||
|
||||
let safe_payload = max_wire_frame
|
||||
.saturating_sub(FRAME_HEADER_SIZE as usize)
|
||||
.saturating_sub((MAX_PADDING_SIZE - 1) as usize);
|
||||
|
||||
let tcp_chunks_count = 65536 / safe_payload;
|
||||
let tcp_buffer = safe_payload * tcp_chunks_count;
|
||||
|
||||
let udp_chunks_count = 16384 / safe_payload;
|
||||
let udp_buffer = safe_payload * udp_chunks_count;
|
||||
|
||||
let channel_cap = 1024;
|
||||
|
||||
netrunner_logger::info!(
|
||||
mtu = system_mtu,
|
||||
payload = safe_payload,
|
||||
tcp_buf = tcp_buffer,
|
||||
"Network Optimizer: Calculations complete for current MTU"
|
||||
);
|
||||
|
||||
Self {
|
||||
mtu: system_mtu,
|
||||
max_wire_frame_size: max_wire_frame,
|
||||
safe_payload_size: safe_payload,
|
||||
tcp_buffer_size: tcp_buffer,
|
||||
udp_buffer_size: udp_buffer,
|
||||
channel_capacity: channel_cap,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn init_global(system_mtu: usize) {
|
||||
let config = Self::new(system_mtu);
|
||||
if GLOBAL_NET_CONFIG.set(config).is_err() {
|
||||
netrunner_logger::warn!("Global network config was already initialized!");
|
||||
}
|
||||
}
|
||||
pub fn global() -> &'static Self {
|
||||
GLOBAL_NET_CONFIG
|
||||
.get()
|
||||
.expect("Global NetworkConfig is not initialized! Call init_global() first.")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user