close to finish a tun/smoltcp

This commit is contained in:
2026-03-07 20:17:49 +07:00
parent 4a789dd128
commit b1cc4c14f4
14 changed files with 912 additions and 279 deletions
+102 -142
View File
@@ -1,19 +1,26 @@
use netrunner_common::protocol::codec::socks::TargetAddress;
use smoltcp::{
iface::{SocketHandle, SocketSet},
socket::{AnySocket, icmp, tcp, udp},
wire::IpListenEndpoint,
wire::{IpAddress, IpListenEndpoint},
};
use std::{
collections::HashMap,
time::{Duration, Instant as StdInstant},
};
use tracing::{debug, info};
use tracing::{debug, info, warn};
use crate::{connections::tcp_connection::TcpConnection, tun::engine::START_TIME};
use crate::{
connections::{
ip_store::FakeIpStore, tcp_connection::TcpConnection, udp_connection::UdpConnection,
},
tun::engine::START_TIME,
};
pub struct ConnectionManager {
last_activity: HashMap<SocketHandle, StdInstant>,
active_tcp_sessions: HashMap<SocketHandle, TcpConnection>,
//active_udp_sessions: HashMap<SocketHandle, UdpSession>,
fake_ip_store: FakeIpStore,
proxy_ip: String,
}
@@ -23,79 +30,113 @@ impl ConnectionManager {
last_activity: HashMap::new(),
active_tcp_sessions: HashMap::new(),
proxy_ip: ip,
fake_ip_store: FakeIpStore::new(),
}
}
/// Основной метод, который вызывается в цикле Engine
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 {
addr: None,
port: 443,
};
match tcp.listen(endpoint) {
Ok(_) => debug!("Socket is now listening"),
Err(e) => warn!(error=?e, "Failed to listen on socket"),
}
}
}
}
}
fn resolve_target(&self, socket: &tcp::Socket) -> TargetAddress {
// Безопасно получаем эндпоинт
let remote_endpoint = match socket.remote_endpoint() {
Some(ep) => ep,
None => {
warn!(handle=?socket, "Attempted to resolve target for an unconnected socket");
// Возвращаем дефолт, чтобы не падать
return TargetAddress::Domain("disconnected".to_string(), 0);
}
};
let port = remote_endpoint.port;
let ip = remote_endpoint.addr;
match ip {
smoltcp::wire::IpAddress::Ipv4(ipv4_addr) => {
let std_ip = std::net::Ipv4Addr::from(ipv4_addr);
// Ищем домен в нашем FakeIpStore
if let Some(domain) = self.fake_ip_store.lookup_by_ip(&std_ip) {
debug!(target=%domain, port=%port, "Resolved fake IP to domain");
return TargetAddress::Domain(domain, port);
}
debug!(ip=%std_ip, port=%port, "Using raw IP target");
TargetAddress::Ipv4(std_ip, port)
}
smoltcp::wire::IpAddress::Ipv6(ipv6_addr) => {
let std_ip = std::net::Ipv6Addr::from(ipv6_addr);
debug!(ip=%std_ip, port=%port, "Using IPv6 target");
TargetAddress::Ipv6(std_ip, port)
}
}
}
pub fn process_sockets(&mut self, socket_set: &mut SocketSet) {
for (handle, socket) in socket_set.iter_mut() {
// 1. Пытаемся даункастить сокет до TCP
if let Some(tcp) = tcp::Socket::downcast_mut(socket) {
self.handle_tcp(handle, tcp);
continue;
}
// 2. До UDP
if let Some(udp) = udp::Socket::downcast_mut(socket) {
} else if let Some(udp) = udp::Socket::downcast_mut(socket) {
self.handle_udp(handle, udp);
continue;
}
// 3. До ICMP
if let Some(icmp) = icmp::Socket::downcast_mut(socket) {
} else if let Some(icmp) = icmp::Socket::downcast_mut(socket) {
self.handle_icmp(handle, icmp);
continue;
}
}
}
fn handle_tcp(&mut self, handle: SocketHandle, socket: &mut tcp::Socket) {
if socket.state() == tcp::State::Established {
let proxy_ip = self.proxy_ip.clone(); // Берем из конфига менеджера
let target = self.resolve_target(socket);
debug!(
handle=%handle,
local=?socket.local_endpoint(),
remote=?socket.remote_endpoint(),
"Socket details"
);
let conn = self.active_tcp_sessions.entry(handle).or_insert_with(|| {
// 1. Получаем endpoint и безопасно распаковываем его
let endpoint = socket
.remote_endpoint()
.expect("TCP socket in Established state must have a remote endpoint");
let mut conn = TcpConnection::new(handle, self.proxy_ip.clone(), target);
// 2. Конвертируем smoltcp::wire::IpAddress в std::net::IpAddr
let ip: std::net::IpAddr = match endpoint.addr {
smoltcp::wire::IpAddress::Ipv4(v4) => std::net::IpAddr::V4(v4.into()),
smoltcp::wire::IpAddress::Ipv6(v6) => std::net::IpAddr::V6(v6.into()),
};
conn.tick(socket);
// 3. Собираем финальный SocketAddr
let target_addr = std::net::SocketAddr::new(ip, endpoint.port);
info!(handle=%handle, target=%target_addr, "Creating new TcpConnection bridge");
TcpConnection::new(handle, proxy_ip, target_addr)
});
conn.poll_and_process(socket);
if conn.is_active() {
self.active_tcp_sessions.insert(handle, conn);
} else {
socket.abort();
}
}
if socket.state() == tcp::State::Closed {
self.active_tcp_sessions.remove(&handle);
if let Some(conn) = self.active_tcp_sessions.get_mut(&handle) {
if !conn.tick(socket) {
self.active_tcp_sessions.remove(&handle);
}
}
}
fn handle_udp(&mut self, handle: SocketHandle, socket: &mut udp::Socket) {
if socket.can_recv() {
match socket.recv() {
Ok((data, endpoint)) => {
info!(handle=%handle, from=?endpoint, len=data.len(), "UDP: packet received");
// МОК: Обработка UDP датаграммы
}
Err(_) => {}
}
}
self.last_activity.insert(handle, StdInstant::now());
UdpConnection::process_incoming(socket, &mut self.fake_ip_store);
}
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, "ICMP: packet received");
// МОК: Ответ на пинг или обработка ошибок
debug!(handle=%handle, from=?endpoint, data=?data, "ICMP: packet received");
}
Err(_) => {}
}
@@ -126,106 +167,25 @@ impl ConnectionManager {
icmp::Socket::new(icmp_rx_buffer, icmp_tx_buffer)
}
pub fn refill_sockets(&mut self, socket_set: &mut SocketSet) {
self.prune_sockets(socket_set);
const TARGET_FREE_TCP: usize = 16;
const TARGET_FREE_UDP: usize = 8;
pub fn setup_sockets(n_tcp: usize, n_udp: usize, n_icmp: usize) -> SocketSet<'static> {
// Создаем хранилище с запасом на все типы сокетов
let mut sockets = SocketSet::new(Vec::with_capacity(n_tcp + n_udp + n_icmp));
for &port in &[80, 443, 8080] {
let current_port_sockets = socket_set
.iter()
.filter(|(_, s)| {
if let Some(tcp) = tcp::Socket::downcast(s) {
// Если эндпоинт есть - сверяем порт
if let Some(endpoint) = tcp.local_endpoint() {
return endpoint.port == port;
}
// КЛЮЧЕВОЙ МОМЕНТ:
// Если эндпоинта НЕТ, но сокет НЕ в состоянии CLOSED,
// или он только что был создан для прослушивания.
// В smoltcp после listen() сокет переходит в LISTEN,
// но local_endpoint может появиться чуть позже.
tcp.state() == tcp::State::Listen
} else {
false
}
})
.count();
if current_port_sockets < 5 {
let mut s = Self::create_tcp_socket();
let endpoint = IpListenEndpoint { addr: None, port };
// Попробуй сначала listen, а потом добавлять
if s.listen(endpoint).is_ok() {
socket_set.add(s);
info!(
"--- REAL ADD --- Port: {}, Sockets for this port: {}",
port,
current_port_sockets + 1
);
}
}
// 1. Добавляем TCP сокеты
for _ in 0..n_tcp {
sockets.add(Self::create_tcp_socket());
}
let udp_active = socket_set
.iter()
.filter(|(_, s)| udp::Socket::downcast(s).is_some())
.count();
let has_icmp = socket_set
.iter()
.any(|(_, s)| icmp::Socket::downcast(s).is_some());
if udp_active < TARGET_FREE_UDP {
let diff = TARGET_FREE_UDP - udp_active;
debug!("Refilling UDP pool: adding {} sockets", diff);
for _ in 0..diff {
let s = Self::create_udp_socket();
socket_set.add(s);
}
// 2. Добавляем UDP сокеты
for _ in 0..n_udp {
sockets.add(Self::create_udp_socket());
}
if !has_icmp {
debug!("Adding ICMP socket for echo requests");
let s = Self::create_icmp_socket();
socket_set.add(s);
}
}
fn prune_sockets(&mut self, socket_set: &mut SocketSet) {
let now = StdInstant::now();
let udp_timeout = Duration::from_secs(60); // 1 минута для UDP
let mut to_remove = Vec::new();
for (handle, socket) in socket_set.iter() {
if let Some(tcp) = tcp::Socket::downcast(socket) {
if tcp.state() == tcp::State::Closed {
to_remove.push(handle);
continue;
}
}
if let Some(udp) = udp::Socket::downcast(socket) {
if udp.endpoint().port == 0 {
continue;
}
let last = self.last_activity.get(&handle).unwrap_or(&START_TIME);
if now.duration_since(*last) > udp_timeout {
debug!(handle=%handle, "UDP socket timeout reached");
to_remove.push(handle);
}
}
if let Some(_icmp) = icmp::Socket::downcast(socket) {
continue;
}
// 3. Добавляем ICMP сокет
for _ in 0..n_icmp {
sockets.add(Self::create_icmp_socket());
}
for handle in to_remove {
socket_set.remove(handle);
self.last_activity.remove(&handle);
}
sockets
}
}
+2
View File
@@ -16,6 +16,7 @@ use smoltcp::{
time::Instant,
};
use tokio::sync::mpsc;
use tracing::trace;
pub struct VirtTunDevice {
capabilities: DeviceCapabilities,
@@ -63,6 +64,7 @@ impl Device for VirtTunDevice {
fn receive(&mut self, _timestamp: Instant) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
if let Ok(buffer) = self.in_buf.try_recv() {
trace!("--- TUN RX: {} bytes ---", buffer.len()); // Посмотри, что именно прилетает
let rx = Self::RxToken {
buffer,
phantom_device: PhantomData,
+35 -38
View File
@@ -1,4 +1,5 @@
use smoltcp::time::Instant;
use smoltcp::wire::{IpAddress, IpCidr};
use smoltcp::{
iface::{Config, Interface, SocketSet},
phy::DeviceCapabilities,
@@ -35,9 +36,10 @@ impl Engine {
pub fn new(config: Config, caps: DeviceCapabilities, ip: String) -> Self {
let now = Engine::current_time();
let (mut device, bridge_rx, bridge_tx, avail) = VirtTunDevice::new(caps);
let mut interface = Interface::new(config, &mut device, now);
let interface = Interface::new(config, &mut device, now);
let socket_set = SocketSet::new(vec![]);
let mut socket_set = ConnectionManager::setup_sockets(16, 4, 1);
let mut manager = ConnectionManager::new(ip);
Self {
interface,
socket_set,
@@ -45,11 +47,12 @@ impl Engine {
bridge_tx,
bridge_rx,
avail,
manager: ConnectionManager::new(ip),
manager,
}
}
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, так как он нам нужен только в одной задаче
@@ -58,25 +61,27 @@ impl Engine {
Self::spawn_engine_to_tun(writer, from_engine);
loop {
self.manager.refill_sockets(&mut self.socket_set);
self.poll();
self.manager.process_sockets(&mut self.socket_set);
if self.avail.load(std::sync::atomic::Ordering::Acquire) {
tokio::task::yield_now().await;
continue;
loop {
self.poll();
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;
}
self.poll_delay().await;
}
}
pub fn poll(&mut self) {
fn poll(&mut self) {
let now = Self::current_time();
// Передаем девайс и сокеты в интерфейс
self.interface
.poll(now, &mut self.device, &mut self.socket_set);
}
pub async fn poll_delay(&mut self) {
async fn poll_delay(&mut self) {
let timestamp = Self::current_time();
let delay = self.interface.poll_delay(timestamp, &self.socket_set);
let sleep_duration = match delay {
@@ -99,6 +104,13 @@ impl Engine {
break;
}
if n >= 20 {
let dst_ip = std::net::Ipv4Addr::new(buf[16], buf[17], buf[18], buf[19]);
let dst_port = u16::from_be_bytes([buf[22], buf[23]]); // Порт для TCP
debug!("Сырой пакет: назначение {}:{}", dst_ip, dst_port);
}
let mut token = TokenBuffer::with_capacity(n);
token.extend_from_slice(&buf[..n]);
@@ -132,42 +144,20 @@ impl Engine {
Instant::from_micros(duration.as_micros() as i64)
}
pub fn add_address(&mut self, address: smoltcp::wire::IpCidr) {
self.interface.update_ip_addrs(|addrs| {
addrs
.push(address)
.expect("Failed to add IP: address list is full");
});
self.interface
.routes_mut()
.add_default_ipv4_route(smoltcp::wire::Ipv4Address::new(10, 0, 0, 1))
.ok();
}
pub fn set_any_ip(&mut self, state: bool) -> () {
self.interface.set_any_ip(state)
}
pub fn set_transparent_mode(&mut self) {
info!("Switching engine to transparent mode (AnyIP + 0.0.0.0/0)");
self.interface.set_any_ip(true);
self.interface.update_ip_addrs(|addrs| {
addrs.clear();
//accept all internet
addrs
.push(smoltcp::wire::IpCidr::new(
smoltcp::wire::IpAddress::v4(0, 0, 0, 0),
0,
))
.expect("Failed to set 0.0.0.0/0");
.push(IpCidr::new(IpAddress::v4(10, 0, 0, 2), 24))
.unwrap();
});
// Устанавливаем дефолтный маршрут "в никуда" (локально)
self.interface.routes_mut().remove_default_ipv4_route();
self.interface
.routes_mut()
.add_default_ipv4_route(smoltcp::wire::Ipv4Address::new(0, 0, 0, 0))
.ok();
}
pub fn set_default_gateway(&mut self, gateway: smoltcp::wire::Ipv4Address) {
@@ -178,4 +168,11 @@ impl Engine {
.add_default_ipv4_route(gateway)
.expect("Failed to set default gateway");
}
pub fn activate(&mut self) {
let now = Self::current_time();
self.interface
.poll(now, &mut self.device, &mut self.socket_set);
self.manager.start_listening(&mut self.socket_set);
}
}
+9 -27
View File
@@ -43,37 +43,19 @@ impl Tun {
Ok((writer, reader))
}
pub fn setup_linux_routes(&self) -> io::Result<()> {
pub fn setup_routing(&self) -> io::Result<()> {
use std::process::Command;
// Направляем трафик через IP, который слушает твой smoltcp Engine
// Предположим, smoltcp настроен на 10.0.0.2
let status = Command::new("ip")
.args(&["route", "add", "default", "via", "10.0.0.2", "dev", "tun0"])
.status()?;
// Теперь, когда мы знаем, что tun0 имеет IP 10.0.0.1,
// мы можем смело ставить его шлюзом.
if !status.success() {
return Err(io::Error::new(
io::ErrorKind::Other,
"Failed to setup routes",
));
}
Ok(())
}
// 1. Устанавливаем маршрут через созданный интерфейс tun0
let _ = Command::new("sudo")
.args(&[
"ip", "route", "add", "default", "via", "10.0.0.2", "dev", "tun0", "metric", "1",
])
.status();
pub fn clear_default_route(&self) -> io::Result<()> {
use std::process::Command;
info!("Removing existing default route...");
// Удаляем текущий маршрут по умолчанию
let status = Command::new("sudo")
.args(&["ip", "route", "del", "default"])
.status()?;
if !status.success() {
warn!("Could not delete default route (maybe it doesn't exist?)");
}
Ok(())
}
}