fixes

Cargo.lock

@@ -1089,6 +1089,7 @@ dependencies = [
  "smoltcp",
  "thiserror 2.0.18",
  "tokio",
+ "tokio-util",
  "tracing",
  "tracing-subscriber",
  "tun",

client/Cargo.toml

@@ -36,3 +36,4 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }
 thiserror = "2.0"
 hickory-proto = "0.25.2"
 lru = "0.16.3"
+tokio-util = "0.7.18"

client/src/connections/tcp_connection.rs

@@ -4,7 +4,8 @@ use smoltcp::socket::tcp;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpStream; // Твой код парсера
 use tokio::sync::mpsc;
-use tracing::{debug, info, trace};
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, info, trace, warn};
 
 pub enum ConnectionState {
     Established,
@@ -18,25 +19,21 @@ pub struct TcpConnection {
     state: ConnectionState,
     tx: mpsc::UnboundedSender<Vec<u8>>,
     rx: mpsc::UnboundedReceiver<Vec<u8>>,
-    pending_data: Option<Vec<u8>>,
+    pending_data: Vec<u8>,
+    token: CancellationToken,
 }
 
+const MAX_PENDING: usize = 256 * 1024;
+
 impl TcpConnection {
     pub fn new(handle: SocketHandle, proxy_addr: String, target_addr: TargetAddress) -> Self {
         let (tx_to_proxy, mut rx_from_smol) = mpsc::unbounded_channel::<Vec<u8>>();
         let (tx_to_smol, rx_from_proxy) = mpsc::unbounded_channel::<Vec<u8>>();
 
-        let proxy_addr_clone = proxy_addr.clone();
+        let token = CancellationToken::new();
-        let target_addr_clone = target_addr.clone();
+        let task_token = token.clone();
 
         tokio::spawn(async move {
-            debug!(
-                %handle,
-                target = ?target_addr_clone,
-                proxy = %proxy_addr_clone,
-                "Attempting to connect through proxy"
-            );
-
             let mut stream = match TcpStream::connect(&proxy_addr).await {
                 Ok(s) => {
                     debug!(%handle, "Connected to proxy successfully");
@@ -49,8 +46,7 @@ impl TcpConnection {
             };
 
             // 2. SOCKS Handshake
-            if let Err(e) =
+            if let Err(e) = SocksRequest::perform_client_handshake(&mut stream, &target_addr).await
-                SocksRequest::perform_client_handshake(&mut stream, &target_addr_clone).await
             {
                 debug!(%handle, error = %e, "SOCKS handshake failed");
                 return;
@@ -72,7 +68,7 @@ impl TcpConnection {
 
             // Читаем из прокси -> Пишем в канал для smoltcp
             let from_proxy = async {
-                let mut buf = [0u8; 4096];
+                let mut buf = [0u8; 65536];
                 loop {
                     match reader.read(&mut buf).await {
                         Ok(0) | Err(_) => break,
@@ -86,53 +82,60 @@ impl TcpConnection {
             };
 
             tokio::select! {
-                _ = to_proxy => {},
+                _ = to_proxy => {}
-                _ = from_proxy => {},
+                _ = from_proxy => {}
+                _ = task_token.cancelled() => { debug!(%handle, "Task cancelled by Manager"); }
             }
         });
 
         Self {
             handle,
-            state: ConnectionState::Handshaking, // Сразу переходим в Active, т.к. задача пошла
+            state: ConnectionState::Active, // Сразу переходим в Active, т.к. задача пошла
             tx: tx_to_proxy,
             rx: rx_from_proxy,
-            pending_data: None,
+            pending_data: vec![],
+            token,
         }
     }
-
     pub fn tick(&mut self, socket: &mut tcp::Socket) -> bool {
-        trace!(handle=%self.handle, state=?socket.state(), "Tick");
+        let state = socket.state();
+        //trace!(handle=%self.handle, ?state, "Tick");
+
         match self.state {
             ConnectionState::Handshaking => {
-                info!("Connection handshaking");
                 return true;
             }
+
             ConnectionState::Active => {
                 self.poll_and_process(socket);
 
-                // Проверяем условия закрытия
+                // FIN от удалённой стороны
+                if state == tcp::State::CloseWait {
+                    socket.close();
+                    self.state = ConnectionState::Closed;
+                    return false;
+                }
+
                 if self.is_finished(socket) {
                     self.state = ConnectionState::Closed;
-                    socket.abort(); // Принудительно разрываем стек
+                    socket.close();
-                    return false; // Сигнализируем, что соединение мертво
+                    return false;
                 }
             }
-            ConnectionState::Closed => return false,
+
+            ConnectionState::Closed => {
+                return false;
+            }
+
             _ => {}
         }
+
         true
     }
-
     pub fn is_finished(&self, socket: &tcp::Socket) -> bool {
         use tcp::State;
-        let socket_closed = matches!(
-            socket.state(),
-            State::Closed | State::TimeWait | State::FinWait1 | State::FinWait2
-        );
 
-        let task_finished = self.rx.is_closed();
+        matches!(socket.state(), State::Closed | State::TimeWait)
-
-        socket_closed || task_finished
     }
 
     pub fn is_active(&self) -> bool {
@@ -140,14 +143,7 @@ impl TcpConnection {
     }
 
     fn poll_and_process(&mut self, socket: &mut tcp::Socket) {
-        if let Some(data) = self.pending_data.take() {
+        // 1. сначала читаем smoltcp
-            if socket.send_slice(&data).is_ok() {
-            } else {
-                self.pending_data = Some(data);
-                return;
-            }
-        }
-
         if socket.can_recv() {
             let _ = socket.recv(|data| {
                 let len = data.len();
@@ -158,22 +154,37 @@ impl TcpConnection {
             });
         }
 
+        // 2. потом отправляем proxy → smoltcp
+        if !self.pending_data.is_empty() {
+            if self.pending_data.len() > MAX_PENDING {
+                warn!(%self.handle, "Buffer overflow! Aborting connection.");
+                socket.abort(); // Убиваем сокет
+                self.token.cancel(); // Говорим tokio-задаче умереть
+                return;
+            }
+
+            match socket.send_slice(&self.pending_data) {
+                Ok(n) => {
+                    self.pending_data.drain(0..n);
+                }
+                Err(_) => {} // Оставляем в pending_data на следующий тик
+            }
+        }
+
         if socket.can_send() {
             while let Ok(data) = self.rx.try_recv() {
                 match socket.send_slice(&data) {
-                    Ok(_) => { /* Успешно */ }
+                    Ok(n) if n < data.len() => {
+                        self.pending_data = data[n..].to_vec();
+                        break;
+                    }
+                    Ok(_) => {}
                     Err(_) => {
-                        // Стек полон, запоминаем данные для следующего вызова
+                        self.pending_data = data;
-                        self.pending_data = Some(data);
                         break;
                     }
                 }
             }
         }
-
-        if socket.state() == tcp::State::CloseWait {
-            debug!(handle=%self.handle, "!!! Triggering socket.close() for CLOSE-WAIT");
-            socket.close();
-        }
     }
 }

client/src/tun/connection_manager.rs

@@ -112,33 +112,52 @@ impl ConnectionManager {
     }
 
     fn handle_tcp(&mut self, handle: SocketHandle, socket: &mut tcp::Socket) {
-        if socket.state() == tcp::State::Established {
+        use tcp::State;
-            let target = self.resolve_target(socket);
-            debug!(
-                handle=%handle,
-                local=?socket.local_endpoint(),
-                remote=?socket.remote_endpoint(),
-                "Socket details"
-            );
 
-            let mut conn = TcpConnection::new(handle, self.proxy_ip.clone(), target);
+        // 1. Если сокет закрыт, просто чистим и возвращаем в LISTEN
-
+        if socket.state() == State::Closed {
-            conn.tick(socket);
+            self.active_tcp_sessions.remove(&handle);
-
+            socket.abort();
-            if conn.is_active() {
+            let _ = socket.listen(443);
-                self.active_tcp_sessions.insert(handle, conn);
+            return;
-            } else {
-                socket.abort();
-            }
         }
 
+        // 2. Если сокет установлен, но в менеджере нет записи
+        if socket.state() == State::Established && !self.active_tcp_sessions.contains_key(&handle) {
+            let target = self.resolve_target(socket);
+
+            // ВАЖНО: Тут можно добавить проверку: если target "плохой" или не резолвится,
+            // сразу убиваем сокет, чтобы не зацикливаться.
+            if let TargetAddress::Domain(d, _) = &target {
+                if d == "disconnected" {
+                    // Или другая логика проверки
+                    socket.abort();
+                    return;
+                }
+            }
+
+            let conn = TcpConnection::new(handle, self.proxy_ip.clone(), target);
+            self.active_tcp_sessions.insert(handle, conn);
+        }
+
+        // 3. Обработка активной сессии
         if let Some(conn) = self.active_tcp_sessions.get_mut(&handle) {
             if !conn.tick(socket) {
+                // Если tick вернул false, значит сессия завершена или произошла ошибка в tokio-задаче
                 self.active_tcp_sessions.remove(&handle);
+                socket.abort(); // Принудительно закрываем "битый" сокет
             }
+        } else if socket.state() == State::Established {
+            // Если мы дошли сюда, значит сессия была, но удалилась (tick вернул false),
+            // а сокет всё еще висит в Established. Убиваем его.
+            socket.abort();
+        }
+
+        // 4. Обработка FIN
+        if socket.state() == State::CloseWait {
+            socket.close();
         }
     }
-
     fn handle_udp(&mut self, handle: SocketHandle, socket: &mut udp::Socket) {
         self.last_activity.insert(handle, StdInstant::now());
 
@@ -201,4 +220,25 @@ impl ConnectionManager {
 
         sockets
     }
+
+    pub fn log_status(&self, socket_set: &SocketSet) {
+        let mut established = 0;
+        let mut total_tcp = 0;
+
+        for (_, socket) in socket_set.iter() {
+            if let Some(tcp) = tcp::Socket::downcast(socket) {
+                total_tcp += 1;
+                if tcp.state() == tcp::State::Established {
+                    established += 1;
+                }
+            }
+        }
+
+        debug!(
+            "TCP Stats: Total_Sockets={}, Established={}, Active_Sessions={}",
+            total_tcp,
+            established,
+            self.active_tcp_sessions.len()
+        );
+    }
 }

client/src/tun/engine.rs

@@ -1,9 +1,11 @@
+use smoltcp::iface::PollResult;
 use smoltcp::time::Instant;
 use smoltcp::wire::{IpAddress, IpCidr};
 use smoltcp::{
     iface::{Config, Interface, SocketSet},
     phy::DeviceCapabilities,
 };
+use std::sync::atomic::Ordering;
 use std::{
     mem,
     sync::{Arc, LazyLock, atomic::AtomicBool},
@@ -37,7 +39,7 @@ impl Engine {
         let (mut device, bridge_rx, bridge_tx, avail) = VirtTunDevice::new(caps);
         let interface = Interface::new(config, &mut device, now);
 
-        let socket_set = ConnectionManager::setup_sockets(16, 4, 1);
+        let socket_set = ConnectionManager::setup_sockets(128, 8, 4);
         let manager = ConnectionManager::new(ip);
         Self {
             interface,
@@ -52,32 +54,41 @@ impl Engine {
 
     pub async fn run(&mut self, tun: Tun) {
         info!("Current routes: {:?}", self.interface.routes());
-        //Bridge from tun to stack and stack to tun
+
         let (writer, reader) = tun.split().expect("Failed to split TUN");
-        // Забираем bridge_rx, так как он нам нужен только в одной задаче
+
         let from_engine = mem::replace(&mut self.bridge_rx, mpsc::unbounded_channel().1);
+
         Self::spawn_tun_to_engine(reader, self.bridge_tx.clone(), self.avail.clone());
         Self::spawn_engine_to_tun(writer, from_engine);
-
+        let mut last_log = StdInstant::now();
         loop {
-            loop {
+            let result = self.poll();
-                self.poll();
 
-                self.manager.process_sockets(&mut self.socket_set);
+            self.manager.process_sockets(&mut self.socket_set);
 
-                if self.avail.load(std::sync::atomic::Ordering::Acquire) {
+            if last_log.elapsed() >= Duration::from_secs(5) {
-                    tokio::task::yield_now().await;
+                self.manager.log_status(&self.socket_set);
-                    continue;
+                last_log = StdInstant::now();
-                }
-                self.poll_delay().await;
             }
+
+            if matches!(result, PollResult::SocketStateChanged) {
+                continue;
+            }
+
+            if self.avail.load(Ordering::Acquire) {
+                tokio::task::yield_now().await;
+                continue;
+            }
+
+            self.poll_delay().await;
         }
     }
 
-    fn poll(&mut self) {
+    fn poll(&mut self) -> PollResult {
         let now = Self::current_time();
         self.interface
-            .poll(now, &mut self.device, &mut self.socket_set);
+            .poll(now, &mut self.device, &mut self.socket_set)
     }
 
     async fn poll_delay(&mut self) {
@@ -97,7 +108,7 @@ impl Engine {
     ) {
         tokio::spawn(async move {
             debug!("TUN-to-Engine bridge task started");
-            let mut buf = [0u8; 4096];
+            let mut buf = [0u8; 65536];
             while let Ok(n) = reader.read(&mut buf).await {
                 if n == 0 {
                     break;

common/src/proxy/connection/bridge.rs

@@ -66,6 +66,6 @@ pub async fn run_proxy_bridge<R, W>(
             data: Bytes::new(),
         })
         .await;
-
+    tokio::time::sleep(std::time::Duration::from_millis(500)).await;
     muxer.remove_stream(stream_id).await;
 }

common/src/proxy/connection/engine.rs

@@ -6,7 +6,7 @@ use tokio::{
     net::tcp::{OwnedReadHalf, OwnedWriteHalf},
     sync::mpsc::Receiver,
 };
-use tracing::error;
+use tracing::{debug, error};
 
 use crate::{
     protocol::{codec::codec::Codec, errors::ErrorAction},
@@ -75,6 +75,7 @@ impl TunnelEngine {
                         break;
                     }
                     // Иначе — это реальная проблема (кривой TLS и т.д.)
+                    error!(error = ?e, "Codec inbound failed");
                     return Err(format!("Codec error: {:?}", e));
                 }
             }
@@ -98,6 +99,7 @@ impl TunnelEngine {
                         error!(stream_id = msg.stream_id, error = %e, "Failed to write encrypted data to network");
                         e.to_string()
                     })?;
+                debug!(stream_id = msg.stream_id, "Outbound packet sent");
                 Ok(())
             }
             Err(e) => {

common/src/proxy/connection/handler.rs

@@ -93,6 +93,5 @@ impl StreamHandler {
 
     async fn on_close(&self, stream_id: u32) {
         self.muxer.dispatch_to_local(stream_id, Bytes::new()).await;
-        self.muxer.remove_stream(stream_id).await;
     }
 }

common/src/proxy/connection/muxer.rs

@@ -5,7 +5,7 @@ use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::Arc;
 use tokio::sync::mpsc::Sender;
 use tokio::sync::RwLock;
-use tracing::error;
+use tracing::{debug, error};
 
 pub struct IdGenerator {
     counter: AtomicU32,
@@ -53,15 +53,16 @@ impl Muxer {
     pub async fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
         let mut lock = self.streams.write().await;
         lock.insert(stream_id, tx);
-        // ДОБАВЬ ЭТО:
         tracing::debug!(
             stream_id,
             total_active = lock.len(),
-            "STREAMS_MAP_UPDATE: Registered new stream"
+            "MUXER: [REGISTER] Stream added"
         );
     }
+
     pub async fn remove_stream(&self, stream_id: u32) {
-        self.streams.write().await.remove(&stream_id);
+        let mut lock = self.streams.write().await;
+        lock.remove(&stream_id); // Просто удаляем, если есть
     }
 
     pub async fn send_control(
@@ -80,17 +81,35 @@ impl Muxer {
             .map_err(|e| e.to_string())
     }
 
-    /// Отправляет входящие данные конкретному локальному обработчику
     pub async fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
-        // Асинхронный лок сам умеет корректно работать с .await
+        let tx = {
-        let tx = self.streams.read().await.get(&stream_id).cloned();
+            let lock = self.streams.read().await;
+            lock.get(&stream_id).cloned()
+        };
 
         if let Some(tx) = tx {
-            if tx.send(data).await.is_err() {
+            if data.is_empty() {
+                tracing::debug!(stream_id, "MUXER: [EOF] Forwarding EOF to local handler");
+            } else {
+                tracing::trace!(
+                    stream_id,
+                    len = data.len(),
+                    "MUXER: [DISPATCH] Sending data"
+                );
+            }
+
+            if let Err(_e) = tx.send(data).await {
+                tracing::debug!(
+                    stream_id,
+                    "MUXER: [WARN] Local channel closed, dropping packet"
+                );
                 self.remove_stream(stream_id).await;
             }
         } else {
-            error!(stream_id, "MUXER: Received data for UNKNOWN stream_id");
+            tracing::trace!(
+                stream_id,
+                "MUXER: [IGNORE] Packet for already closed stream"
+            );
         }
     }
 }