242 lines
9.1 KiB
Rust
242 lines
9.1 KiB
Rust
//! Реестр виртуальных соединений userspace-стека smoltcp.
|
||
//!
|
||
//! [`SessionTracker`] — это «NAT-таблица» клиента: он связывает хендлы сокетов
|
||
//! smoltcp ([`SocketHandle`]) с логическими id, держит активные TCP/UDP-соединения
|
||
//! и каналы доставки входящих данных, следит за активностью и убирает «призраков».
|
||
//!
|
||
//! Ключевые обязанности:
|
||
//! - **Реестр** TCP/UDP-соединений и двусторонний маппинг `handle ⇄ id`.
|
||
//! - **Pending TCP** — полуоткрытые соединения с удерживаемым семафор-пермитом
|
||
//! (ограничение числа одновременных установок).
|
||
//! - **Idle-выметание** ([`enforce_idle_timeouts`](SessionTracker::enforce_idle_timeouts))
|
||
//! и **LRU-эвикт** ([`evict_oldest_socket`](SessionTracker::evict_oldest_socket))
|
||
//! при достижении лимита сокетов — с защитой системных/слушающих сокетов.
|
||
//! - **Отложенное удаление**: `queue_removal` + `cleanup` (нельзя трогать
|
||
//! `SocketSet` во время итерации по нему).
|
||
|
||
use std::{
|
||
collections::HashMap,
|
||
time::{Duration, Instant as StdInstant},
|
||
};
|
||
|
||
use bytes::Bytes;
|
||
use dashmap::DashMap;
|
||
use smoltcp::{
|
||
iface::{SocketHandle, SocketSet},
|
||
socket::Socket,
|
||
wire::IpAddress,
|
||
};
|
||
use std::sync::Arc;
|
||
use tokio::sync::{OwnedSemaphorePermit, mpsc};
|
||
|
||
use crate::net::connection::{TcpConnection, UdpConnection};
|
||
|
||
/// Состояние всех виртуальных соединений и их маппингов на хендлы smoltcp.
|
||
pub struct SessionTracker {
|
||
last_activity: HashMap<SocketHandle, StdInstant>,
|
||
active_tcp: HashMap<SocketHandle, TcpConnection>,
|
||
active_udp: HashMap<SocketHandle, UdpConnection>,
|
||
pub inbound_tx: Arc<DashMap<u64, (mpsc::Sender<Bytes>, Arc<std::sync::atomic::AtomicBool>)>>,
|
||
handle_to_id: HashMap<SocketHandle, u64>,
|
||
|
||
id_to_handle: HashMap<u64, SocketHandle>,
|
||
pending_tcp: HashMap<SocketHandle, (StdInstant, OwnedSemaphorePermit)>,
|
||
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: Arc::new(DashMap::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, permit: OwnedSemaphorePermit) {
|
||
self.pending_tcp.insert(handle, (StdInstant::now(), permit));
|
||
self.last_activity.insert(handle, StdInstant::now());
|
||
}
|
||
|
||
pub fn pop_pending_permit(&mut self, handle: SocketHandle) -> Option<OwnedSemaphorePermit> {
|
||
self.pending_tcp.remove(&handle).map(|(_, permit)| permit)
|
||
}
|
||
|
||
pub fn register_tcp(
|
||
&mut self,
|
||
handle: SocketHandle,
|
||
id: u64,
|
||
conn: TcpConnection,
|
||
tx: mpsc::Sender<Bytes>,
|
||
is_saturated: Arc<std::sync::atomic::AtomicBool>,
|
||
) {
|
||
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, is_saturated));
|
||
self.last_activity.insert(handle, StdInstant::now());
|
||
}
|
||
|
||
pub fn register_udp(
|
||
&mut self,
|
||
handle: SocketHandle,
|
||
id: u64,
|
||
conn: UdpConnection,
|
||
tx: mpsc::Sender<Bytes>,
|
||
is_saturated: Arc<std::sync::atomic::AtomicBool>,
|
||
) {
|
||
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, is_saturated));
|
||
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 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 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);
|
||
}
|
||
}
|
||
|
||
/// LRU-эвикт для освобождения слота при достижении лимита сокетов.
|
||
///
|
||
/// Сначала ищет уже «мёртвый» TCP-сокет (Closed/TimeWait/CloseWait/FinWait);
|
||
/// если таких нет — закрывает самый давно неактивный пользовательский сокет,
|
||
/// **не трогая** системные/слушающие. Возвращает `true`, если кого-то закрыл.
|
||
pub fn evict_oldest_socket(&mut self, socket_set: &mut SocketSet) -> bool {
|
||
let mut victim = None;
|
||
|
||
for (handle, socket) in socket_set.iter() {
|
||
if let smoltcp::socket::Socket::Tcp(tcp_socket) = socket {
|
||
let state = tcp_socket.state();
|
||
if matches!(
|
||
state,
|
||
smoltcp::socket::tcp::State::Closed
|
||
| smoltcp::socket::tcp::State::TimeWait
|
||
| smoltcp::socket::tcp::State::CloseWait
|
||
| smoltcp::socket::tcp::State::FinWait1
|
||
| smoltcp::socket::tcp::State::FinWait2
|
||
) {
|
||
victim = Some(handle);
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
|
||
if victim.is_none() {
|
||
// 🔥 ИСПРАВЛЕНИЕ: Защищаем системные (слушающие) сокеты от удаления
|
||
victim = self
|
||
.last_activity
|
||
.iter()
|
||
.filter(|(h, _)| {
|
||
self.active_tcp.contains_key(h)
|
||
|| self.active_udp.contains_key(h)
|
||
|| self.pending_tcp.contains_key(h)
|
||
})
|
||
.min_by_key(|&(_, &time)| time)
|
||
.map(|(&handle, _)| handle);
|
||
}
|
||
|
||
if let Some(handle) = victim {
|
||
netrunner_logger::info!(
|
||
"🔪 LRU Eviction: Force closing socket {:?} to free up a slot",
|
||
handle
|
||
);
|
||
self.queue_removal(handle);
|
||
self.cleanup(socket_set);
|
||
true
|
||
} else {
|
||
false
|
||
}
|
||
}
|
||
|
||
/// Фактически удаляет все сокеты из очереди `to_remove` из `SocketSet` и всех
|
||
/// внутренних таблиц. Вызывается вне итерации по сокетам (см. отложенность).
|
||
pub fn cleanup(&mut self, socket_set: &mut SocketSet) {
|
||
for handle in self.to_remove.drain(..) {
|
||
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);
|
||
}
|
||
}
|
||
}
|
||
|
||
pub fn get_app_buffer_info(&self, handle: SocketHandle) -> usize {
|
||
self.active_tcp
|
||
.get(&handle)
|
||
.map(|conn| conn.app_pending_out_size())
|
||
.unwrap_or(0)
|
||
}
|
||
}
|