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

connection manager refactoring

Browse Source
This commit is contained in:
Kirill
2026-03-24 13:39:21 +07:00
parent 9177b61a02
commit 7603d2c92e
5 changed files with 313 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/src/main.rs
+26 -53
View File
@@ -1,27 +1,14 @@
use netrunner_client::{
connections::dns::DnsHandler,
tun::{
engine::Engine,
routing::{reset_platform_routing, setup_platform_routing},
tun::Tun,
},
};
use netrunner_core::proxy::connection::connection::ClientHandler;
use netrunner_client::tun::{engine::EngineBuilder, routing::reset_platform_routing, tun::Tun};
use netrunner_logger::{error, info};
use smoltcp::{iface::Config, phy::DeviceCapabilities};
use std::net::Ipv4Addr;
use tokio_util::sync::CancellationToken;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
netrunner_logger::Logger::init();
info!("Initializing NetRunner Stack...");
let mut dns_handler = DnsHandler::new(".");
if let Err(e) = dns_handler.init().await {
error!("Failed to initialize DNS blocklist: {}", e);
}
let remote_address = "147.45.43.70:443";
// 1. Создаем TUN интерфейс (зависит от платформы, поэтому делаем тут)
let tun_device = Tun::create(|config| {
config
.tun_name("netr0")
@@ -32,51 +19,37 @@ async fn main() -> anyhow::Result<()> {
})
.expect("Failed to initialize TUN device");
let remote_address: String = "147.45.43.70:443".into();
let _ = setup_platform_routing(&remote_address);
info!("TUN interface is UP: 10.0.0.1/24");
let config = Config::new(smoltcp::wire::HardwareAddress::Ip);
let mut caps = DeviceCapabilities::default();
caps.max_transmission_unit = 1350;
caps.medium = smoltcp::phy::Medium::Ip;
// 2. Собираем движок через наш новый Builder
let builder_result = EngineBuilder::new(remote_address)
.with_cache_path(".")
.with_tun(tun_device)
.build()
.await;
let network_token = CancellationToken::new();
// 3. Обрабатываем результат и запускаем цикл
match builder_result {
Ok((mut engine, tun)) => {
info!("Engine starting process loop...");
info!("Establishing secure tunnel to proxy server...");
let ctrl_c = tokio::signal::ctrl_c();
let muxer = match ClientHandler::connect(&remote_address).await {
Ok(m) => m,
Err(e) => {
error!("Failed to establish secure tunnel to server: {}", e);
return Err(anyhow::anyhow!("Failed to establish secure tunnel: {}", e));
tokio::select! {
res = engine.run(tun) => {
error!("Engine loop error: {:?}", res);
},
_ = ctrl_c => {
info!("Ctrl+C received, shutting down...");
}
}
}
};
info!("Secure tunnel established, Muxer is ready.");
let mut engine = Engine::new(config, caps, dns_handler, muxer);
engine.set_any_ip(true);
engine.set_transparent_mode();
engine.set_default_gateway(Ipv4Addr::new(10, 0, 0, 2));
engine.activate();
info!("Stack IP initialized: 10.0.0.2");
info!("Engine starting process loop...");
let ctrl_c = tokio::signal::ctrl_c();
tokio::select! {
res = engine.run(tun_device) => {
error!("Engine loop error: {:?}", res);
},
_ = ctrl_c => {
info!("Ctrl+C received, shutting down...");
Err(e) => {
error!("Failed to build VPN Engine: {}", e);
}
}
network_token.cancel();
// 4. Очистка системных роутов при любом сценарии выхода (ошибка или Ctrl+C)
info!("Restoring system routing...");
let addr: std::net::SocketAddr = remote_address.parse().expect("Invalid address format");
let p_ip = addr.ip().to_string();
@@ -85,7 +58,7 @@ async fn main() -> anyhow::Result<()> {
error!("Failed to restore routing: {}", e);
} else {
info!("System routing restored successfully.");
};
}
Ok(())
}
client/src/tun/connection_manager.rs
+282 -304
View File
@@ -17,92 +17,270 @@ use crate::connections::{
udp_connection::UdpConnection,
};
pub struct ConnectionManager {
// ============================================================================
// 1. УПРАВЛЕНИЕ СОСТОЯНИЕМ СЕССИЙ (SessionTracker)
// ============================================================================
struct SessionTracker {
last_activity: HashMap<SocketHandle, StdInstant>,
active_tcp_sessions: HashMap<SocketHandle, TcpConnection>,
active_udp_sessions: HashMap<SocketHandle, UdpConnection>,
active_tcp: HashMap<SocketHandle, TcpConnection>,
active_udp: HashMap<SocketHandle, UdpConnection>,
failed_until: HashMap<SocketHandle, StdInstant>,
to_remove: Vec<SocketHandle>,
}
impl SessionTracker {
fn new() -> Self {
Self {
last_activity: HashMap::new(),
active_tcp: HashMap::new(),
active_udp: HashMap::new(),
failed_until: HashMap::new(),
to_remove: Vec::new(),
}
}
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 from SocketSet and internal maps");
socket_set.remove(handle);
self.last_activity.remove(&handle);
self.failed_until.remove(&handle);
self.active_tcp.remove(&handle);
self.active_udp.remove(&handle);
}
}
fn has_tcp(
&self,
dst_addr: smoltcp::wire::IpAddress,
dst_port: u16,
socket_set: &SocketSet,
) -> bool {
socket_set.iter().any(|(_, s)| {
if let Some(tcp) = tcp::Socket::downcast(s) {
if let Some(ep) = tcp.local_endpoint() {
return ep.addr == dst_addr && ep.port == dst_port;
}
}
false
})
}
fn has_udp(
&self,
dst_addr: smoltcp::wire::IpAddress,
dst_port: u16,
socket_set: &SocketSet,
) -> bool {
socket_set.iter().any(|(_, s)| {
if let Some(udp) = udp::Socket::downcast(s) {
let ep = udp.endpoint();
return ep.addr == Some(dst_addr) && ep.port == dst_port;
}
false
})
}
}
// ============================================================================
// 2. РАЗРЕШЕНИЕ АДРЕСОВ И DNS (TargetResolver)
// ============================================================================
struct TargetResolver {
dns_handler: DnsHandler,
fake_ip_store: FakeIpStore,
failed_until: HashMap<SocketHandle, StdInstant>,
sockets_to_remove: Vec<SocketHandle>,
}
impl TargetResolver {
fn new(dns_handler: DnsHandler) -> Self {
Self {
dns_handler,
fake_ip_store: FakeIpStore::new(),
}
}
fn resolve_tcp(&self, socket: &tcp::Socket) -> TargetAddress {
let ep = match socket.local_endpoint() {
Some(ep) => ep,
None => {
warn!(handle=?socket, "Target resolution failed: no local endpoint");
return TargetAddress::Domain("disconnected".to_string(), 0);
}
};
match ep.addr {
smoltcp::wire::IpAddress::Ipv4(ip) => {
let std_ip = std::net::Ipv4Addr::from(ip);
if let Some(domain) = self.fake_ip_store.lookup_by_ip(&std_ip) {
TargetAddress::Domain(domain, ep.port)
} else {
TargetAddress::Ipv4(std_ip, ep.port)
}
}
smoltcp::wire::IpAddress::Ipv6(ip) => {
TargetAddress::Ipv6(std::net::Ipv6Addr::from(ip), ep.port)
}
}
}
fn process_dns_query(&mut self, data: &[u8]) -> Option<Vec<u8>> {
self.dns_handler.handle_query(data, &mut self.fake_ip_store)
}
}
// ============================================================================
// 3. ФАБРИКА СОКЕТОВ (SocketFactory)
// ============================================================================
struct SocketFactory;
impl SocketFactory {
fn create_tcp<'a>(port: u16) -> tcp::Socket<'a> {
let buf_size = match port {
443 | 80 => 1024 * 1024 * 2,
22 => 32 * 1024,
53 => 16 * 1024,
_ => 128 * 1024,
};
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 (buf_size, packet_count) = match port {
443 => (512 * 1024, 390),
53 => (64 * 1024, 32),
_ => (128 * 1024, 100),
};
udp::Socket::new(
udp::PacketBuffer::new(
vec![udp::PacketMetadata::EMPTY; packet_count],
vec![0; buf_size],
),
udp::PacketBuffer::new(
vec![udp::PacketMetadata::EMPTY; packet_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]),
)
}
}
// ============================================================================
// 4. ГЛАВНЫЙ КООРДИНАТОР (ConnectionManager)
// ============================================================================
pub struct ConnectionManager {
tracker: SessionTracker,
resolver: TargetResolver,
muxer: Muxer,
}
impl ConnectionManager {
pub fn new(dns_handler: DnsHandler, muxer: Muxer) -> Self {
Self {
last_activity: HashMap::new(),
active_tcp_sessions: HashMap::new(),
active_udp_sessions: HashMap::new(),
fake_ip_store: FakeIpStore::new(),
failed_until: HashMap::new(),
dns_handler,
sockets_to_remove: Vec::new(),
tracker: SessionTracker::new(),
resolver: TargetResolver::new(dns_handler),
muxer,
}
}
pub fn setup_sockets(n_icmp: usize) -> SocketSet<'static> {
let mut sockets = SocketSet::new(Vec::with_capacity(48));
for _ in 0..n_icmp {
sockets.add(SocketFactory::create_icmp());
}
sockets
}
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 endpoint = IpListenEndpoint {
let _ = tcp.listen(IpListenEndpoint {
addr: None,
port: 443,
};
let _ = tcp.listen(endpoint);
});
}
} else if let Some(udp) = udp::Socket::downcast_mut(socket) {
if !udp.is_open() {
let endpoint = IpListenEndpoint {
let _ = udp.bind(IpListenEndpoint {
addr: None,
port: 53,
};
match udp.bind(endpoint) {
Ok(_) => debug!("UDP socket bound to port 53"),
Err(e) => warn!(error=?e, "Failed to bind UDP socket"),
}
});
}
}
}
}
fn resolve_target(&self, socket: &tcp::Socket) -> TargetAddress {
let local_endpoint = match socket.local_endpoint() {
Some(ep) => ep,
None => {
warn!(handle=?socket, "Target resolution failed: socket disconnected or no local endpoint");
return TargetAddress::Domain("disconnected".to_string(), 0);
}
pub fn try_create_socket_from_packet(&mut self, packet: &[u8], socket_set: &mut SocketSet) {
let Ok(ip_packet) = Ipv4Packet::new_checked(packet) else {
return;
};
let port = local_endpoint.port;
let ip = local_endpoint.addr;
match ip_packet.next_header() {
IpProtocol::Tcp => {
if let Ok(tcp_packet) = TcpPacket::new_checked(ip_packet.payload()) {
if tcp_packet.syn() && !tcp_packet.ack() {
let dst_port = tcp_packet.dst_port();
let dst_addr = ip_packet.dst_addr();
match ip {
smoltcp::wire::IpAddress::Ipv4(ipv4_addr) => {
let std_ip = std::net::Ipv4Addr::from(ipv4_addr);
let start = StdInstant::now();
if let Some(domain) = self.fake_ip_store.lookup_by_ip(&std_ip) {
debug!(ip=%std_ip, domain=%domain, elapsed=?start.elapsed(), "FakeIP match found");
TargetAddress::Domain(domain, port)
} else {
warn!(ip=%std_ip, port=%port, "IP not found in FakeIpStore. Using raw IP (possible DNS lag)");
TargetAddress::Ipv4(std_ip, port)
if !self.tracker.has_tcp(dst_addr.into(), dst_port, socket_set) {
let mut socket = SocketFactory::create_tcp(dst_port);
let endpoint = IpListenEndpoint {
addr: Some(dst_addr.into()),
port: dst_port,
};
if socket.listen(endpoint).is_ok() {
socket_set.add(socket);
}
}
}
}
}
smoltcp::wire::IpAddress::Ipv6(ipv6_addr) => {
let std_ip = std::net::Ipv6Addr::from(ipv6_addr);
debug!(ip=%std_ip, port=%port, "Using raw IPv6 target");
TargetAddress::Ipv6(std_ip, port)
IpProtocol::Udp => {
if let Ok(udp_packet) = UdpPacket::new_checked(ip_packet.payload()) {
let dst_port = udp_packet.dst_port();
if dst_port == 0 || dst_port == 137 || dst_port == 138 {
return;
}
let dst_addr = ip_packet.dst_addr();
if !self.tracker.has_udp(dst_addr.into(), dst_port, socket_set) {
let mut socket = SocketFactory::create_udp(dst_port);
let endpoint = IpListenEndpoint {
addr: Some(dst_addr.into()),
port: dst_port,
};
if socket.bind(endpoint).is_ok() {
socket_set.add(socket);
}
}
}
}
_ => {}
}
}
pub fn process_sockets(&mut self, socket_set: &mut SocketSet) {
if !self.active_tcp_sessions.is_empty() || !self.active_udp_sessions.is_empty() {}
for (handle, socket) in socket_set.iter_mut() {
if let Some(tcp) = tcp::Socket::downcast_mut(socket) {
self.handle_tcp(handle, tcp);
@@ -115,34 +293,27 @@ impl ConnectionManager {
}
fn handle_tcp(&mut self, handle: SocketHandle, socket: &mut tcp::Socket) {
use tcp::State;
if socket.state() == State::Closed {
if self.active_tcp_sessions.remove(&handle).is_some() {
info!(%handle, "TCP Session removed: local socket closed");
}
if !self.sockets_to_remove.contains(&handle) {
self.sockets_to_remove.push(handle);
}
if socket.state() == tcp::State::Closed {
self.tracker.active_tcp.remove(&handle);
self.tracker.queue_removal(handle);
return;
}
if socket.state() == State::Established && !self.active_tcp_sessions.contains_key(&handle) {
let start_establish = StdInstant::now();
let target = self.resolve_target(socket);
if let TargetAddress::Domain(d, _) = &target {
if socket.state() == tcp::State::Established
&& !self.tracker.active_tcp.contains_key(&handle)
{
let target = self.resolver.resolve_tcp(socket);
if let TargetAddress::Domain(ref d, _) = target {
if d == "disconnected" {
warn!(%handle, "Aborting socket: target unresolved");
socket.abort();
return;
}
}
info!(%handle, target = %target, "New TCP session: Local established, starting Muxer handshake");
info!(%handle, target = %target, "New TCP session established");
let (conn, mut rx_from_smol, tx_to_smol, handshake_tx) = TcpConnection::new(handle);
self.active_tcp_sessions.insert(handle, conn);
self.tracker.active_tcp.insert(handle, conn);
let muxer = self.muxer.clone();
let stream_id = muxer.next_id();
@@ -152,7 +323,6 @@ impl ConnectionManager {
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(CHANNEL_CAPACITY);
muxer.register_stream(stream_id, v_tx);
let send_start = StdInstant::now();
if muxer
.send_to_netwrok(MuxMessage {
stream_id,
@@ -162,24 +332,15 @@ impl ConnectionManager {
.await
.is_err()
{
error!(stream_id, "Muxer: Failed to send Connect frame");
muxer.remove_stream(stream_id);
return;
}
match tokio::time::timeout(Duration::from_secs(10), v_rx.recv()).await {
Ok(Some(data)) => {
if data.len() >= 2 && data[1] == 0x00 {
info!(stream_id, %handle, elapsed=?send_start.elapsed(), "Muxer: Handshake success");
let _ = handshake_tx.send(());
} else {
warn!(stream_id, "Muxer: Connection rejected by proxy server");
muxer.remove_stream(stream_id);
return;
}
Ok(Some(data)) if data.len() >= 2 && data[1] == 0x00 => {
let _ = handshake_tx.send(());
}
_ => {
error!(stream_id, "Muxer: Handshake timeout (10s)");
muxer.remove_stream(stream_id);
return;
}
@@ -209,10 +370,7 @@ impl ConnectionManager {
}
};
tokio::select! {
_ = to_proxy => debug!(stream_id, "Stream direction: SMOL -> Proxy closed"),
_ = from_proxy => debug!(stream_id, "Stream direction: Proxy -> SMOL closed"),
}
tokio::select! { _ = to_proxy => {}, _ = from_proxy => {} }
let _ = muxer
.send_to_netwrok(MuxMessage {
@@ -222,63 +380,50 @@ impl ConnectionManager {
})
.await;
muxer.remove_stream(stream_id);
debug!(stream_id, "Muxer stream fully removed");
});
}
if let Some(conn) = self.active_tcp_sessions.get_mut(&handle) {
if let Some(conn) = self.tracker.active_tcp.get_mut(&handle) {
if !conn.tick(socket) {
debug!(%handle, "TCP Connection tick failed (buffer full or error), aborting");
socket.abort();
}
}
if socket.state() == State::CloseWait {
debug!(%handle, "Socket in CloseWait, closing locally");
if socket.state() == tcp::State::CloseWait {
socket.close();
}
}
fn handle_udp(&mut self, handle: SocketHandle, socket: &mut udp::Socket) {
self.last_activity.insert(handle, StdInstant::now());
self.tracker.last_activity.insert(handle, StdInstant::now());
let local_port = socket.endpoint().port;
if local_port == 53 {
if socket.endpoint().port == 53 {
while socket.can_recv() {
let (data, meta) = match socket.recv() {
Ok(res) => res,
Err(_) => break,
};
if let Some(response) = self.dns_handler.handle_query(data, &mut self.fake_ip_store)
{
netrunner_logger::debug!(to = %meta.endpoint, "Sending DNS response (FakeIP/Filtered)");
let _ = socket.send_slice(&response, meta);
if let Ok((data, meta)) = socket.recv() {
if let Some(response) = self.resolver.process_dns_query(data) {
let _ = socket.send_slice(&response, meta);
}
} else {
break;
}
}
return;
}
if socket.is_open() && !self.active_udp_sessions.contains_key(&handle) {
let endpoint = socket.endpoint();
let target = match endpoint.addr {
Some(smoltcp::wire::IpAddress::Ipv4(ipv4_addr)) => {
TargetAddress::Ipv4(std::net::Ipv4Addr::from(ipv4_addr), endpoint.port)
if socket.is_open() && !self.tracker.active_udp.contains_key(&handle) {
let ep = socket.endpoint();
let target = match ep.addr {
Some(smoltcp::wire::IpAddress::Ipv4(ip)) => {
TargetAddress::Ipv4(std::net::Ipv4Addr::from(ip), ep.port)
}
Some(smoltcp::wire::IpAddress::Ipv6(ipv6_addr)) => {
TargetAddress::Ipv6(std::net::Ipv6Addr::from(ipv6_addr), endpoint.port)
}
None => {
netrunner_logger::warn!(%handle, "UDP socket endpoint has no IP address bound");
return;
Some(smoltcp::wire::IpAddress::Ipv6(ip)) => {
TargetAddress::Ipv6(std::net::Ipv6Addr::from(ip), ep.port)
}
None => return,
};
netrunner_logger::info!(%handle, target = %target, "New UDP proxied session established");
let (conn, mut rx_from_smol, tx_to_smol) = UdpConnection::new(handle);
self.active_udp_sessions.insert(handle, conn);
self.tracker.active_udp.insert(handle, conn);
let muxer = self.muxer.clone();
let stream_id = muxer.next_id();
@@ -298,12 +443,15 @@ impl ConnectionManager {
let to_proxy = async {
while let Some(data) = rx_from_smol.recv().await {
let msg = MuxMessage {
stream_id,
frame_type: FrameType::UdpData,
data,
};
if muxer.send_to_netwrok(msg).await.is_err() {
if muxer
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::UdpData,
data,
})
.await
.is_err()
{
break;
}
}
@@ -311,19 +459,13 @@ impl ConnectionManager {
let from_proxy = async {
while let Some(data) = v_rx.recv().await {
if data.is_empty() {
break;
}
if tx_to_smol.send(data).await.is_err() {
if data.is_empty() || tx_to_smol.send(data).await.is_err() {
break;
}
}
};
tokio::select! {
_ = to_proxy => {}
_ = from_proxy => {}
}
tokio::select! { _ = to_proxy => {}, _ = from_proxy => {} }
let _ = muxer
.send_to_netwrok(MuxMessage {
@@ -336,204 +478,40 @@ impl ConnectionManager {
});
}
if let Some(conn) = self.active_udp_sessions.get_mut(&handle) {
if let Some(conn) = self.tracker.active_udp.get_mut(&handle) {
if !conn.tick(socket) {
self.sockets_to_remove.push(handle);
self.active_udp_sessions.remove(&handle);
self.tracker.queue_removal(handle);
self.tracker.active_udp.remove(&handle);
}
}
}
fn handle_icmp(&mut self, handle: SocketHandle, socket: &mut icmp::Socket) {
fn handle_icmp(&mut self, _handle: SocketHandle, socket: &mut icmp::Socket) {
if socket.can_recv() {
match socket.recv() {
Ok((data, endpoint)) => {
debug!(handle=%handle, from=?endpoint, data=?data, "ICMP: packet received");
}
Err(_) => {}
}
let _ = socket.recv();
}
}
fn create_dynamic_tcp_socket<'a>(port: u16) -> tcp::Socket<'a> {
let buf_size = match port {
443 | 80 => 1024 * 1024 * 2,
22 => 32 * 1024,
53 => 16 * 1024,
_ => 128 * 1024,
};
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
}
pub fn try_create_socket_from_packet(&mut self, packet: &[u8], socket_set: &mut SocketSet) {
let Ok(ip_packet) = Ipv4Packet::new_checked(packet) else {
return;
};
match ip_packet.next_header() {
IpProtocol::Tcp => {
let Ok(tcp_packet) = TcpPacket::new_checked(ip_packet.payload()) else {
return;
};
if tcp_packet.syn() && !tcp_packet.ack() {
let dst_port = tcp_packet.dst_port();
let dst_addr = ip_packet.dst_addr();
if !self.has_active_tcp_session(socket_set, dst_addr.into(), dst_port) {
let start_create = StdInstant::now();
let mut socket = Self::create_dynamic_tcp_socket(dst_port);
let endpoint = IpListenEndpoint {
addr: Some(dst_addr.into()),
port: dst_port,
};
if socket.listen(endpoint).is_ok() {
let handle = socket_set.add(socket);
debug!(%handle, "Dynamic TCP: Created for {}:{} in {:?}", dst_addr, dst_port, start_create.elapsed());
} else {
warn!("Dynamic TCP: Failed to listen on {}:{}", dst_addr, dst_port);
}
}
}
}
IpProtocol::Udp => {
let Ok(udp_packet) = UdpPacket::new_checked(ip_packet.payload()) else {
return;
};
let dst_port = udp_packet.dst_port();
let dst_addr = ip_packet.dst_addr();
if dst_port == 0 || dst_port == 137 || dst_port == 138 {
return;
}
if !self.has_active_udp_session(socket_set, dst_addr.into(), dst_port) {
let start_create = StdInstant::now();
let mut socket = Self::create_dynamic_udp_socket(dst_port);
let endpoint = IpListenEndpoint {
addr: Some(dst_addr.into()),
port: dst_port,
};
if socket.bind(endpoint).is_ok() {
let handle = socket_set.add(socket);
debug!(%handle, "Dynamic UDP: Created for {}:{} in {:?}", dst_addr, dst_port, start_create.elapsed());
}
}
}
_ => {}
}
}
fn create_dynamic_udp_socket<'a>(port: u16) -> udp::Socket<'a> {
let (buf_size, packet_count) = match port {
443 => (512 * 1024, 390),
53 => (64 * 1024, 32),
_ => (128 * 1024, 100),
};
udp::Socket::new(
udp::PacketBuffer::new(
vec![udp::PacketMetadata::EMPTY; packet_count],
vec![0; buf_size],
),
udp::PacketBuffer::new(
vec![udp::PacketMetadata::EMPTY; packet_count],
vec![0; buf_size],
),
)
}
fn create_icmp_socket<'a>() -> icmp::Socket<'a> {
let icmp_rx_buffer =
icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; 8], vec![0; 2048]);
let icmp_tx_buffer =
icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; 8], vec![0; 2048]);
icmp::Socket::new(icmp_rx_buffer, icmp_tx_buffer)
}
fn has_active_tcp_session(
&self,
socket_set: &SocketSet,
dst_addr: smoltcp::wire::IpAddress,
dst_port: u16,
) -> bool {
for (_, socket) in socket_set.iter() {
if let Some(tcp) = tcp::Socket::downcast(socket) {
if let Some(endpoint) = tcp.local_endpoint() {
if endpoint.addr == dst_addr && endpoint.port == dst_port {
return true;
}
}
}
}
false
}
fn has_active_udp_session(
&self,
socket_set: &SocketSet,
dst_addr: smoltcp::wire::IpAddress,
dst_port: u16,
) -> bool {
for (_, socket) in socket_set.iter() {
if let Some(udp) = udp::Socket::downcast(socket) {
let endpoint = udp.endpoint();
if endpoint.addr == Some(dst_addr) && endpoint.port == dst_port {
return true;
}
}
}
false
}
pub fn cleanup(&mut self, socket_set: &mut SocketSet) {
for handle in self.sockets_to_remove.drain(..) {
debug!(%handle, "Cleanup: Removing socket from SocketSet and internal maps");
socket_set.remove(handle);
self.last_activity.remove(&handle);
self.failed_until.remove(&handle);
self.active_tcp_sessions.remove(&handle);
self.active_udp_sessions.remove(&handle);
}
}
pub fn setup_sockets(n_icmp: usize) -> SocketSet<'static> {
let mut sockets = SocketSet::new(Vec::with_capacity(48));
for _ in 0..n_icmp {
sockets.add(Self::create_icmp_socket());
}
sockets
self.tracker.cleanup(socket_set);
}
pub fn log_status(&self, socket_set: &SocketSet) {
let mut established = 0;
let mut total_tcp = 0;
let mut est = 0;
let mut total = 0;
for (_, socket) in socket_set.iter() {
if let Some(tcp) = tcp::Socket::downcast(socket) {
total_tcp += 1;
total += 1;
if tcp.state() == tcp::State::Established {
established += 1;
est += 1;
}
}
}
debug!(
"TCP Stats: Total_Sockets={}, Established={}, Active_Sessions={}",
total_tcp,
established,
self.active_tcp_sessions.len()
"TCP Stats: Total={}, Established={}, Active={}",
total,
est,
self.tracker.active_tcp.len()
);
}
}