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

buffers update, warning fixes, ghosts killed

Browse Source
This commit is contained in:
Kirill
2026-04-03 18:18:10 +07:00
parent e53a95692e
commit 780750171b
29 changed files with 1205 additions and 1044 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/src/net/connection.rs
+155 -5
View File
@@ -1,14 +1,20 @@
use bytes::{Buf, Bytes, BytesMut};
use netrunner_core::net::network::NetworkConfig;
use netrunner_core::{
net::NetworkConfig,
rawcast::{LocalProtocol, RawCastFrame},
};
use smoltcp::{
iface::SocketHandle,
socket::{tcp, udp},
wire::IpEndpoint,
wire::{
Icmpv4Message, Icmpv4Packet, Icmpv6Message, Icmpv6Packet, IpAddress, IpEndpoint,
Ipv6Address,
},
};
use std::time::Duration;
use tokio::sync::{mpsc, oneshot};
use netrunner_logger::debug;
use netrunner_logger::{debug, info};
pub struct ConnectionCore<T> {
pub handle: SocketHandle,
@@ -83,7 +89,7 @@ impl TcpConnection {
match rx.try_recv() {
Ok(_) => {
debug!(%self.core.handle, "TCP Handshake successful, State -> Active");
self.state = ConnectionState::Active;
self.state = ConnectionState::Established;
self.handshake_rx = None;
return true;
}
@@ -112,7 +118,14 @@ impl TcpConnection {
return false;
}
_ => {}
ConnectionState::Established => {
info!(
"✅ [TCP {}] Connection fully established and ready for data",
self.core.handle
);
self.state = ConnectionState::Active;
return true;
}
}
true
@@ -185,6 +198,50 @@ impl TcpConnection {
socket.close();
}
}
pub fn spawn(
socket_id: u64,
dst_ip: std::net::Ipv4Addr,
dst_port: u16,
target: String,
mut rx_smol: mpsc::Receiver<Bytes>,
handshake_tx: oneshot::Sender<()>,
tx_tunnel: mpsc::Sender<RawCastFrame>,
) {
tokio::spawn(async move {
// 1. Формируем и отправляем кадр на установку соединения
let mut frame = RawCastFrame::connect(LocalProtocol::Tcp, socket_id, dst_ip, dst_port);
frame.payload = Bytes::from(target);
if tx_tunnel.send(frame).await.is_err() {
netrunner_logger::error!("❌ [TCP {}] Failed to send CONNECT to tunnel", socket_id);
return;
}
// 2. Сигнализируем ConnectionManager, что запрос в туннель ушел успешно
let _ = handshake_tx.send(());
// 3. Цикл пересылки данных из smoltcp -> туннель
while let Some(data) = rx_smol.recv().await {
let data_frame = RawCastFrame::data(
LocalProtocol::Tcp,
socket_id,
dst_ip,
dst_port,
data.to_vec(),
);
if tx_tunnel.send(data_frame).await.is_err() {
break; // Туннель закрыт
}
}
let close_frame = RawCastFrame::close(LocalProtocol::Tcp, socket_id, dst_ip, dst_port);
let _ = tx_tunnel.send(close_frame).await;
debug!("🏁 [TCP {}] Spawned task finished", socket_id);
});
}
}
const UDP_TIMEOUT: Duration = Duration::from_secs(60);
@@ -259,4 +316,97 @@ impl UdpConnection {
true
}
pub fn spawn(
socket_id: u64,
dst_ip: std::net::Ipv4Addr,
dst_port: u16,
target: String,
mut rx_smol: mpsc::Receiver<UdpPacketTarget>,
tx_tunnel: mpsc::Sender<RawCastFrame>,
) {
tokio::spawn(async move {
debug!("📡 [UDP {}] Task started for {}", socket_id, target);
// 1. Регистрация UDP сессии в туннеле
let mut frame = RawCastFrame::connect(LocalProtocol::Udp, socket_id, dst_ip, dst_port);
frame.payload = Bytes::from(target);
if tx_tunnel.send(frame).await.is_err() {
netrunner_logger::error!("❌ [UDP {}] Failed to send CONNECT to tunnel", socket_id);
return;
}
// 2. Цикл пересылки пакетов
// rx_smol отдает кортеж (данные, ip, порт)
while let Some((data, ip, port)) = rx_smol.recv().await {
let data_frame =
RawCastFrame::data(LocalProtocol::Udp, socket_id, ip, port, data.to_vec());
if tx_tunnel.send(data_frame).await.is_err() {
break;
}
}
let close_frame = RawCastFrame::close(LocalProtocol::Udp, socket_id, dst_ip, dst_port);
let _ = tx_tunnel.send(close_frame).await;
info!("🛑 [UDP {}] Task stopped", socket_id);
});
}
}
use smoltcp::socket::icmp;
pub struct IcmpResponder;
impl IcmpResponder {
pub fn handle(socket: &mut icmp::Socket) {
if !socket.can_recv() {
return;
}
// 1. Сначала достаем данные и адрес
let result = socket.recv();
if let Ok((data, src_addr)) = result {
// 2. Копируем данные в Vec, чтобы разорвать связь с буфером сокета.
// Теперь заимствование `socket` от метода .recv() закончилось.
let payload = data.to_vec();
match src_addr {
IpAddress::Ipv4(_) => Self::reply_v4(socket, payload, src_addr),
IpAddress::Ipv6(v6) => Self::reply_v6(socket, payload, v6),
}
}
}
// Принимаем Vec<u8>, чтобы не делать to_vec() второй раз внутри
fn reply_v4(socket: &mut icmp::Socket, mut payload: Vec<u8>, src: IpAddress) {
if let Ok(pkt) = Icmpv4Packet::new_checked(&payload) {
if pkt.msg_type() == Icmpv4Message::EchoRequest {
let mut reply_pkt = Icmpv4Packet::new_unchecked(&mut payload);
reply_pkt.set_msg_type(Icmpv4Message::EchoReply);
reply_pkt.fill_checksum();
let _ = socket.send_slice(&payload, src);
info!("🏓 [ICMPv4] Echo Reply -> {}", src);
}
}
}
fn reply_v6(socket: &mut icmp::Socket, mut payload: Vec<u8>, src: Ipv6Address) {
if let Ok(pkt) = Icmpv6Packet::new_checked(&payload) {
if pkt.msg_type() == Icmpv6Message::EchoRequest {
let mut reply_pkt = Icmpv6Packet::new_unchecked(&mut payload);
reply_pkt.set_msg_type(Icmpv6Message::EchoReply);
let gateway = Ipv6Address::new(0xfe80, 0, 0, 0, 0, 0, 0, 1);
reply_pkt.fill_checksum(&gateway, &src);
let _ = socket.send_slice(&payload, src.into());
info!("🏓 [ICMPv6] Echo Reply -> {}", src);
}
}
}
}
client/src/net/connection_manager.rs
+150 -511
View File
@@ -1,106 +1,26 @@
use bytes::Bytes;
use netrunner_core::{
net::network::NetworkConfig,
rawcast::{LocalProtocol, RawCastEvent, RawCastFrame},
};
use netrunner_logger::{debug, error, info, trace, warn};
use netrunner_core::rawcast::{RawCastEvent, RawCastFrame};
use netrunner_logger::{info, trace};
use smoltcp::{
iface::{SocketHandle, SocketSet},
socket::{AnySocket, icmp, tcp, udp},
wire::{
IpListenEndpoint, IpProtocol, Ipv4Packet, Ipv6Address, Ipv6Packet, TcpPacket, UdpPacket,
},
socket::{Socket, tcp, udp},
wire::{IpAddress, IpListenEndpoint, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket},
};
use std::{collections::HashMap, time::Instant as StdInstant};
use std::{sync::Arc, time::Duration};
use tokio::sync::mpsc;
use crate::net::{
connection::{TcpConnection, UdpConnection},
dns::DnsHandler,
ip_store::FakeIpStore,
connection::{IcmpResponder, TcpConnection, UdpConnection},
dns::{DnsHandler, FakeIpStore},
session_tracker::SessionTracker,
socket_factory::SocketProvider,
};
struct SessionTracker {
last_activity: HashMap<SocketHandle, StdInstant>,
active_tcp: HashMap<SocketHandle, TcpConnection>,
active_udp: HashMap<SocketHandle, UdpConnection>,
inbound_tx: HashMap<u64, mpsc::Sender<Bytes>>,
failed_until: HashMap<SocketHandle, StdInstant>,
to_remove: Vec<SocketHandle>,
next_socket_id: u64,
handle_to_id: HashMap<SocketHandle, u64>,
pending_tcp: HashMap<SocketHandle, StdInstant>,
}
impl SessionTracker {
fn new() -> Self {
Self {
last_activity: HashMap::new(),
active_tcp: HashMap::new(),
active_udp: HashMap::new(),
inbound_tx: HashMap::new(),
failed_until: HashMap::new(),
to_remove: Vec::new(),
next_socket_id: 1,
handle_to_id: HashMap::new(),
pending_tcp: HashMap::new(),
}
}
fn generate_socket_id(&mut self) -> u64 {
let id = self.next_socket_id;
self.next_socket_id = self.next_socket_id.wrapping_add(1);
id
}
fn queue_removal(&mut self, handle: SocketHandle) {
if !self.to_remove.contains(&handle) {
self.to_remove.push(handle);
}
}
fn cleanup(&mut self, socket_set: &mut SocketSet) {
for handle in self.to_remove.drain(..) {
debug!(%handle, "Cleanup: Removing socket");
socket_set.remove(handle);
self.last_activity.remove(&handle);
self.failed_until.remove(&handle);
self.active_tcp.remove(&handle);
self.active_udp.remove(&handle);
self.pending_tcp.remove(&handle);
if let Some(socket_id) = self.handle_to_id.remove(&handle) {
self.inbound_tx.remove(&socket_id);
}
}
}
fn has_connection_from(
&self,
src_addr: smoltcp::wire::IpAddress,
src_port: u16,
socket_set: &SocketSet,
) -> bool {
socket_set.iter().any(|(_, s)| {
if let Some(tcp) = tcp::Socket::downcast(s) {
if let Some(remote) = tcp.remote_endpoint() {
if remote.addr == src_addr && remote.port == src_port {
return true;
}
}
}
false
})
}
fn get_id_by_port(&self, port: u16) -> Option<u64> {
self.active_udp.iter().find_map(|(handle, conn)| {
if conn.has_client(port) {
self.handle_to_id.get(handle).copied()
} else {
None
}
})
}
struct Flow {
src: IpAddress,
dst: IpAddress,
src_p: u16,
dst_p: u16,
}
struct TargetResolver {
@@ -119,54 +39,22 @@ impl TargetResolver {
fn process_dns_query(&mut self, data: &[u8]) -> Option<Vec<u8>> {
self.dns_handler.handle_query(data, &mut self.fake_ip_store)
}
}
struct SocketFactory;
impl SocketFactory {
fn create_tcp<'a>(port: u16) -> tcp::Socket<'a> {
let cfg = NetworkConfig::global();
let buf_size = match port {
443 | 80 | 8080 | 1935 => cfg.tcp_buf_heavy,
_ => cfg.tcp_buf_light,
};
let mut socket = tcp::Socket::new(
tcp::SocketBuffer::new(vec![0; buf_size]),
tcp::SocketBuffer::new(vec![0; buf_size]),
);
socket.set_nagle_enabled(false);
socket.set_ack_delay(None);
socket
}
fn create_udp<'a>(port: u16) -> udp::Socket<'a> {
let cfg = NetworkConfig::global();
let (buf_size, meta_count) = match port {
443 | 4433 | 10000..=60000 => (cfg.udp_buf_heavy, cfg.udp_meta_heavy),
53 | 123 => (cfg.udp_buf_light, cfg.udp_meta_light),
_ => (cfg.udp_buf_heavy / 2, cfg.udp_meta_heavy / 2),
};
udp::Socket::new(
udp::PacketBuffer::new(
vec![udp::PacketMetadata::EMPTY; meta_count],
vec![0; buf_size],
pub fn resolve_destination(&self, addr: IpAddress, port: u16) -> (std::net::Ipv4Addr, String) {
match addr {
IpAddress::Ipv4(ip) => {
let std_ip = std::net::Ipv4Addr::from(ip);
if let Some(domain) = self.fake_ip_store.lookup_by_ip(&std_ip) {
(std_ip, format!("{}:{}", domain, port))
} else {
(std_ip, format!("{}:{}", std_ip, port))
}
}
IpAddress::Ipv6(ip) => (
std::net::Ipv4Addr::UNSPECIFIED,
format!("[{}]:{}", ip, port),
),
udp::PacketBuffer::new(
vec![udp::PacketMetadata::EMPTY; meta_count],
vec![0; buf_size],
),
)
}
fn create_icmp<'a>() -> icmp::Socket<'a> {
icmp::Socket::new(
icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; 8], vec![0; 2048]),
icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; 8], vec![0; 2048]),
)
}
}
}
@@ -174,87 +62,71 @@ pub struct ConnectionManager {
tracker: SessionTracker,
resolver: TargetResolver,
tx_to_tunnel: mpsc::Sender<RawCastFrame>,
factory: Arc<dyn SocketProvider>,
}
impl ConnectionManager {
pub fn new(dns_handler: DnsHandler, tx_to_tunnel: mpsc::Sender<RawCastFrame>) -> Self {
pub fn new(
dns_handler: DnsHandler,
tx_to_tunnel: mpsc::Sender<RawCastFrame>,
factory: Arc<dyn SocketProvider>,
) -> Self {
Self {
tracker: SessionTracker::new(),
resolver: TargetResolver::new(dns_handler),
tx_to_tunnel,
factory,
}
}
pub fn try_inject_inbound(&mut self, frame: RawCastFrame) -> Result<(), RawCastFrame> {
if frame.event != RawCastEvent::Data {
if frame.event == RawCastEvent::Close {
info!("💀 [Stream {}] Received CLOSE from tunnel", frame.socket_id);
self.tracker.inbound_tx.remove(&frame.socket_id);
}
if frame.event == RawCastEvent::Close {
info!("💀 [Stream {}] Received CLOSE from tunnel", frame.socket_id);
self.tracker.close_tunnel_session(frame.socket_id);
return Ok(());
}
if let Some(tx) = self.tracker.inbound_tx.get(&frame.socket_id) {
match tx.try_send(frame.payload.clone()) {
Ok(_) => {
trace!(
"📥 [Stream {}] Inbound data: {} bytes",
frame.socket_id,
frame.payload.len()
);
Ok(())
if frame.event != RawCastEvent::Data {
return Ok(());
}
if let Some(tx) = self.tracker.get_inbound_tx(frame.socket_id) {
tx.try_send(frame.payload.clone()).map_err(|e| {
if matches!(e, mpsc::error::TrySendError::Closed(_)) {
self.tracker.close_tunnel_session(frame.socket_id);
}
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
warn!("🟡 [Stream {}] Inbound channel FULL", frame.socket_id);
Err(frame)
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
error!("🔴 [Stream {}] Inbound channel CLOSED", frame.socket_id);
self.tracker.inbound_tx.remove(&frame.socket_id);
Ok(())
}
}
frame
})
} else {
trace!(
"👻 [Stream {}] ORPHAN packet from tunnel. ID mismatch?",
frame.socket_id
);
trace!("👻 [Stream {}] Orphan packet", frame.socket_id);
Ok(())
}
}
pub fn setup_sockets(n_icmp: usize) -> SocketSet<'static> {
let mut sockets = SocketSet::new(Vec::with_capacity(64));
let mut dns_socket = SocketFactory::create_udp(53);
let _ = dns_socket.bind(IpListenEndpoint {
addr: None,
port: 53,
});
sockets.add(dns_socket);
for _ in 0..n_icmp {
sockets.add(SocketFactory::create_icmp());
}
sockets
pub fn setup_sockets(factory: &dyn SocketProvider, n_icmp: usize) -> SocketSet<'static> {
factory.create_base_set(n_icmp)
}
pub fn start_listening(&mut self, socket_set: &mut SocketSet) {
for (_, socket) in socket_set.iter_mut() {
if let Some(tcp) = tcp::Socket::downcast_mut(socket) {
if !tcp.is_open() {
let _ = tcp.listen(IpListenEndpoint {
addr: None,
port: 443,
});
match socket {
Socket::Tcp(tcp) => {
if !tcp.is_open() {
let _ = tcp.listen(IpListenEndpoint {
addr: None,
port: 443,
});
}
}
} else if let Some(udp) = udp::Socket::downcast_mut(socket) {
if !udp.is_open() {
let _ = udp.bind(IpListenEndpoint {
addr: None,
port: 53,
});
Socket::Udp(udp) => {
if !udp.is_open() {
let _ = udp.bind(IpListenEndpoint {
addr: None,
port: 53,
});
}
}
_ => {}
}
}
}
@@ -274,20 +146,29 @@ impl ConnectionManager {
let Ok(ip) = Ipv4Packet::new_checked(packet) else {
return;
};
let (src, dst) = (ip.src_addr().into(), ip.dst_addr().into());
let mut flow = Flow {
src: ip.src_addr().into(),
dst: ip.dst_addr().into(),
src_p: 0,
dst_p: 0,
};
match ip.next_header() {
IpProtocol::Tcp => {
if let Ok(p) = TcpPacket::new_checked(ip.payload()) {
self.intercept_tcp(src, dst, p.src_port(), p.dst_port(), socket_set);
flow.src_p = p.src_port();
flow.dst_p = p.dst_port();
self.intercept_tcp(flow, socket_set);
}
}
IpProtocol::Udp => {
if let Ok(p) = UdpPacket::new_checked(ip.payload()) {
self.intercept_udp(src, dst, p.src_port(), p.dst_port(), socket_set);
flow.src_p = p.src_port();
flow.dst_p = p.dst_port();
self.intercept_udp(flow, socket_set);
}
}
p => trace!("Skip IPv4 protocol {:?}", p),
_ => {}
}
}
@@ -295,387 +176,145 @@ impl ConnectionManager {
let Ok(ip) = Ipv6Packet::new_checked(packet) else {
return;
};
let (src, dst) = (ip.src_addr().into(), ip.dst_addr().into());
let mut flow = Flow {
src: ip.src_addr().into(),
dst: ip.dst_addr().into(),
src_p: 0,
dst_p: 0,
};
match ip.next_header() {
IpProtocol::Tcp => {
if let Ok(p) = TcpPacket::new_checked(ip.payload()) {
self.intercept_tcp(src, dst, p.src_port(), p.dst_port(), socket_set);
flow.src_p = p.src_port();
flow.dst_p = p.dst_port();
self.intercept_tcp(flow, socket_set);
}
}
IpProtocol::Udp => {
if let Ok(p) = UdpPacket::new_checked(ip.payload()) {
self.intercept_udp(src, dst, p.src_port(), p.dst_port(), socket_set);
flow.src_p = p.src_port();
flow.dst_p = p.dst_port();
self.intercept_udp(flow, socket_set);
}
}
_ => {}
}
}
fn intercept_tcp(
&mut self,
src: smoltcp::wire::IpAddress,
dst: smoltcp::wire::IpAddress,
src_p: u16,
dst_p: u16,
socket_set: &mut SocketSet,
) {
if !self.tracker.has_connection_from(src, src_p, socket_set) {
info!(
"🆕 [TCP] New session detected: {}:{} -> {}:{}",
src, src_p, dst, dst_p
);
let mut socket = SocketFactory::create_tcp(dst_p);
let _ = socket.listen(IpListenEndpoint {
addr: Some(dst),
port: dst_p,
});
fn intercept_tcp(&mut self, f: Flow, socket_set: &mut SocketSet) {
if !self.tracker.has_connection_from(f.src, f.src_p, socket_set) {
let socket = self.factory.create_listening_tcp(Some(f.dst), f.dst_p);
let handle = socket_set.add(socket);
self.tracker.pending_tcp.insert(handle, StdInstant::now());
self.tracker.add_pending_tcp(handle);
}
}
fn intercept_udp(
&mut self,
src: smoltcp::wire::IpAddress,
dst: smoltcp::wire::IpAddress,
src_p: u16,
dst_p: u16,
socket_set: &mut SocketSet,
) {
if dst_p == 0
|| dst_p == 137
|| dst_p == 138
|| self.tracker.get_id_by_port(src_p).is_some()
fn intercept_udp(&mut self, f: Flow, socket_set: &mut SocketSet) {
if f.dst_p == 0
|| f.dst_p == 137
|| f.dst_p == 138
|| f.dst_p == 53
|| self.tracker.is_client_known(f.src_p)
{
return;
}
if dst_p == 53 {
return;
}
let socket_id = self.tracker.next_id();
let (dst_ip, target) = self.resolver.resolve_destination(f.dst, f.dst_p);
let socket_id = self.tracker.generate_socket_id();
let (dst_ip, target) = match dst {
smoltcp::wire::IpAddress::Ipv4(ip) => {
let std_ip = std::net::Ipv4Addr::from(ip);
if let Some(domain) = self.resolver.fake_ip_store.lookup_by_ip(&std_ip) {
info!(
"🔍 [UDP {}] Reverse lookup MATCH: {} -> {}",
socket_id, std_ip, domain
);
(std_ip, format!("{}:{}", domain, dst_p))
} else {
info!(
"🔍 [UDP {}] Reverse lookup MISS: using raw IP {}",
socket_id, std_ip
);
(std_ip, format!("{}:{}", std_ip, dst_p))
}
}
smoltcp::wire::IpAddress::Ipv6(ip) => {
let std_ip = std::net::Ipv4Addr::new(0, 0, 0, 0);
(std_ip, format!("[{}]:{}", ip, dst_p))
}
};
info!(
"🚀 [UDP MASTER] ID:{} | Intercepted: {}:{} -> Target: {}",
socket_id, src, src_p, target
);
let mut socket = SocketFactory::create_udp(dst_p);
if socket
.bind(IpListenEndpoint {
addr: Some(dst),
port: dst_p,
})
.is_ok()
{
let socket = self.factory.create_bound_udp(Some(f.dst), f.dst_p);
if socket.is_open() {
let handle = socket_set.add(socket);
let (conn, mut rx_smol, tx_smol) = UdpConnection::new(handle, src, src_p);
let (conn, rx_smol, tx_smol) = UdpConnection::new(handle, f.src, f.src_p);
self.tracker.handle_to_id.insert(handle, socket_id);
self.tracker.active_udp.insert(handle, conn);
self.tracker.inbound_tx.insert(socket_id, tx_smol);
let tx_tunnel = self.tx_to_tunnel.clone();
tokio::spawn(async move {
info!("📡 [UDP {}] Task started for {}", socket_id, target);
let mut frame = RawCastFrame::connect(LocalProtocol::Udp, socket_id, dst_ip, dst_p);
frame.payload = bytes::Bytes::from(target.clone());
info!(
"📦 [UDP {}] Packed Connect frame: dst_ip={}, dst_port={}, payload={}",
socket_id, dst_ip, dst_p, target
);
if tx_tunnel.send(frame).await.is_err() {
error!("❌ [UDP {}] Failed to send CONNECT to tunnel", socket_id);
return;
}
let mut pkt_count = 0;
while let Some((data, ip, port)) = rx_smol.recv().await {
pkt_count += 1;
if pkt_count == 1 {
info!("📤 [UDP {}] First data packet sent to tunnel", socket_id);
}
let df =
RawCastFrame::data(LocalProtocol::Udp, socket_id, ip, port, data.to_vec());
if tx_tunnel.send(df).await.is_err() {
break;
}
}
info!(
"🛑 [UDP {}] Task stopped. Sent {} packets",
socket_id, pkt_count
);
});
self.tracker.register_udp(handle, socket_id, conn, tx_smol);
UdpConnection::spawn(
socket_id,
dst_ip,
f.dst_p,
target,
rx_smol,
self.tx_to_tunnel.clone(),
);
}
}
pub fn process_sockets(&mut self, socket_set: &mut SocketSet) {
for (handle, socket) in socket_set.iter_mut() {
if let Some(s) = tcp::Socket::downcast_mut(socket) {
self.handle_tcp(handle, s);
} else if let Some(s) = udp::Socket::downcast_mut(socket) {
self.handle_udp(handle, s);
} else if let Some(s) = icmp::Socket::downcast_mut(socket) {
self.handle_icmp(handle, s);
self.handle_icmpv6(handle, s);
match socket {
Socket::Tcp(s) => self.handle_tcp(handle, s),
Socket::Udp(s) => self.handle_udp(handle, s),
Socket::Icmp(s) => IcmpResponder::handle(s),
}
}
}
fn handle_tcp(&mut self, handle: SocketHandle, socket: &mut tcp::Socket) {
self.tracker.update_activity(handle);
if socket.state() == tcp::State::Closed {
if let Some(id) = self.tracker.handle_to_id.get(&handle) {
info!("🏁 [TCP {}] Connection closed", id);
}
self.tracker.pending_tcp.remove(&handle);
self.tracker.queue_removal(handle);
return;
}
if matches!(socket.state(), tcp::State::Listen | tcp::State::SynReceived) {
if let Some(created_at) = self.tracker.pending_tcp.get(&handle) {
if created_at.elapsed() > std::time::Duration::from_secs(20) {
warn!(
"🧹 [TCP] Cleaning up ghost socket stuck in {:?}",
socket.state()
);
socket.abort();
self.tracker.queue_removal(handle);
}
}
if self
.tracker
.check_pending_timeout(handle, Duration::from_secs(20))
{
socket.abort();
self.tracker.queue_removal(handle);
return;
}
if socket.state() == tcp::State::Established {
if self.tracker.pending_tcp.remove(&handle).is_some()
|| !self.tracker.active_tcp.contains_key(&handle)
{
let Some(ep) = socket.local_endpoint() else {
return;
};
let socket_id = self.tracker.generate_socket_id();
let (conn, mut rx_smol, tx_smol, handshake_tx) = TcpConnection::new(handle);
if socket.state() == tcp::State::Established && self.tracker.should_init_tcp(handle) {
let ep = socket.local_endpoint().unwrap();
let socket_id = self.tracker.next_id();
let (dst_ip, target) = self.resolver.resolve_destination(ep.addr, ep.port);
self.tracker.handle_to_id.insert(handle, socket_id);
self.tracker.active_tcp.insert(handle, conn);
self.tracker.inbound_tx.insert(socket_id, tx_smol);
let (conn, rx_smol, tx_smol, handshake_tx) = TcpConnection::new(handle);
self.tracker.register_tcp(handle, socket_id, conn, tx_smol);
let (dst_ip, target) = match ep.addr {
smoltcp::wire::IpAddress::Ipv4(ip) => {
let std_ip = std::net::Ipv4Addr::from(ip);
if let Some(domain) = self.resolver.fake_ip_store.lookup_by_ip(&std_ip) {
(std_ip, format!("{}:{}", domain, ep.port))
} else {
(std_ip, format!("{}:{}", std_ip, ep.port))
}
}
smoltcp::wire::IpAddress::Ipv6(ip) => (
std::net::Ipv4Addr::new(0, 0, 0, 0),
format!("[{}]:{}", ip, ep.port),
),
};
info!("🔗 [TCP {}] established -> {}", socket_id, target);
let tx_tunnel = self.tx_to_tunnel.clone();
tokio::spawn(async move {
let mut frame =
RawCastFrame::connect(LocalProtocol::Tcp, socket_id, dst_ip, ep.port);
frame.payload = bytes::Bytes::from(target.clone());
if tx_tunnel.send(frame).await.is_ok() {
let _ = handshake_tx.send(());
while let Some(data) = rx_smol.recv().await {
let _ = tx_tunnel
.send(RawCastFrame::data(
LocalProtocol::Tcp,
socket_id,
dst_ip,
ep.port,
data.to_vec(),
))
.await;
}
}
});
}
TcpConnection::spawn(
socket_id,
dst_ip,
ep.port,
target,
rx_smol,
handshake_tx,
self.tx_to_tunnel.clone(),
);
}
if let Some(conn) = self.tracker.active_tcp.get_mut(&handle) {
if let Some(conn) = self.tracker.get_tcp_mut(handle) {
if !conn.tick(socket) {
info!("⚠️ [TCP] Tick failed, aborting handle {:?}", handle);
socket.abort();
}
}
}
fn handle_udp(&mut self, handle: SocketHandle, socket: &mut udp::Socket) {
self.tracker.last_activity.insert(handle, StdInstant::now());
self.tracker.update_activity(handle);
if socket.endpoint().port == 53 {
while let Ok((data, meta)) = socket.recv() {
if let Some(res) = self.resolver.process_dns_query(data) {
debug!("🔍 [DNS] Resolved query for client {}", meta.endpoint);
if let Err(e) = socket.send_slice(&res, meta) {
warn!("❌ [DNS] Failed to send response: {:?}", e);
}
let _ = socket.send_slice(&res, meta);
}
}
return;
}
if let Some(conn) = self.tracker.active_udp.get_mut(&handle) {
if let Some(conn) = self.tracker.get_udp_mut(handle) {
if !conn.tick(socket) {
if let Some(socket_id) = self.tracker.handle_to_id.get(&handle) {
info!("🛑 [UDP {}] Session expired or closed by tick", socket_id);
}
self.tracker.queue_removal(handle);
}
}
}
fn handle_icmp(&mut self, _handle: SocketHandle, socket: &mut icmp::Socket) {
if !socket.can_recv() {
return;
}
match socket.recv() {
Ok((data, src_addr)) => {
let Ok(pkt) = smoltcp::wire::Icmpv4Packet::new_checked(data) else {
return;
};
match pkt.msg_type() {
smoltcp::wire::Icmpv4Message::EchoRequest => {
info!("🏓 [ICMPv4] Ping Request from {}", src_addr);
let mut reply_data = data.to_vec();
let mut reply_pkt =
smoltcp::wire::Icmpv4Packet::new_unchecked(&mut reply_data);
reply_pkt.set_msg_type(smoltcp::wire::Icmpv4Message::EchoReply);
reply_pkt.fill_checksum();
if let Err(e) = socket.send_slice(&reply_data, src_addr) {
warn!("❌ [ICMPv4] Reply failed to {}: {:?}", src_addr, e);
} else {
info!("✅ [ICMPv4] Echo Reply sent to {}", src_addr);
}
}
smoltcp::wire::Icmpv4Message::DstUnreachable => {
warn!(
"🚫 [ICMPv4] Destination Unreachable from {}. Check MTU!",
src_addr
);
}
_ => debug!("📡 [ICMPv4] Message {:?} from {}", pkt.msg_type(), src_addr),
}
}
Err(e) => trace!("ICMPv4 recv error: {:?}", e),
}
}
fn handle_icmpv6(&mut self, _handle: SocketHandle, socket: &mut icmp::Socket) {
if !socket.can_recv() {
return;
}
match socket.recv() {
Ok((data, src_addr)) => {
let smoltcp::wire::IpAddress::Ipv6(ipv6_src) = src_addr else {
return;
};
let Ok(pkt) = smoltcp::wire::Icmpv6Packet::new_checked(data) else {
return;
};
match pkt.msg_type() {
smoltcp::wire::Icmpv6Message::EchoRequest => {
info!("🏓 [ICMPv6] Ping Request from {}", ipv6_src);
let mut reply_data = data.to_vec();
let mut reply_pkt =
smoltcp::wire::Icmpv6Packet::new_unchecked(&mut reply_data);
reply_pkt.set_msg_type(smoltcp::wire::Icmpv6Message::EchoReply);
let my_v6_gateway =
smoltcp::wire::Ipv6Address::new(0xfe80, 0, 0, 0, 0, 0, 0, 1);
reply_pkt.fill_checksum(&my_v6_gateway, &ipv6_src);
if let Err(e) = socket.send_slice(&reply_data, src_addr) {
warn!("❌ [ICMPv6] Failed to send reply: {:?}", e);
} else {
info!("✅ [ICMPv6] Echo Reply sent to {}", src_addr);
}
}
smoltcp::wire::Icmpv6Message::DstUnreachable => {
warn!("🚫 [ICMPv6] Destination Unreachable from {}", ipv6_src);
}
smoltcp::wire::Icmpv6Message::PktTooBig => {
warn!(
"📏 [ICMPv6] Packet Too Big! MTU issue detected at {}",
ipv6_src
);
}
_ => {}
}
}
Err(e) => trace!("ICMPv6 recv error: {:?}", e),
}
}
pub fn cleanup(&mut self, socket_set: &mut SocketSet) {
self.tracker.enforce_idle_timeouts(Duration::from_secs(120));
self.tracker.cleanup(socket_set);
}
pub fn log_status(&self, socket_set: &SocketSet) {
let mut est = 0;
let mut total = 0;
for (_, socket) in socket_set.iter() {
if let Some(tcp) = tcp::Socket::downcast(socket) {
total += 1;
if tcp.state() == tcp::State::Established {
est += 1;
}
}
}
debug!(
"TCP Stats: Total={}, Established={}, Active={}",
total,
est,
self.tracker.active_tcp.len()
);
}
}
client/src/net/dns.rs
+70 -14
View File
@@ -1,16 +1,70 @@
use crate::net::ip_store::FakeIpStore;
use anyhow::Result;
use hickory_proto::op::{Message, MessageType, ResponseCode};
use hickory_proto::rr::{RData, Record, RecordType};
use lru::LruCache;
use netrunner_logger::{debug, error, info};
use std::collections::HashSet;
use std::net::Ipv4Addr;
use std::num::NonZeroUsize;
use std::path::PathBuf;
use std::time::{Duration, SystemTime};
use tokio::fs::{self, File};
use tokio::io::{AsyncBufReadExt, BufReader};
// --- Fake IP Storage ---
const START_IP: u32 = 0x64400001; // 100.64.0.1
pub struct FakeIpStore {
cache: LruCache<String, Ipv4Addr>,
rev_cache: LruCache<Ipv4Addr, String>,
next_ip: u32,
}
impl FakeIpStore {
pub fn new() -> Self {
info!("Initializing FakeIpStore starting at 100.64.0.1");
Self {
cache: LruCache::new(NonZeroUsize::new(2000).unwrap()),
rev_cache: LruCache::new(NonZeroUsize::new(2000).unwrap()),
next_ip: START_IP,
}
}
pub fn get_or_assign(&mut self, host: &str) -> Ipv4Addr {
if let Some(&ip) = self.cache.get(host) {
debug!(host = %host, ip = %ip, "Cache hit: IP already assigned");
return ip;
}
let ip = Ipv4Addr::from(self.next_ip);
self.next_ip += 1;
self.cache.put(host.to_string(), ip);
self.rev_cache.put(ip, host.to_string());
debug!(host = %host, ip = %ip, "Assigned new fake IP");
ip
}
pub fn lookup_by_ip(&self, ip: &Ipv4Addr) -> Option<String> {
if let Some(host) = self.rev_cache.peek(ip) {
debug!(ip = %ip, host = %host, "Reverse lookup successful");
Some(host.clone())
} else {
debug!(ip = %ip, "Reverse lookup miss");
None
}
}
}
// --- DNS Handler & Blocklist Logic ---
pub struct DnsHandler {
block_list: HashSet<String>,
forbidden_suffixes: Vec<String>,
cache_path: String,
}
impl DnsHandler {
pub fn new(cache_dir: &str) -> Self {
Self {
@@ -22,8 +76,9 @@ impl DnsHandler {
cache_path: format!("{}/hosts_cache.txt", cache_dir),
}
}
pub async fn init(&mut self) -> anyhow::Result<()> {
let path = std::path::PathBuf::from(&self.cache_path);
pub async fn init(&mut self) -> Result<()> {
let path = PathBuf::from(&self.cache_path);
if path.exists() {
let _ = self.load_from_file().await;
@@ -33,19 +88,17 @@ impl DnsHandler {
SystemTime::now()
.duration_since(meta.modified()?)
.unwrap_or_default()
> Duration::from_secs(604800)
> Duration::from_secs(604800) // 1 неделя
} else {
true
};
if needs_update {
let p_clone = self.cache_path.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(10)).await;
if let Err(e) = Self::download_blocklist_async(p_clone).await {
netrunner_logger::error!("DNS: Background update failed: {}", e);
error!("DNS: Background update failed: {}", e);
}
});
}
@@ -53,9 +106,8 @@ impl DnsHandler {
Ok(())
}
async fn download_blocklist_async(cache_path: String) -> anyhow::Result<()> {
async fn download_blocklist_async(cache_path: String) -> Result<()> {
info!("DNS: Starting background download to {}", cache_path);
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()?;
@@ -67,16 +119,15 @@ impl DnsHandler {
if resp.status().is_success() {
let bytes = resp.bytes().await?;
tokio::fs::write(&cache_path, bytes).await?;
info!("DNS: Blocklist downloaded successfully to {}", cache_path);
fs::write(&cache_path, bytes).await?;
info!("DNS: Blocklist downloaded successfully.");
} else {
error!("DNS: Download failed with status {}", resp.status());
}
Ok(())
}
async fn load_from_file(&mut self) -> anyhow::Result<()> {
async fn load_from_file(&mut self) -> Result<()> {
let file = File::open(&self.cache_path).await?;
let mut lines = BufReader::new(file).lines();
let mut count = 0;
@@ -92,7 +143,7 @@ impl DnsHandler {
count += 1;
}
}
info!("DNS: Loaded {} domains from cache.", count);
info!("DNS: Loaded {} domains from blocklist.", count);
Ok(())
}
@@ -111,6 +162,7 @@ impl DnsHandler {
.set_recursion_available(true)
.add_query(query.clone());
// 1. Проверка блокировок
if self.forbidden_suffixes.iter().any(|s| name.ends_with(s))
|| self.block_list.contains(&name)
{
@@ -119,6 +171,7 @@ impl DnsHandler {
return res.to_vec().ok();
}
// 2. Генерация Fake IP для A-записей
if query.query_type() == RecordType::A {
let fake_ip = store.get_or_assign(&name);
res.add_answer(Record::from_rdata(
@@ -127,6 +180,9 @@ impl DnsHandler {
RData::A(fake_ip.into()),
));
res.set_response_code(ResponseCode::NoError);
} else {
// Для остальных типов записей (AAAA и т.д.) просто возвращаем пустой NoError или NXDomain
res.set_response_code(ResponseCode::NoError);
}
res.to_vec().ok()
client/src/net/engine.rs
+45 -44
View File
@@ -1,5 +1,5 @@
use netrunner_core::net::ClientHandler;
use netrunner_core::net::network::NetworkConfig;
use netrunner_core::net::NetworkConfig;
use netrunner_core::rawcast::RawCastFrame;
use smoltcp::iface::PollResult;
use smoltcp::time::Instant;
@@ -23,6 +23,7 @@ use netrunner_logger::{debug, error, info, warn};
use crate::net::connection_manager::ConnectionManager;
use crate::net::dns::DnsHandler;
use crate::net::socket_factory::{SmolSocketFactory, SocketProvider};
use crate::tun::device::{TokenBuffer, VirtTunDevice};
use crate::tun::routing::setup_platform_routing;
use crate::tun::tun::Tun;
@@ -39,6 +40,7 @@ pub struct Engine {
avail: Arc<AtomicBool>,
rx_from_tunnel: mpsc::Receiver<RawCastFrame>,
}
impl Engine {
pub fn new(
config: Config,
@@ -46,15 +48,16 @@ impl Engine {
dns_handler: DnsHandler,
tx_to_tunnel: mpsc::Sender<RawCastFrame>,
rx_from_tunnel: mpsc::Receiver<RawCastFrame>,
factory: Arc<dyn SocketProvider>,
) -> Self {
let now = Engine::current_time();
let (mut device, to_smoltcp_tx, from_smoltcp_rx, avail) = VirtTunDevice::new(caps);
let interface = Interface::new(config, &mut device, now);
let socket_set = ConnectionManager::setup_sockets(2);
let socket_set = ConnectionManager::setup_sockets(factory.as_ref(), 2);
let manager = ConnectionManager::new(dns_handler, tx_to_tunnel);
let manager = ConnectionManager::new(dns_handler, tx_to_tunnel, factory.clone());
Self {
interface,
@@ -78,13 +81,13 @@ impl Engine {
let from_smoltcp_rx = self.from_smoltcp_rx.take().expect("Engine started twice");
Self::spawn_tun_writer(writer, from_smoltcp_rx);
let mut last_log = StdInstant::now();
let mut stuck_frame: Option<RawCastFrame> = None;
loop {
while let Ok(token) = tun_to_engine_rx.try_recv() {
self.manager
.try_create_socket_from_packet(&token, &mut self.socket_set);
if self.to_smoltcp_tx.send(token).is_ok() {
self.device.mark_rx_available();
}
@@ -106,12 +109,8 @@ impl Engine {
}
self.manager.process_sockets(&mut self.socket_set);
let result = self.poll();
if last_log.elapsed() >= Duration::from_secs(5) {
self.manager.log_status(&self.socket_set);
last_log = StdInstant::now();
}
let result = self.poll();
if matches!(result, PollResult::SocketStateChanged) {
self.manager.cleanup(&mut self.socket_set);
@@ -132,29 +131,17 @@ impl Engine {
}
};
if stuck_frame.is_none() {
tokio::select! {
_ = sleep_fut => {}
Some(token) = tun_to_engine_rx.recv() => {
self.manager.try_create_socket_from_packet(&token, &mut self.socket_set);
if self.to_smoltcp_tx.send(token).is_ok() {
self.device.mark_rx_available();
}
}
Some(frame) = self.rx_from_tunnel.recv() => {
if let Err(returned_frame) = self.manager.try_inject_inbound(frame) {
stuck_frame = Some(returned_frame);
}
tokio::select! {
_ = sleep_fut => {}
Some(token) = tun_to_engine_rx.recv() => {
self.manager.try_create_socket_from_packet(&token, &mut self.socket_set);
if self.to_smoltcp_tx.send(token).is_ok() {
self.device.mark_rx_available();
}
}
} else {
tokio::select! {
_ = sleep_fut => {}
Some(token) = tun_to_engine_rx.recv() => {
self.manager.try_create_socket_from_packet(&token, &mut self.socket_set);
if self.to_smoltcp_tx.send(token).is_ok() {
self.device.mark_rx_available();
}
Some(frame) = self.rx_from_tunnel.recv(), if stuck_frame.is_none() => {
if let Err(returned_frame) = self.manager.try_inject_inbound(frame) {
stuck_frame = Some(returned_frame);
}
}
}
@@ -302,6 +289,8 @@ impl EngineConfig {
pub struct EngineBuilder {
config: EngineConfig,
tun_device: Option<Tun>,
// Добавляем возможность прокинуть кастомную фабрику сокетов
socket_factory: Option<Arc<dyn SocketProvider>>,
}
impl EngineBuilder {
@@ -309,6 +298,7 @@ impl EngineBuilder {
Self {
config,
tun_device: None,
socket_factory: None,
}
}
@@ -320,10 +310,7 @@ impl EngineBuilder {
pub async fn build(self) -> Result<(Engine, Tun), String> {
let tun = self.tun_device.ok_or("TUN device is required")?;
info!(
"Initializing Engine components with config: {:?}",
self.config
);
info!("Initializing Engine with config: {:?}", self.config);
let mut dns_handler = DnsHandler::new(&self.config.cache_path);
if let Err(e) = dns_handler.init().await {
@@ -332,11 +319,8 @@ impl EngineBuilder {
if self.config.setup_routing {
info!("Applying platform routing rules...");
if let Err(e) = setup_platform_routing(&self.config.remote_address) {
return Err(format!("Routing setup failed: {}", e));
}
} else {
info!("Platform routing setup skipped via config.");
setup_platform_routing(&self.config.remote_address)
.map_err(|e| format!("Routing setup failed: {}", e))?;
}
let smol_config = Config::new(smoltcp::wire::HardwareAddress::Ip);
@@ -351,7 +335,6 @@ impl EngineBuilder {
mpsc::channel(NetworkConfig::global().muxer_capacity);
info!("Establishing secure tunnel to proxy server...");
ClientHandler::connect(
&self.config.remote_address,
rx_for_client_handler,
@@ -360,12 +343,27 @@ impl EngineBuilder {
.await
.map_err(|e| format!("Failed to establish secure tunnel: {}", e))?;
info!("Secure tunnel established, Muxer is ready.");
let factory = self.socket_factory.unwrap_or_else(|| {
// 1. Разыменовываем ссылку и клонируем данные в новый объект
let config_owned = (*NetworkConfig::global()).clone();
let mut engine = Engine::new(smol_config, caps, dns_handler, tx_to_tunnel, rx_from_tunnel);
// 2. Оборачиваем во владеющий Arc
let config = Arc::new(config_owned);
// Теперь типы совпадают: Arc<NetworkConfig> -> SmolSocketFactory::new
Arc::new(SmolSocketFactory::new(config))
});
let mut engine = Engine::new(
smol_config,
caps,
dns_handler,
tx_to_tunnel,
rx_from_tunnel,
factory,
);
engine.set_any_ip(self.config.any_ip);
if self.config.transparent_mode {
engine.set_transparent_mode();
}
@@ -373,7 +371,10 @@ impl EngineBuilder {
engine.set_default_gateway(self.config.default_gateway);
engine.activate();
info!("Stack IP initialized: {}", self.config.default_gateway);
info!(
"Engine successfully built. Stack IP: {}",
self.config.default_gateway
);
Ok((engine, tun))
}
client/src/net/ip_store.rs
-46
View File
@@ -1,46 +0,0 @@
use lru::LruCache;
use netrunner_logger::{debug, info};
use std::net::Ipv4Addr;
use std::num::NonZeroUsize;
pub struct FakeIpStore {
cache: LruCache<String, Ipv4Addr>,
rev_cache: LruCache<Ipv4Addr, String>,
next_ip: u32,
}
const START_IP: u32 = 0x64400001;
impl FakeIpStore {
pub fn new() -> Self {
info!("Initializing FakeIpStore starting at 100.64.0.1");
Self {
cache: LruCache::new(NonZeroUsize::new(2000).unwrap()),
rev_cache: LruCache::new(NonZeroUsize::new(2000).unwrap()),
next_ip: START_IP,
}
}
pub fn get_or_assign(&mut self, host: &str) -> Ipv4Addr {
if let Some(&ip) = self.cache.get(host) {
debug!(host = %host, ip = %ip, "Cache hit: IP already assigned");
return ip;
}
let ip = Ipv4Addr::from(self.next_ip);
self.next_ip += 1;
self.cache.put(host.to_string(), ip);
self.rev_cache.put(ip, host.to_string());
debug!(host = %host, ip = %ip, "Assigned new fake IP");
ip
}
pub fn lookup_by_ip(&self, ip: &Ipv4Addr) -> Option<String> {
if let Some(host) = self.rev_cache.peek(ip) {
debug!(ip = %ip, host = %host, "Reverse lookup successful");
Some(host.clone())
} else {
debug!(ip = %ip, "Reverse lookup miss");
None
}
}
}
client/src/net/mod.rs
+3 -2
View File
@@ -1,5 +1,6 @@
mod connection;
pub mod connection_manager;
pub mod dns;
mod dns;
pub mod engine;
pub mod ip_store;
mod session_tracker;
mod socket_factory;
client/src/net/session_tracker.rs
+177
View File
@@ -0,0 +1,177 @@
use std::{
collections::HashMap,
time::{Duration, Instant as StdInstant},
};
use bytes::Bytes;
use smoltcp::{
iface::{SocketHandle, SocketSet},
socket::Socket,
wire::IpAddress,
};
use tokio::sync::mpsc;
use crate::net::connection::{TcpConnection, UdpConnection};
pub struct SessionTracker {
last_activity: HashMap<SocketHandle, StdInstant>,
active_tcp: HashMap<SocketHandle, TcpConnection>,
active_udp: HashMap<SocketHandle, UdpConnection>,
inbound_tx: HashMap<u64, mpsc::Sender<Bytes>>,
handle_to_id: HashMap<SocketHandle, u64>,
// 🔥 ДОБАВЛЕНО: Обратный маппинг, чтобы находить сокет по ID из туннеля
id_to_handle: HashMap<u64, SocketHandle>,
pending_tcp: HashMap<SocketHandle, StdInstant>,
to_remove: Vec<SocketHandle>,
next_socket_id: u64,
}
impl SessionTracker {
pub fn new() -> Self {
Self {
last_activity: HashMap::new(),
active_tcp: HashMap::new(),
active_udp: HashMap::new(),
inbound_tx: HashMap::new(),
handle_to_id: HashMap::new(),
id_to_handle: HashMap::new(), // Инициализация
pending_tcp: HashMap::new(),
to_remove: Vec::new(),
next_socket_id: 1,
}
}
pub fn next_id(&mut self) -> u64 {
let id = self.next_socket_id;
self.next_socket_id = self.next_socket_id.wrapping_add(1);
id
}
pub fn add_pending_tcp(&mut self, handle: SocketHandle) {
self.pending_tcp.insert(handle, StdInstant::now());
}
pub fn register_tcp(
&mut self,
handle: SocketHandle,
id: u64,
conn: TcpConnection,
tx: mpsc::Sender<Bytes>,
) {
self.pending_tcp.remove(&handle);
self.handle_to_id.insert(handle, id);
self.id_to_handle.insert(id, handle); // 🔥 Регистрация обратного индекса
self.active_tcp.insert(handle, conn);
self.inbound_tx.insert(id, tx);
self.last_activity.insert(handle, StdInstant::now()); // Для TCP тоже полезно
}
pub fn register_udp(
&mut self,
handle: SocketHandle,
id: u64,
conn: UdpConnection,
tx: mpsc::Sender<Bytes>,
) {
self.handle_to_id.insert(handle, id);
self.id_to_handle.insert(id, handle); // 🔥 Регистрация обратного индекса
self.active_udp.insert(handle, conn);
self.inbound_tx.insert(id, tx);
self.last_activity.insert(handle, StdInstant::now());
}
pub fn has_connection_from(
&self,
src_addr: IpAddress,
src_port: u16,
socket_set: &SocketSet,
) -> bool {
socket_set.iter().any(|(_, s)| {
if let Socket::Tcp(tcp) = s {
if let Some(remote) = tcp.remote_endpoint() {
return remote.addr == src_addr && remote.port == src_port;
}
}
false
})
}
pub fn is_client_known(&self, port: u16) -> bool {
self.active_udp.values().any(|conn| conn.has_client(port))
}
pub fn should_init_tcp(&mut self, handle: SocketHandle) -> bool {
self.pending_tcp.contains_key(&handle) && !self.active_tcp.contains_key(&handle)
}
pub fn check_pending_timeout(&self, handle: SocketHandle, timeout: Duration) -> bool {
self.pending_tcp
.get(&handle)
.map_or(false, |t| t.elapsed() > timeout)
}
pub fn get_tcp_mut(&mut self, handle: SocketHandle) -> Option<&mut TcpConnection> {
self.active_tcp.get_mut(&handle)
}
pub fn get_udp_mut(&mut self, handle: SocketHandle) -> Option<&mut UdpConnection> {
self.active_udp.get_mut(&handle)
}
pub fn update_activity(&mut self, handle: SocketHandle) {
self.last_activity.insert(handle, StdInstant::now());
}
pub fn get_inbound_tx(&self, id: u64) -> Option<&mpsc::Sender<Bytes>> {
self.inbound_tx.get(&id)
}
pub fn close_tunnel_session(&mut self, id: u64) {
self.inbound_tx.remove(&id);
// 🔥 ИСПРАВЛЕНО: Теперь мы реально находим сокет и ставим его в очередь на удаление
if let Some(handle) = self.id_to_handle.remove(&id) {
self.queue_removal(handle);
}
}
pub fn queue_removal(&mut self, handle: SocketHandle) {
if !self.to_remove.contains(&handle) {
self.to_remove.push(handle);
}
}
// 🔥 ДОБАВЛЕНО: Сборщик мусора для залипших сокетов
pub fn enforce_idle_timeouts(&mut self, timeout: Duration) {
let now = StdInstant::now();
let mut ghosts = Vec::new();
for (&handle, &last_seen) in &self.last_activity {
if now.duration_since(last_seen) > timeout {
ghosts.push(handle);
}
}
for handle in ghosts {
netrunner_logger::debug!("🧹 Sweeping idle ghost socket: {:?}", handle);
self.queue_removal(handle);
}
}
pub fn cleanup(&mut self, socket_set: &mut SocketSet) {
for handle in self.to_remove.drain(..) {
// Удаление из socket_set автоматически делает abort() внутри smoltcp
// и освобождает память пула/буферов
socket_set.remove(handle);
self.active_tcp.remove(&handle);
self.active_udp.remove(&handle);
self.pending_tcp.remove(&handle);
self.last_activity.remove(&handle);
if let Some(id) = self.handle_to_id.remove(&handle) {
self.inbound_tx.remove(&id);
self.id_to_handle.remove(&id);
}
}
}
}
client/src/net/socket_factory.rs
+153
View File
@@ -0,0 +1,153 @@
use crossbeam_queue::ArrayQueue;
use netrunner_core::net::NetworkConfig;
use smoltcp::{
iface::SocketSet,
socket::{icmp, tcp, udp},
time::Duration,
wire::{IpAddress, IpListenEndpoint},
};
use std::sync::Arc;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TrafficProfile {
Interactive,
Bulk,
Dns,
Default,
}
impl TrafficProfile {
pub fn guess_from_port(port: u16, is_tcp: bool) -> Self {
match (port, is_tcp) {
(22, true) | (3389, true) | (5900, true) => Self::Interactive,
(443, true) | (80, true) | (8080, true) | (1935, true) => Self::Bulk,
(53, false) | (123, false) => Self::Dns,
_ => Self::Default,
}
}
}
pub trait SocketProvider: Send + Sync {
fn create_tcp(&self, profile: TrafficProfile) -> tcp::Socket<'static>;
fn create_udp(&self, profile: TrafficProfile) -> udp::Socket<'static>;
fn create_icmp(&self) -> icmp::Socket<'static>;
fn create_listening_tcp(&self, addr: Option<IpAddress>, port: u16) -> tcp::Socket<'static>;
fn create_bound_udp(&self, addr: Option<IpAddress>, port: u16) -> udp::Socket<'static>;
fn create_base_set(&self, n_icmp: usize) -> SocketSet<'static>;
fn reconfigure_tcp(&self, socket: &mut tcp::Socket, profile: TrafficProfile);
}
pub struct SmolSocketFactory {
config: Arc<NetworkConfig>,
heavy_pool: Arc<ArrayQueue<Vec<u8>>>,
}
impl SmolSocketFactory {
pub fn new(config: Arc<NetworkConfig>) -> Self {
Self {
config,
// Пул на 128 тяжелых буферов. Если нужно больше — queue вернет ошибку,
// и мы просто отдадим память ОС (что безопасно).
heavy_pool: Arc::new(ArrayQueue::new(128)),
}
}
/// Берем буфер из пула или создаем новый, если пул пуст
fn alloc_buf(&self, size: usize) -> Vec<u8> {
if size == self.config.tcp_buf_heavy {
if let Some(mut buf) = self.heavy_pool.pop() {
buf.fill(0); // Очищаем данные от прошлой сессии
return buf;
}
}
vec![0u8; size]
}
}
impl SocketProvider for SmolSocketFactory {
fn create_tcp(&self, profile: TrafficProfile) -> tcp::Socket<'static> {
let (rx_size, tx_size) = match profile {
TrafficProfile::Bulk => (self.config.tcp_buf_heavy, self.config.tcp_buf_heavy),
TrafficProfile::Interactive => (self.config.tcp_buf_light, self.config.tcp_buf_light),
_ => (self.config.tcp_buf_light * 2, self.config.tcp_buf_light * 2),
};
let rx_buffer = tcp::SocketBuffer::new(self.alloc_buf(rx_size));
let tx_buffer = tcp::SocketBuffer::new(self.alloc_buf(tx_size));
let mut socket = tcp::Socket::new(rx_buffer, tx_buffer);
self.reconfigure_tcp(&mut socket, profile);
socket.set_keep_alive(Some(Duration::from_secs(30)));
socket.set_timeout(Some(Duration::from_secs(60)));
socket
}
fn reconfigure_tcp(&self, socket: &mut tcp::Socket, profile: TrafficProfile) {
match profile {
TrafficProfile::Interactive => {
socket.set_nagle_enabled(false);
socket.set_ack_delay(None);
}
TrafficProfile::Bulk | TrafficProfile::Default => {
socket.set_nagle_enabled(true);
socket.set_ack_delay(Some(Duration::from_millis(10)));
}
_ => {}
}
}
fn create_udp(&self, profile: TrafficProfile) -> udp::Socket<'static> {
let (buf_size, meta_count) = match profile {
TrafficProfile::Bulk => (self.config.udp_buf_heavy, self.config.udp_meta_heavy),
TrafficProfile::Dns => (self.config.udp_buf_light, self.config.udp_meta_light),
_ => (
self.config.udp_buf_heavy / 2,
self.config.udp_meta_heavy / 2,
),
};
udp::Socket::new(
udp::PacketBuffer::new(
vec![udp::PacketMetadata::EMPTY; meta_count],
self.alloc_buf(buf_size),
),
udp::PacketBuffer::new(
vec![udp::PacketMetadata::EMPTY; meta_count],
self.alloc_buf(buf_size),
),
)
}
fn create_icmp(&self) -> icmp::Socket<'static> {
icmp::Socket::new(
icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; 4], vec![0; 512]),
icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; 4], vec![0; 512]),
)
}
fn create_listening_tcp(&self, addr: Option<IpAddress>, port: u16) -> tcp::Socket<'static> {
let profile = TrafficProfile::guess_from_port(port, true);
let mut socket = self.create_tcp(profile);
let _ = socket.listen(IpListenEndpoint { addr, port });
socket
}
fn create_bound_udp(&self, addr: Option<IpAddress>, port: u16) -> udp::Socket<'static> {
let profile = TrafficProfile::guess_from_port(port, false);
let mut socket = self.create_udp(profile);
let _ = socket.bind(IpListenEndpoint { addr, port });
socket
}
fn create_base_set(&self, n_icmp: usize) -> SocketSet<'static> {
let mut sockets = SocketSet::new(Vec::with_capacity(128));
sockets.add(self.create_bound_udp(None, 53)); // DNS
for _ in 0..n_icmp {
sockets.add(self.create_icmp());
}
sockets
}
}