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

fixes

Browse Source
This commit is contained in:
Kirill
2026-03-08 17:07:15 +07:00
parent 676683d29c
commit 2069f1a62c
9 changed files with 179 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/src/connections/tcp_connection.rs
+60 -49
View File
@@ -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 target_addr_clone = target_addr.clone();
let token = CancellationToken::new();
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) =
SocksRequest::perform_client_handshake(&mut stream, &target_addr_clone).await
if let Err(e) = SocksRequest::perform_client_handshake(&mut stream, &target_addr).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 => {},
_ = from_proxy => {},
_ = to_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(); // Принудительно разрываем стек
return false; // Сигнализируем, что соединение мертво
socket.close();
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();
socket_closed || task_finished
matches!(socket.state(), State::Closed | State::TimeWait)
}
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() {
if socket.send_slice(&data).is_ok() {
} else {
self.pending_data = Some(data);
return;
}
}
// 1. сначала читаем smoltcp
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 = Some(data);
self.pending_data = 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
+58 -18
View File
@@ -112,33 +112,52 @@ impl ConnectionManager {
}
fn handle_tcp(&mut self, handle: SocketHandle, socket: &mut tcp::Socket) {
if socket.state() == tcp::State::Established {
let target = self.resolve_target(socket);
debug!(
handle=%handle,
local=?socket.local_endpoint(),
remote=?socket.remote_endpoint(),
"Socket details"
);
use tcp::State;
let mut conn = TcpConnection::new(handle, self.proxy_ip.clone(), target);
conn.tick(socket);
if conn.is_active() {
self.active_tcp_sessions.insert(handle, conn);
} else {
socket.abort();
}
// 1. Если сокет закрыт, просто чистим и возвращаем в LISTEN
if socket.state() == State::Closed {
self.active_tcp_sessions.remove(&handle);
socket.abort();
let _ = socket.listen(443);
return;
}
// 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
+26 -15
View File
@@ -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 {
self.poll();
let result = 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) {
tokio::task::yield_now().await;
continue;
}
self.poll_delay().await;
if last_log.elapsed() >= Duration::from_secs(5) {
self.manager.log_status(&self.socket_set);
last_log = StdInstant::now();
}
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;