connection refactoring

Cargo.lock

@@ -1217,6 +1217,7 @@ name = "netrunner-core"
 version = "0.1.0"
 dependencies = [
  "aead",
+ "async-trait",
  "bytes",
  "chacha20poly1305",
  "hex",

core/Cargo.toml

@@ -16,4 +16,5 @@ tokio = "1.49.0"
 rand = "0.10.0"
 hex = "0.4.3"
 tokio-util = { version = "0.7", features = ["rt"] }
+async-trait = "0.1.89"
 netrunner-logger = { path = "../tools/log" }

core/src/proxy/connection/bridge.rs

@@ -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_to_netwrok(MuxMessage {
-        .send(MuxMessage {
             stream_id,
             frame_type: FrameType::Close,
             data: Bytes::new(),

core/src/proxy/connection/connection.rs

@@ -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 {
-        loop {
+        Self {
-            match SocksRequest::parse(&mut self.read_buf) {
+            addr: "0.0.0.0:0".parse().unwrap(), // заглушка, если адрес не важен
-                Ok(Some(req)) => {
+            inbound,
-                    info!(client = %self.addr, request = ?req, "SOCKS request successfully parsed");
+            outbound,
-                    return Ok(req);
+            read_buf: BytesMut::with_capacity(BUF_SIZE),
-                }
+            codec: Codec::new(false),
-                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));
         }
     }
 
+    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| {
+                .map_err(|e| e.to_string())?;
-                    error!(client = %self.addr, error = %e, "Failed to read from socket during SOCKS handshake");
-                    e.to_string()
-                })?;
-
             if n == 0 {
-                warn!(client = %self.addr, "Client closed connection prematurely during SOCKS handshake");
+                return Err("Client closed connection".into());
-                return Err("Client closed connection during SOCKS handshake".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
-        self.outbound.write_all(&buf).await.map_err(|e| {
+            .write_all(&buf)
-            error!(client = %self.addr, error = %e, "Failed to send SOCKS reply");
+            .await
-            e.to_string()
+            .map_err(|e| e.to_string())
-        })?;
+    }
-
-        Ok(())
 }
 
-    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");
-
+        self.conn.read_socks_request().await?;
-        debug!("Reading SOCKS handshake request");
+        self.conn
-        let _ = self.read_socks_request().await.map_err(|e| {
+            .send_socks_reply(SocksReply::HandshakeSelect { method: 0x00 })
-            error!("SOCKS handshake failed: {}", e);
-            e
-        })?;
-
-        self.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 stream_id = self.muxer.next_id();
-        let target_str = target.to_string();
+        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);
+        self.muxer
-        muxer.register_stream(stream_id, v_tx).await;
+            .send_to_netwrok(MuxMessage {
-
-        muxer
-            .to_network
-            .send(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 =
+        let first_payload = tokio::time::timeout(std::time::Duration::from_secs(10), v_rx.recv())
-            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();
+            .map_err(|_| "Timeout")?
-                    return Err("Target connection failed".into());
+            .ok_or("No data")?;
-                }
-            };
 
         if first_payload.len() >= 2 && first_payload[1] == 0x00 {
-            debug!(
+            self.conn
-                stream_id,
+                .outbound
-                "Server confirmed connection, forwarding SOCKS reply to browser"
-            );
-
-            self.outbound
                 .write_all(&first_payload)
                 .await
                 .map_err(|e| e.to_string())?;
         } else {
-            self.outbound.write_all(&first_payload).await.ok();
+            self.conn.outbound.write_all(&first_payload).await.ok();
-            return Err("Server rejected connection".into());
+            return Err("Rejected".into());
         }
 
-        let Self {
+        let browser_in = self.conn.inbound;
-            inbound: browser_in,
+        let browser_out = self.conn.outbound;
-            outbound: browser_out,
+        let muxer = self.muxer;
-            ..
-        } = self;
-
-        let muxer_clone = muxer.clone();
         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> {
+pub struct ServerHandler {
-        info!("Acting as TLS Server, waiting for ClientHello");
+    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 {
+        let hello = loop {
-            match self.codec.make_server_handshake(&mut self.read_buf) {
+            match self
-                Ok(bytes) => {
+                .conn
-                    info!("ClientHello received, sending ServerHello");
+                .codec
-                    break bytes;
+                .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))?;
+                        .map_err(|e| e.to_string())?
-
+                        == 0
-                    if n == 0 {
+                    {
-                        return Err("Client closed connection".into());
+                        return Err("Closed".into());
                     }
                 }
-                Err(e) => return Err(format!("TLS error: {:?}", e)),
+                Err(e) => return Err(format!("{:?}", e)),
             }
         };
 
-        self.outbound
+        self.conn
-            .write_all(&server_hello_bytes)
+            .outbound
+            .write_all(&hello)
             .await
             .map_err(|e| e.to_string())?;
-        info!("TLS Tunnel established as server");
 
-        let handler =
+        let handler = std::sync::Arc::new(StreamHandler::new(muxer, ConnectionRole::Server));
-            std::sync::Arc::new(StreamHandler::new(muxer.clone(), ConnectionRole::Server));
+        TunnelEngine {
-
+            inbound: self.conn.inbound,
-        debug!("Handover to TunnelEngine");
+            outbound: self.conn.outbound,
-        let engine = TunnelEngine {
+            codec: self.conn.codec,
-            inbound: self.inbound,
+            read_buf: self.conn.read_buf,
-            outbound: self.outbound,
-            codec: self.codec,
-            read_buf: self.read_buf,
             mux_rx,
             handler,
-            token,
+            token: self.token,
-        };
+        }
-
+        .run()
-        engine.run().await.map_err(|e| {
+        .await
-            error!("TunnelEngine error: {}", e);
-            e
-        })
     }
 }

core/src/proxy/connection/handler.rs

@@ -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

@@ -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

@@ -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,65 +43,53 @@ impl Network {
 
         match self.role {
             ConnectionRole::Client => {
-                info!("Starting Client mode: Initializing persistent tunnel to proxy...");
+                info!("Starting Client mode");
-
+                let server_addr = self
-                let muxer = match self.initialize_client_tunnel(token.clone()).await {
+                    .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() => {
+                        _ = token.cancelled() => break,
-                            info!("SOCKS listener: Shutting down...");
-                            break;
-                        }
                         res = listener.accept() => {
-                            if let Ok((stream, client_addr)) = listener.accept().await {
+                            if let Ok((stream, client_addr)) = res {
-                                let current_muxer = muxer.clone();
+                                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);
+                                    if let Err(e) = handler.run().await {
-                                    let _ = connection.handle_socks_client(current_muxer).await;
+                                        error!(error = %e, "Client handler error");
+                                    }
                                 });
                             }
                         }
                     }
                 }
             }
-
             ConnectionRole::Server => {
-                let listener = TcpListener::bind(&addr)
+                let listener = TcpListener::bind(&addr).await.expect("Server bind failed");
-                    .await
-                    .expect("Failed to bind Server port");
-                info!(listen_addr = %addr, "Proxy Server listening for incoming tunnels");
-
                 loop {
                     tokio::select! {
-                        _ = token.cancelled() => {
+                        _ = token.cancelled() => break,
-                            info!("Server listener: Shutting down...");
-                            break;
-                        }
                         res = listener.accept() => {
-                            match res {
+                            if let Ok((stream, client_addr)) = res {
-                                Ok((stream, client_addr)) => {
+                                let conn = Connection::new(stream, client_addr, true);
-                                    let connection = Connection::new(stream, client_addr, true);
+                                let handler = ServerHandler { conn, token: token.clone() };
-                                    let connection_token = token.clone();
                                 tokio::spawn(async move {
-                                        if let Err(e) = connection.handle_server_tunnel(connection_token).await {
+                                    if let Err(e) = handler.run().await {
-                                            error!(client = %client_addr, error = %e, "Tunnel error");
+                                        error!(client = %client_addr, error = %e, "Server handler error");
                                     }
                                 });
                             }
-                                Err(e) => {
-                                    error!(error = %e, "Failed to accept connection");
-                                }
-                            }
                         }
                     }
                 }