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

connection refactoring

Browse Source
This commit is contained in:
Kirill
2026-03-17 14:27:07 +07:00
parent b7c254e9c5
commit bf6e621f14
7 changed files with 210 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/src/proxy/connection/bridge.rs
+2 -3
View File
@@ -30,7 +30,7 @@ pub async fn run_proxy_bridge<R, W>(
frame_type: FrameType::Data,
data: buf.split().freeze(),
};
if muxer.to_network.send(msg).await.is_err() { break; }
if muxer.send_to_netwrok(msg).await.is_err() { break; }
}
Err(e) => {
error!(stream_id, error = %e, "Socket read error");
@@ -58,8 +58,7 @@ pub async fn run_proxy_bridge<R, W>(
}
let _ = muxer
.to_network
.send(MuxMessage {
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::Close,
data: Bytes::new(),
core/src/proxy/connection/connection.rs
+170 -135
View File
@@ -1,16 +1,3 @@
use bytes::{Bytes, BytesMut};
use netrunner_logger::{debug, error, info, instrument, trace, warn};
use std::net::SocketAddr;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{
tcp::{OwnedReadHalf, OwnedWriteHalf},
TcpStream,
},
sync::mpsc::{self},
};
use tokio_util::sync::CancellationToken;
use crate::{
protocol::{
codec::{
@@ -27,7 +14,20 @@ use crate::{
handler::StreamHandler,
muxer::{MuxMessage, Muxer},
},
tlseng::profile::BrowserProfile,
};
use bytes::{Bytes, BytesMut};
use netrunner_logger::{debug, error, info, trace, warn};
use std::net::SocketAddr;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{
tcp::{OwnedReadHalf, OwnedWriteHalf},
TcpStream,
},
sync::mpsc,
};
use tokio_util::sync::CancellationToken;
pub const BUF_SIZE: usize = 16384;
@@ -37,6 +37,11 @@ pub enum ConnectionRole {
Server,
}
#[async_trait::async_trait]
pub trait TunnelHandler {
async fn run(self) -> Result<(), String>;
}
pub struct Connection {
addr: SocketAddr,
pub inbound: OwnedReadHalf,
@@ -57,188 +62,218 @@ impl Connection {
}
}
async fn read_socks_request(&mut self) -> Result<SocksRequest, String> {
pub fn new_raw(inbound: OwnedReadHalf, outbound: OwnedWriteHalf) -> Self {
Self {
addr: "0.0.0.0:0".parse().unwrap(), // заглушка, если адрес не важен
inbound,
outbound,
read_buf: BytesMut::with_capacity(BUF_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)) => {
info!(client = %self.addr, request = ?req, "SOCKS request successfully parsed");
return Ok(req);
}
Ok(None) => {
trace!(client = %self.addr, buffer_len = self.read_buf.len(), "SOCKS parse: need more data");
}
Err(e) => {
error!(client = %self.addr, error = %e, "SOCKS protocol violation");
return Err(format!("Socks parse error: {}", e));
}
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| {
error!(client = %self.addr, error = %e, "Failed to read from socket during SOCKS handshake");
e.to_string()
})?;
.map_err(|e| e.to_string())?;
if n == 0 {
warn!(client = %self.addr, "Client closed connection prematurely during SOCKS handshake");
return Err("Client closed connection during SOCKS handshake".into());
return Err("Client closed connection".into());
}
trace!(client = %self.addr, read_bytes = n, "Read data from client for SOCKS handshake");
}
}
async fn send_socks_reply(&mut self, reply: SocksReply) -> Result<(), String> {
pub async fn send_socks_reply(&mut self, reply: SocksReply) -> Result<(), String> {
let mut buf = BytesMut::with_capacity(24);
debug!(client = %self.addr, reply = ?reply, "Sending SOCKS reply to client");
reply.write_to(&mut buf);
self.outbound.write_all(&buf).await.map_err(|e| {
error!(client = %self.addr, error = %e, "Failed to send SOCKS reply");
e.to_string()
})?;
Ok(())
self.outbound
.write_all(&buf)
.await
.map_err(|e| e.to_string())
}
}
pub async fn handle_socks_client(mut self, muxer: Muxer) -> Result<(), String> {
pub struct ClientHandler {
pub conn: Connection,
pub muxer: Muxer,
}
// Внутри connection.rs или где у тебя ClientHandler
impl ClientHandler {
pub async fn connect(
remote_proxy_addr: &str,
token: CancellationToken,
) -> 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);
// TLS Handshake
let ch = conn
.codec
.make_client_handshake(&BrowserProfile::CHROME_131, "google.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 (mux_tx, mux_rx) = mpsc::channel(BUF_SIZE);
let muxer = Muxer::new(mux_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,
mux_rx,
handler,
token: token.clone(),
};
tokio::spawn(async move { engine.run().await });
Ok(muxer)
}
}
#[async_trait::async_trait]
impl TunnelHandler for ClientHandler {
async fn run(mut self) -> Result<(), String> {
info!("Starting SOCKS multiplexed handling");
debug!("Reading SOCKS handshake request");
let _ = self.read_socks_request().await.map_err(|e| {
error!("SOCKS handshake failed: {}", e);
e
})?;
self.send_socks_reply(SocksReply::HandshakeSelect { method: 0x00 })
self.conn.read_socks_request().await?;
self.conn
.send_socks_reply(SocksReply::HandshakeSelect { method: 0x00 })
.await?;
let req = self.read_socks_request().await?;
let req = self.conn.read_socks_request().await?;
let target = if let SocksRequest::Connect { target, .. } = req {
target
} else {
return Err("Expected Connect".into());
};
let stream_id = muxer.next_id();
let target_str = target.to_string();
let stream_id = self.muxer.next_id();
let (v_tx, mut v_rx) = mpsc::channel::<bytes::Bytes>(1024);
self.muxer.register_stream(stream_id, v_tx).await;
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(1024);
muxer.register_stream(stream_id, v_tx).await;
muxer
.to_network
.send(MuxMessage {
self.muxer
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::Connect,
data: Bytes::from(target_str),
data: bytes::Bytes::from(target.to_string()),
})
.await
.map_err(|e| e.to_string())?;
let first_payload =
match tokio::time::timeout(std::time::Duration::from_secs(10), v_rx.recv()).await {
Ok(Some(data)) => data,
_ => {
error!(
stream_id,
"Server timeout or failed to send Connect confirmation"
);
self.send_socks_reply(SocksReply::ConnectResult {
reply_code: 0x01,
atyp: 0x01,
addr: [0, 0, 0, 0],
port: 0,
})
.await
.ok();
return Err("Target connection failed".into());
}
};
let first_payload = tokio::time::timeout(std::time::Duration::from_secs(10), v_rx.recv())
.await
.map_err(|_| "Timeout")?
.ok_or("No data")?;
if first_payload.len() >= 2 && first_payload[1] == 0x00 {
debug!(
stream_id,
"Server confirmed connection, forwarding SOCKS reply to browser"
);
self.outbound
self.conn
.outbound
.write_all(&first_payload)
.await
.map_err(|e| e.to_string())?;
} else {
self.outbound.write_all(&first_payload).await.ok();
return Err("Server rejected connection".into());
self.conn.outbound.write_all(&first_payload).await.ok();
return Err("Rejected".into());
}
let Self {
inbound: browser_in,
outbound: browser_out,
..
} = self;
let muxer_clone = muxer.clone();
let browser_in = self.conn.inbound;
let browser_out = self.conn.outbound;
let muxer = self.muxer;
tokio::spawn(async move {
run_proxy_bridge(stream_id, browser_in, browser_out, muxer_clone, v_rx).await;
run_proxy_bridge(stream_id, browser_in, browser_out, muxer, v_rx).await;
});
Ok(())
}
}
pub async fn handle_server_tunnel(mut self, token: CancellationToken) -> Result<(), String> {
info!("Acting as TLS Server, waiting for ClientHello");
pub struct ServerHandler {
pub conn: Connection,
pub token: CancellationToken,
}
#[async_trait::async_trait]
impl TunnelHandler for ServerHandler {
async fn run(mut self) -> Result<(), String> {
info!("Acting as TLS Server");
let (mux_tx, mux_rx) = mpsc::channel(BUF_SIZE);
let muxer = Muxer::new(mux_tx.clone(), false);
let muxer = Muxer::new(mux_tx, false);
let server_hello_bytes = loop {
match self.codec.make_server_handshake(&mut self.read_buf) {
Ok(bytes) => {
info!("ClientHello received, sending ServerHello");
break bytes;
}
let hello = loop {
match self
.conn
.codec
.make_server_handshake(&mut self.conn.read_buf)
{
Ok(b) => break b,
Err(e) if e.action == ErrorAction::Wait => {
let n = self
if self
.conn
.inbound
.read_buf(&mut self.read_buf)
.read_buf(&mut self.conn.read_buf)
.await
.map_err(|err| format!("Read error: {}", err))?;
if n == 0 {
return Err("Client closed connection".into());
.map_err(|e| e.to_string())?
== 0
{
return Err("Closed".into());
}
}
Err(e) => return Err(format!("TLS error: {:?}", e)),
Err(e) => return Err(format!("{:?}", e)),
}
};
self.outbound
.write_all(&server_hello_bytes)
self.conn
.outbound
.write_all(&hello)
.await
.map_err(|e| e.to_string())?;
info!("TLS Tunnel established as server");
let handler =
std::sync::Arc::new(StreamHandler::new(muxer.clone(), ConnectionRole::Server));
debug!("Handover to TunnelEngine");
let engine = TunnelEngine {
inbound: self.inbound,
outbound: self.outbound,
codec: self.codec,
read_buf: self.read_buf,
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,
mux_rx,
handler,
token,
};
engine.run().await.map_err(|e| {
error!("TunnelEngine error: {}", e);
e
})
token: self.token,
}
.run()
.await
}
}
core/src/proxy/connection/handler.rs
+1 -1
View File
@@ -35,7 +35,7 @@ impl StreamHandler {
let target_str = String::from_utf8_lossy(&payload).to_string();
let muxer = self.muxer.clone();
let (v_tx, v_rx) = tokio::sync::mpsc::channel(100);
let (v_tx, v_rx) = tokio::sync::mpsc::channel(1024);
muxer.register_stream(stream_id, v_tx).await;
tokio::spawn(async move {
core/src/proxy/connection/muxer.rs
+6 -2
View File
@@ -1,9 +1,9 @@
use crate::protocol::codec::frame::FrameType;
use bytes::Bytes;
use netrunner_logger::{debug, error};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc::error::SendError;
use tokio::sync::mpsc::Sender;
use tokio::sync::RwLock;
@@ -32,7 +32,7 @@ pub struct MuxMessage {
#[derive(Clone)]
pub struct Muxer {
pub to_network: Sender<MuxMessage>,
to_network: Sender<MuxMessage>,
streams: Arc<RwLock<HashMap<u32, Sender<Bytes>>>>,
id_gen: Arc<IdGenerator>,
}
@@ -50,6 +50,10 @@ 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 async fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
let mut lock = self.streams.write().await;
lock.insert(stream_id, tx);
core/src/proxy/network.rs
+29 -40
View File
@@ -1,14 +1,16 @@
use crate::{
protocol::errors::ErrorAction,
proxy::connection::{
connection::{Connection, ConnectionRole, BUF_SIZE},
connection::{
ClientHandler, Connection, ConnectionRole, ServerHandler, TunnelHandler, BUF_SIZE,
},
engine::TunnelEngine,
muxer::Muxer,
},
tlseng::profile::BrowserProfile,
};
use bytes::BytesMut;
use netrunner_logger::{error, info, instrument};
use netrunner_logger::{error, info};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
@@ -21,7 +23,6 @@ pub struct Network {
role: ConnectionRole,
remote_proxy_addr: Option<String>,
}
impl Network {
pub fn new(
host: String,
@@ -42,64 +43,52 @@ impl Network {
match self.role {
ConnectionRole::Client => {
info!("Starting Client mode: Initializing persistent tunnel to proxy...");
let muxer = match self.initialize_client_tunnel(token.clone()).await {
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, token.clone()).await {
Ok(m) => m,
Err(e) => {
error!(error = %e, "Global tunnel failed. Exit.");
error!(error = %e, "Global tunnel failed.");
return;
}
};
let listener = TcpListener::bind(&addr).await.expect("SOCKS bind failed");
info!(socks_addr = %addr, "SOCKS5 ready");
loop {
tokio::select! {
_ = token.cancelled() => {
info!("SOCKS listener: Shutting down...");
break;
}
_ = token.cancelled() => break,
res = listener.accept() => {
if let Ok((stream, client_addr)) = listener.accept().await {
let current_muxer = muxer.clone();
if let Ok((stream, client_addr)) = res {
let conn = Connection::new(stream, client_addr, false);
let handler = ClientHandler{ conn, muxer: muxer.clone() };
tokio::spawn(async move {
let connection = Connection::new(stream, client_addr, false);
let _ = connection.handle_socks_client(current_muxer).await;
if let Err(e) = handler.run().await {
error!(error = %e, "Client handler error");
}
});
}
}
}
}
}
ConnectionRole::Server => {
let listener = TcpListener::bind(&addr)
.await
.expect("Failed to bind Server port");
info!(listen_addr = %addr, "Proxy Server listening for incoming tunnels");
let listener = TcpListener::bind(&addr).await.expect("Server bind failed");
loop {
tokio::select! {
_ = token.cancelled() => {
info!("Server listener: Shutting down...");
break;
}
_ = token.cancelled() => break,
res = listener.accept() => {
match res {
Ok((stream, client_addr)) => {
let connection = Connection::new(stream, client_addr, true);
let connection_token = token.clone();
tokio::spawn(async move {
if let Err(e) = connection.handle_server_tunnel(connection_token).await {
error!(client = %client_addr, error = %e, "Tunnel error");
}
});
}
Err(e) => {
error!(error = %e, "Failed to accept connection");
}
if let Ok((stream, client_addr)) = res {
let conn = Connection::new(stream, client_addr, 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");
}
});
}
}
}