structure refactoring

This commit is contained in:
2026-03-24 15:32:52 +07:00
parent 7603d2c92e
commit fc93665fb2
14 changed files with 401 additions and 332 deletions
+292
View File
@@ -0,0 +1,292 @@
use crate::net::CHANNEL_CAPACITY;
use bytes::{Buf, Bytes, BytesMut};
use smoltcp::{
iface::SocketHandle,
socket::{tcp, udp},
wire::IpEndpoint,
};
use std::time::Duration;
use tokio::{
sync::{mpsc, oneshot},
time::Instant,
};
// ============================================================================
// 1. БАЗОВАЯ СТРУКТУРА (ConnectionCore)
// ============================================================================
/// Фундамент для любого соединения.
/// Инициализирует и хранит каналы связи между smoltcp и Muxer'ом.
pub struct ConnectionCore {
pub handle: SocketHandle,
pub tx: mpsc::Sender<Bytes>,
pub rx: mpsc::Receiver<Bytes>,
}
impl ConnectionCore {
pub fn new(handle: SocketHandle) -> (Self, mpsc::Receiver<Bytes>, mpsc::Sender<Bytes>) {
let (tx_to_net, rx_from_smol) = mpsc::channel::<Bytes>(CHANNEL_CAPACITY);
let (tx_to_smol, rx_from_net) = mpsc::channel::<Bytes>(CHANNEL_CAPACITY);
let core = Self {
handle,
tx: tx_to_net,
rx: rx_from_net,
};
(core, rx_from_smol, tx_to_smol)
}
}
// ============================================================================
// 2. TCP СОЕДИНЕНИЕ (TcpConnection)
// ============================================================================
pub enum ConnectionState {
Established,
Handshaking,
Active,
Closed,
}
const MAX_PENDING: usize = 64 * 1024;
const TCP_CHUNK_SIZE: usize = 1024 * 16;
pub struct TcpConnection {
core: ConnectionCore,
state: ConnectionState,
pending_data: BytesMut,
handshake_rx: Option<oneshot::Receiver<()>>,
}
impl TcpConnection {
pub fn new(
handle: SocketHandle,
) -> (
Self,
mpsc::Receiver<Bytes>,
mpsc::Sender<Bytes>,
oneshot::Sender<()>,
) {
let (core, rx_from_smol, tx_to_smol) = ConnectionCore::new(handle);
let (handshake_tx, handshake_rx) = oneshot::channel();
let conn = Self {
core,
state: ConnectionState::Handshaking,
pending_data: BytesMut::new(),
handshake_rx: Some(handshake_rx),
};
(conn, rx_from_smol, tx_to_smol, handshake_tx)
}
pub fn is_finished(&self, socket: &tcp::Socket) -> bool {
matches!(socket.state(), tcp::State::Closed | tcp::State::TimeWait)
}
pub fn _is_active(&self) -> bool {
matches!(self.state, ConnectionState::Active)
}
pub fn tick(&mut self, socket: &mut tcp::Socket) -> bool {
let state = socket.state();
match self.state {
ConnectionState::Handshaking => {
if let Some(rx) = &mut self.handshake_rx {
match rx.try_recv() {
Ok(_) => {
self.state = ConnectionState::Active;
self.handshake_rx = None;
return true;
}
Err(oneshot::error::TryRecvError::Empty) => return true,
Err(oneshot::error::TryRecvError::Closed) => {
self.state = ConnectionState::Closed;
return false;
}
}
} else {
return false;
}
}
ConnectionState::Active => {
self.poll_and_process(socket);
if state == tcp::State::CloseWait {
socket.close();
self.state = ConnectionState::Closed;
return false;
}
if self.is_finished(socket) {
self.state = ConnectionState::Closed;
socket.close();
return false;
}
}
ConnectionState::Closed => {
return false;
}
_ => {}
}
true
}
fn poll_and_process(&mut self, socket: &mut tcp::Socket) {
// 1. Вычитываем данные из smoltcp и шлем в Muxer
while socket.can_recv() {
let mut full = false;
let mut temp = [0u8; TCP_CHUNK_SIZE];
if let Ok(n) = socket.peek_slice(&mut temp) {
if n == 0 {
break;
}
let chunk = Bytes::copy_from_slice(&temp[..n]);
match self.core.tx.try_send(chunk) {
Ok(_) => {
socket.recv_slice(&mut temp[..n]).unwrap();
}
Err(mpsc::error::TrySendError::Full(_)) => {
full = true;
}
Err(_) => {
self.state = ConnectionState::Closed;
return;
}
}
} else {
break;
}
if full {
break;
}
}
// 2. Читаем данные из Muxer'а с учетом Backpressure
let current_pending = self.pending_data.len();
let fill_ratio = (current_pending as f32 / MAX_PENDING as f32) * 100.0;
if current_pending >= MAX_PENDING {
netrunner_logger::warn!(
%self.core.handle,
"Backpressure ACTIVE: Buffer is FULL ({} bytes). Stalling RX channel.",
current_pending
);
} else if fill_ratio > 80.0 {
netrunner_logger::info!(
%self.core.handle,
"Bufferbloat Warning: Buffer {:.1}% full ({} bytes). Latency increasing.",
fill_ratio, current_pending
);
while let Ok(data) = self.core.rx.try_recv() {
self.pending_data.extend_from_slice(&data);
if self.pending_data.len() >= MAX_PENDING {
break;
}
}
} else {
while let Ok(data) = self.core.rx.try_recv() {
self.pending_data.extend_from_slice(&data);
if self.pending_data.len() >= MAX_PENDING {
break;
}
}
}
// 3. Отправляем буферизированные данные в smoltcp
if !self.pending_data.is_empty() && socket.can_send() {
match socket.send_slice(&self.pending_data) {
Ok(n) => {
self.pending_data.advance(n);
if n > 0 && self.pending_data.len() < (MAX_PENDING / 2) && fill_ratio > 90.0 {
netrunner_logger::info!(
%self.core.handle,
"Backpressure RELIEVED: Buffer drained to {} bytes",
self.pending_data.len()
);
}
}
Err(e) => {
netrunner_logger::debug!(%self.core.handle, "Smoltcp socket send error: {:?}", e);
}
}
}
}
}
// ============================================================================
// 3. UDP СОЕДИНЕНИЕ (UdpConnection)
// ============================================================================
const UDP_TIMEOUT: Duration = Duration::from_secs(60);
pub struct UdpConnection {
core: ConnectionCore,
client_endpoint: Option<IpEndpoint>,
last_activity: Instant,
}
impl UdpConnection {
pub fn new(handle: SocketHandle) -> (Self, mpsc::Receiver<Bytes>, mpsc::Sender<Bytes>) {
let (core, rx_from_smol, tx_to_smol) = ConnectionCore::new(handle);
let conn = Self {
core,
client_endpoint: None,
last_activity: Instant::now(),
};
(conn, rx_from_smol, tx_to_smol)
}
pub fn tick(&mut self, socket: &mut udp::Socket) -> bool {
// Проверка таймаута
if self.last_activity.elapsed() > UDP_TIMEOUT {
netrunner_logger::debug!(%self.core.handle, "UDP Session closed due to timeout");
socket.close();
return false;
}
// Читаем из smoltcp и шлем в сеть
if socket.can_recv() {
while let Ok((data, metadata)) = socket.recv() {
self.client_endpoint = Some(metadata.endpoint);
if self.core.tx.try_send(Bytes::copy_from_slice(data)).is_ok() {
self.last_activity = Instant::now();
}
}
}
// Читаем из сети и шлем в smoltcp
if socket.can_send() {
if let Some(endpoint) = self.client_endpoint {
while let Ok(data) = self.core.rx.try_recv() {
if data.is_empty() {
socket.close();
return false;
}
if socket.send_slice(&data, endpoint).is_ok() {
self.last_activity = Instant::now();
} else {
break;
}
}
}
}
true
}
}
+519
View File
@@ -0,0 +1,519 @@
use bytes::Bytes;
use netrunner_core::{
protocol::codec::{frame::FrameType, socks::TargetAddress},
proxy::connection::muxer::{MuxMessage, Muxer},
};
use netrunner_logger::{debug, info, warn};
use smoltcp::{
iface::{SocketHandle, SocketSet},
socket::{AnySocket, icmp, tcp, udp},
wire::{IpListenEndpoint, IpProtocol, Ipv4Packet, TcpPacket, UdpPacket},
};
use std::{collections::HashMap, time::Duration, time::Instant as StdInstant};
use tokio::sync::mpsc;
use crate::net::{
CHANNEL_CAPACITY,
connection::{TcpConnection, UdpConnection},
dns::DnsHandler,
ip_store::FakeIpStore,
};
// ============================================================================
// 1. УПРАВЛЕНИЕ СОСТОЯНИЕМ СЕССИЙ (SessionTracker)
// ============================================================================
struct SessionTracker {
last_activity: HashMap<SocketHandle, StdInstant>,
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,
}
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 {
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 _ = 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,
});
}
}
}
}
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 => {
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();
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);
}
}
}
}
}
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) {
for (handle, socket) in socket_set.iter_mut() {
if let Some(tcp) = tcp::Socket::downcast_mut(socket) {
self.handle_tcp(handle, tcp);
} else if let Some(udp) = udp::Socket::downcast_mut(socket) {
self.handle_udp(handle, udp);
} else if let Some(icmp) = icmp::Socket::downcast_mut(socket) {
self.handle_icmp(handle, icmp);
}
}
}
fn handle_tcp(&mut self, handle: SocketHandle, socket: &mut tcp::Socket) {
if socket.state() == tcp::State::Closed {
self.tracker.active_tcp.remove(&handle);
self.tracker.queue_removal(handle);
return;
}
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" {
socket.abort();
return;
}
}
info!(%handle, target = %target, "New TCP session established");
let (conn, mut rx_from_smol, tx_to_smol, handshake_tx) = TcpConnection::new(handle);
self.tracker.active_tcp.insert(handle, conn);
let muxer = self.muxer.clone();
let stream_id = muxer.next_id();
let connect_payload = target.to_string();
tokio::spawn(async move {
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(CHANNEL_CAPACITY);
muxer.register_stream(stream_id, v_tx);
if muxer
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::Connect,
data: Bytes::from(connect_payload),
})
.await
.is_err()
{
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 => {
let _ = handshake_tx.send(());
}
_ => {
muxer.remove_stream(stream_id);
return;
}
}
let to_proxy = async {
while let Some(data) = rx_from_smol.recv().await {
if muxer
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::Data,
data,
})
.await
.is_err()
{
break;
}
}
};
let from_proxy = async {
while let Some(data) = v_rx.recv().await {
if data.is_empty() || tx_to_smol.send(data).await.is_err() {
break;
}
}
};
tokio::select! { _ = to_proxy => {}, _ = from_proxy => {} }
let _ = muxer
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::Close,
data: Bytes::new(),
})
.await;
muxer.remove_stream(stream_id);
});
}
if let Some(conn) = self.tracker.active_tcp.get_mut(&handle) {
if !conn.tick(socket) {
socket.abort();
}
}
if socket.state() == tcp::State::CloseWait {
socket.close();
}
}
fn handle_udp(&mut self, handle: SocketHandle, socket: &mut udp::Socket) {
self.tracker.last_activity.insert(handle, StdInstant::now());
if socket.endpoint().port == 53 {
while socket.can_recv() {
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.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(ip)) => {
TargetAddress::Ipv6(std::net::Ipv6Addr::from(ip), ep.port)
}
None => return,
};
let (conn, mut rx_from_smol, tx_to_smol) = UdpConnection::new(handle);
self.tracker.active_udp.insert(handle, conn);
let muxer = self.muxer.clone();
let stream_id = muxer.next_id();
let connect_payload = target.to_string();
tokio::spawn(async move {
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(CHANNEL_CAPACITY);
muxer.register_stream(stream_id, v_tx);
let _ = muxer
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::UdpConnect,
data: Bytes::from(connect_payload),
})
.await;
let to_proxy = async {
while let Some(data) = rx_from_smol.recv().await {
if muxer
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::UdpData,
data,
})
.await
.is_err()
{
break;
}
}
};
let from_proxy = async {
while let Some(data) = v_rx.recv().await {
if data.is_empty() || tx_to_smol.send(data).await.is_err() {
break;
}
}
};
tokio::select! { _ = to_proxy => {}, _ = from_proxy => {} }
let _ = muxer
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::Close,
data: Bytes::new(),
})
.await;
muxer.remove_stream(stream_id);
});
}
if let Some(conn) = self.tracker.active_udp.get_mut(&handle) {
if !conn.tick(socket) {
self.tracker.queue_removal(handle);
self.tracker.active_udp.remove(&handle);
}
}
}
fn handle_icmp(&mut self, _handle: SocketHandle, socket: &mut icmp::Socket) {
if socket.can_recv() {
let _ = socket.recv();
}
}
pub fn cleanup(&mut self, socket_set: &mut SocketSet) {
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()
);
}
}
+134
View File
@@ -0,0 +1,134 @@
use crate::net::ip_store::FakeIpStore;
use hickory_proto::op::{Message, MessageType, ResponseCode};
use hickory_proto::rr::{RData, Record, RecordType};
use netrunner_logger::{debug, error, info};
use std::collections::HashSet;
use std::time::{Duration, SystemTime};
use tokio::fs::{self, File};
use tokio::io::{AsyncBufReadExt, BufReader};
pub struct DnsHandler {
block_list: HashSet<String>,
forbidden_suffixes: Vec<String>,
cache_path: String,
}
impl DnsHandler {
pub fn new(cache_dir: &str) -> Self {
Self {
block_list: HashSet::new(),
forbidden_suffixes: vec![".lan", ".local", ".home", ".arpa"]
.into_iter()
.map(String::from)
.collect(),
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);
if path.exists() {
let _ = self.load_from_file().await;
}
let needs_update = if let Ok(meta) = fs::metadata(&path).await {
SystemTime::now()
.duration_since(meta.modified()?)
.unwrap_or_default()
> Duration::from_secs(604800)
} 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);
}
});
}
Ok(())
}
async fn download_blocklist_async(cache_path: String) -> anyhow::Result<()> {
info!("DNS: Starting background download to {}", cache_path);
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()?;
let resp = client
.get("https://raw.githubusercontent.com/StevenBlack/hosts/master/hosts")
.send()
.await?;
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);
} else {
error!("DNS: Download failed with status {}", resp.status());
}
Ok(())
}
async fn load_from_file(&mut self) -> anyhow::Result<()> {
let file = File::open(&self.cache_path).await?;
let mut lines = BufReader::new(file).lines();
let mut count = 0;
while let Some(line) = lines.next_line().await? {
let line = line.trim();
if line.is_empty() || line.starts_with('#') {
continue;
}
if let Some(domain) = line.split_whitespace().nth(1) {
self.block_list.insert(domain.to_lowercase());
count += 1;
}
}
info!("DNS: Loaded {} domains from cache.", count);
Ok(())
}
pub fn handle_query(&self, data: &[u8], store: &mut FakeIpStore) -> Option<Vec<u8>> {
let req = Message::from_vec(data).ok()?;
let query = req.queries().first()?;
let name = query
.name()
.to_string()
.trim_end_matches('.')
.to_lowercase();
let mut res = Message::new();
res.set_id(req.id())
.set_message_type(MessageType::Response)
.set_recursion_available(true)
.add_query(query.clone());
if self.forbidden_suffixes.iter().any(|s| name.ends_with(s))
|| self.block_list.contains(&name)
{
debug!(domain = %name, "DNS: Blocked");
res.set_response_code(ResponseCode::NXDomain);
return res.to_vec().ok();
}
if query.query_type() == RecordType::A {
let fake_ip = store.get_or_assign(&name);
res.add_answer(Record::from_rdata(
query.name().clone(),
60,
RData::A(fake_ip.into()),
));
res.set_response_code(ResponseCode::NoError);
}
res.to_vec().ok()
}
}
+340
View File
@@ -0,0 +1,340 @@
use netrunner_core::proxy::connection::connection::ClientHandler;
use netrunner_core::proxy::connection::muxer::Muxer;
use smoltcp::iface::PollResult;
use smoltcp::time::Instant;
use smoltcp::wire::{IpAddress, IpCidr};
use smoltcp::{
iface::{Config, Interface, SocketSet},
phy::DeviceCapabilities,
};
use std::net::Ipv4Addr;
use std::sync::atomic::Ordering;
use std::{
sync::{Arc, LazyLock, atomic::AtomicBool},
time::Instant as StdInstant,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
use tokio::time::{Duration, sleep};
use tun::{DeviceReader, DeviceWriter};
use netrunner_logger::{debug, error, info, warn};
use crate::net::connection_manager::ConnectionManager;
use crate::net::dns::DnsHandler;
use crate::tun::device::{TokenBuffer, VirtTunDevice};
use crate::tun::routing::setup_platform_routing;
use crate::tun::tun::Tun;
pub static START_TIME: LazyLock<StdInstant> = LazyLock::new(StdInstant::now);
pub struct Engine {
interface: Interface,
socket_set: SocketSet<'static>,
manager: ConnectionManager,
device: VirtTunDevice,
to_smoltcp_tx: UnboundedSender<TokenBuffer>,
from_smoltcp_rx: Option<UnboundedReceiver<TokenBuffer>>,
avail: Arc<AtomicBool>,
}
impl Engine {
pub fn new(
config: Config,
caps: DeviceCapabilities,
dns_handler: DnsHandler,
muxer: Muxer,
) -> 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 manager = ConnectionManager::new(dns_handler, muxer);
Self {
interface,
socket_set,
device,
to_smoltcp_tx,
from_smoltcp_rx: Some(from_smoltcp_rx),
avail,
manager,
}
}
pub async fn run(&mut self, tun: Tun) {
info!("Current routes: {:?}", self.interface.routes());
let (writer, reader) = tun.split().expect("Failed to split TUN");
let (tun_to_engine_tx, mut tun_to_engine_rx) = mpsc::unbounded_channel();
Self::spawn_tun_reader(reader, tun_to_engine_tx, self.avail.clone());
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();
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();
}
}
let result = self.poll();
self.manager.process_sockets(&mut self.socket_set);
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.swap(false, Ordering::Acquire) {
tokio::task::yield_now().await;
continue;
}
self.manager.cleanup(&mut self.socket_set);
self.poll_delay().await;
}
}
fn poll(&mut self) -> PollResult {
let now = Self::current_time();
self.interface
.poll(now, &mut self.device, &mut self.socket_set)
}
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 {
Some(d) => Duration::from_micros(d.micros()),
None => Duration::from_millis(10),
};
sleep(sleep_duration).await;
}
fn spawn_tun_reader(
mut reader: DeviceReader,
to_engine: UnboundedSender<TokenBuffer>,
is_avail: Arc<AtomicBool>,
) {
tokio::spawn(async move {
debug!("TUN Reader task started");
let mut buf = [0u8; 65536];
while let Ok(n) = reader.read(&mut buf).await {
if n == 0 {
break;
}
let first_byte = buf[0];
let version = first_byte >> 4;
if version == 4 {
if buf[12..16] == [0, 0, 0, 0] {
continue;
}
if buf[16] >= 224 && buf[16] <= 239 {
continue;
}
} else if version == 6 {
if first_byte == 0xff {
continue;
}
}
let mut token = TokenBuffer::with_capacity(n);
token.extend_from_slice(&buf[..n]);
if to_engine.send(token).is_ok() {
is_avail.store(true, Ordering::Release);
} else {
break;
}
}
warn!("TUN Reader task stopped");
});
}
fn spawn_tun_writer(
mut writer: DeviceWriter,
mut from_smoltcp: UnboundedReceiver<TokenBuffer>,
) {
tokio::spawn(async move {
debug!("TUN Writer task started");
while let Some(token) = from_smoltcp.recv().await {
if writer.write_all(&token).await.is_err() {
break;
}
}
warn!("TUN Writer task stopped");
});
}
fn current_time() -> Instant {
let duration = StdInstant::now().duration_since(*START_TIME);
Instant::from_micros(duration.as_micros() as i64)
}
pub fn set_any_ip(&mut self, state: bool) {
self.interface.set_any_ip(state)
}
pub fn set_transparent_mode(&mut self) {
self.interface.update_ip_addrs(|addrs| {
addrs.clear();
addrs
.push(IpCidr::new(IpAddress::v4(10, 0, 0, 2), 24))
.unwrap();
});
self.interface.routes_mut().remove_default_ipv4_route();
}
pub fn set_default_gateway(&mut self, gateway: smoltcp::wire::Ipv4Address) {
info!("Setting default IPv4 gateway to: {}", gateway);
self.interface.routes_mut().remove_default_ipv4_route();
self.interface
.routes_mut()
.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);
}
}
// ============================================================================
// КОНФИГУРАЦИЯ ДВИЖКА
// ============================================================================
#[derive(Clone, Debug)]
pub struct EngineConfig {
pub remote_address: String,
pub cache_path: String,
pub mtu: usize,
pub setup_routing: bool,
pub any_ip: bool,
pub transparent_mode: bool,
pub default_gateway: Ipv4Addr,
}
impl EngineConfig {
pub fn new(remote_address: impl Into<String>) -> Self {
Self {
remote_address: remote_address.into(),
cache_path: ".".to_string(),
mtu: 1350,
setup_routing: true,
any_ip: true,
transparent_mode: true,
default_gateway: Ipv4Addr::new(10, 0, 0, 2),
}
}
pub fn with_cache_path(mut self, path: impl Into<String>) -> Self {
self.cache_path = path.into();
self
}
pub fn with_mtu(mut self, mtu: usize) -> Self {
self.mtu = mtu;
self
}
pub fn _disable_routing(mut self) -> Self {
self.setup_routing = false;
self
}
}
// ============================================================================
// БИЛДЕР ДВИЖКА
// ============================================================================
pub struct EngineBuilder {
config: EngineConfig,
tun_device: Option<Tun>,
}
impl EngineBuilder {
/// Инициализируем билдер на основе готового конфига
pub fn new(config: EngineConfig) -> Self {
Self {
config,
tun_device: None,
}
}
/// Передаем TUN интерфейс (зависит от платформы, поэтому не в конфиге)
pub fn with_tun(mut self, tun: Tun) -> Self {
self.tun_device = Some(tun);
self
}
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
);
// 1. Инициализация DNS
let mut dns_handler = DnsHandler::new(&self.config.cache_path);
if let Err(e) = dns_handler.init().await {
error!("Failed to initialize DNS blocklist: {}", e);
}
// 2. Настройка системного роутинга (Опционально)
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.");
}
// 3. Конфигурация интерфейса smoltcp
let smol_config = Config::new(smoltcp::wire::HardwareAddress::Ip);
let mut caps = DeviceCapabilities::default();
caps.max_transmission_unit = self.config.mtu; // Берем из конфига
caps.medium = smoltcp::phy::Medium::Ip;
// 4. Подключение к серверу
info!("Establishing secure tunnel to proxy server...");
let muxer = ClientHandler::connect(&self.config.remote_address)
.await
.map_err(|e| format!("Failed to establish secure tunnel: {}", e))?;
info!("Secure tunnel established, Muxer is ready.");
// 5. Инициализация и настройка Engine
let mut engine = Engine::new(smol_config, caps, dns_handler, muxer);
engine.set_any_ip(self.config.any_ip);
if self.config.transparent_mode {
engine.set_transparent_mode();
}
engine.set_default_gateway(self.config.default_gateway);
engine.activate();
info!("Stack IP initialized: {}", self.config.default_gateway);
Ok((engine, tun))
}
}
+46
View File
@@ -0,0 +1,46 @@
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
}
}
}
+7
View File
@@ -0,0 +1,7 @@
mod connection;
pub mod connection_manager;
pub mod dns;
pub mod engine;
pub mod ip_store;
pub const CHANNEL_CAPACITY: usize = 16;