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

core engine rewrite and client side changes

Browse Source
This commit is contained in:
Kirill
2026-03-21 22:24:21 +07:00
parent af5ab22712
commit 46e8b4ae73
13 changed files with 573 additions and 376 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/src/connections/dns.rs
+2 -2
View File
@@ -1,7 +1,7 @@
use crate::connections::ip_store::FakeIpStore;
use hickory_proto::op::{Message, MessageType, ResponseCode};
use hickory_proto::rr::{RData, Record, RecordType};
use netrunner_logger::{error, info};
use netrunner_logger::{debug, error, info};
use std::collections::HashSet;
use std::time::{Duration, SystemTime};
use tokio::fs::{self, File};
@@ -114,7 +114,7 @@ impl DnsHandler {
if self.forbidden_suffixes.iter().any(|s| name.ends_with(s))
|| self.block_list.contains(&name)
{
info!(domain = %name, "DNS: Blocked");
debug!(domain = %name, "DNS: Blocked");
res.set_response_code(ResponseCode::NXDomain);
return res.to_vec().ok();
}
client/src/connections/ip_store.rs
+1 -1
View File
@@ -30,7 +30,7 @@ impl FakeIpStore {
self.next_ip += 1;
self.cache.put(host.to_string(), ip);
self.rev_cache.put(ip, host.to_string());
info!(host = %host, ip = %ip, "Assigned new fake IP");
debug!(host = %host, ip = %ip, "Assigned new fake IP");
ip
}
client/src/connections/mod.rs
+1 -1
View File
@@ -3,4 +3,4 @@ pub mod ip_store;
pub mod tcp_connection;
pub mod udp_connection;
pub const CHANNEL_CAPACITY: usize = 2048;
pub const CHANNEL_CAPACITY: usize = 16;
client/src/connections/tcp_connection.rs
+33 -115
View File
@@ -1,16 +1,9 @@
use std::time::Duration;
use bytes::Bytes;
use netrunner_core::protocol::codec::frame::FrameType;
use netrunner_core::protocol::codec::socks::TargetAddress;
use netrunner_core::proxy::connection::BUF_SIZE;
use netrunner_core::proxy::connection::muxer::{MuxMessage, Muxer};
use bytes::{Buf, Bytes, BytesMut};
use netrunner_core::proxy::connection::TCP_BUF_SIZE;
use smoltcp::iface::SocketHandle;
use smoltcp::socket::tcp;
use tokio::sync::{mpsc, oneshot};
use crate::connections::CHANNEL_CAPACITY;
pub enum ConnectionState {
Established,
Handshaking,
@@ -19,118 +12,44 @@ pub enum ConnectionState {
}
pub struct TcpConnection {
handle: SocketHandle,
pub handle: SocketHandle,
state: ConnectionState,
// UPLOAD: Ограниченный канал (Bounded) для Backpressure (128)
tx: mpsc::Sender<Vec<u8>>,
// DOWNLOAD: Безлимитный канал (Unbounded), чтобы Муксер НИКОГДА не зависал
rx: mpsc::UnboundedReceiver<Vec<u8>>,
pending_data: Vec<u8>,
// UPLOAD: Ограниченный канал для передачи данных наружу
tx: mpsc::Sender<Bytes>,
// DOWNLOAD: Канал для приема данных из сети
rx: mpsc::Receiver<Bytes>,
pending_data: BytesMut,
handshake_rx: Option<oneshot::Receiver<()>>,
}
const MAX_PENDING: usize = 4 * 1024 * 1024;
const MAX_PENDING: usize = 32 * 512 * 1024;
const TCP_CHUNK_SIZE: usize = 65536;
impl TcpConnection {
pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self {
// UPLOAD: Ограничиваем очередь
let (tx_to_mux, mut rx_from_smol) = mpsc::channel::<Vec<u8>>(CHANNEL_CAPACITY);
// DOWNLOAD: Безлимитный канал до синхронного tick()
let (tx_to_smol, rx_from_proxy) = mpsc::unbounded_channel::<Vec<u8>>();
pub fn new(
handle: SocketHandle,
) -> (
Self,
mpsc::Receiver<Bytes>,
mpsc::Sender<Bytes>,
oneshot::Sender<()>,
) {
let (tx_to_net, rx_from_smol) = mpsc::channel::<Bytes>(TCP_BUF_SIZE);
let (tx_to_smol, rx_from_net) = mpsc::channel::<Bytes>(TCP_BUF_SIZE);
let (handshake_tx, handshake_rx) = oneshot::channel();
let stream_id = muxer.next_id();
tokio::spawn(async move {
// ИСПРАВЛЕНИЕ: Даем Муксеру ограниченный канал, как он и просит (тип mpsc::Sender)
// Делаем его достаточно вместительным (1024)
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(BUF_SIZE);
muxer.register_stream(stream_id, v_tx);
let connect_payload = target_addr.to_string();
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;
}
let first_payload = tokio::time::timeout(Duration::from_secs(10), v_rx.recv()).await;
match first_payload {
Ok(Some(data)) => {
if data.len() >= 2 && data[1] == 0x00 {
let _ = handshake_tx.send(());
} else {
netrunner_logger::warn!(stream_id, "Server rejected TCP connection");
muxer.remove_stream(stream_id);
return;
}
}
_ => {
netrunner_logger::error!(stream_id, "Timeout waiting for proxy response");
muxer.remove_stream(stream_id);
return;
}
}
let to_proxy = async {
while let Some(data) = rx_from_smol.recv().await {
let msg = MuxMessage {
stream_id,
frame_type: FrameType::Data,
data: Bytes::from(data),
};
if muxer.send_to_netwrok(msg).await.is_err() {
break;
}
}
};
// МАГИЯ ЗДЕСЬ: Мы мгновенно читаем из ограниченного v_rx и переливаем
// в безлимитный tx_to_smol. send() у безлимитного канала никогда не блокируется.
// Поэтому v_rx всегда пустой, и Муксер никогда не зависнет!
let from_proxy = async {
while let Some(data) = v_rx.recv().await {
if data.is_empty() {
break;
}
if tx_to_smol.send(data.to_vec()).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);
});
Self {
let conn = Self {
handle,
state: ConnectionState::Handshaking,
tx: tx_to_mux,
rx: rx_from_proxy, // Это UnboundedReceiver, здесь ничего менять не надо
pending_data: vec![],
tx: tx_to_net,
rx: rx_from_net,
pending_data: BytesMut::new(),
handshake_rx: Some(handshake_rx),
}
};
(conn, rx_from_smol, tx_to_smol, handshake_tx)
}
pub fn tick(&mut self, socket: &mut tcp::Socket) -> bool {
let state = socket.state();
@@ -188,7 +107,7 @@ impl TcpConnection {
}
fn poll_and_process(&mut self, socket: &mut tcp::Socket) {
// 1. UPLOAD: Чанкинг и Backpressure (Защита от краша сервера)
// 1. UPLOAD: Чанкинг и Backpressure
if socket.can_recv() {
while socket.can_recv() {
let mut channel_full = false;
@@ -199,10 +118,10 @@ impl TcpConnection {
return (0, ());
}
let chunk_size = std::cmp::min(data.len(), 16000);
let chunk_size = std::cmp::min(data.len(), TCP_CHUNK_SIZE);
let chunk = &data[..chunk_size];
match self.tx.try_send(chunk.to_vec()) {
match self.tx.try_send(Bytes::copy_from_slice(chunk)) {
Ok(_) => (chunk_size, ()),
Err(mpsc::error::TrySendError::Full(_)) => {
channel_full = true;
@@ -221,7 +140,7 @@ impl TcpConnection {
}
}
// 2. DOWNLOAD: Сброс безлимитного канала в буфер
// 2. DOWNLOAD: Сброс ограниченного канала в буфер
while let Ok(data) = self.rx.try_recv() {
if self.pending_data.is_empty() && socket.can_send() {
match socket.send_slice(&data) {
@@ -237,7 +156,6 @@ impl TcpConnection {
self.pending_data.extend_from_slice(&data);
}
// Защита от OOM. 32MB более чем достаточно.
if self.pending_data.len() > MAX_PENDING {
netrunner_logger::error!(
%self.handle,
@@ -254,7 +172,7 @@ impl TcpConnection {
if !self.pending_data.is_empty() && socket.can_send() {
match socket.send_slice(&self.pending_data) {
Ok(n) => {
self.pending_data.drain(..n);
self.pending_data.advance(n);
}
Err(_) => {}
}
client/src/connections/udp_connection.rs
+30 -85
View File
@@ -3,128 +3,71 @@ use smoltcp::socket::udp;
use smoltcp::wire::IpEndpoint;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use bytes::Bytes;
use netrunner_core::{
protocol::codec::{frame::FrameType, socks::TargetAddress},
proxy::connection::muxer::{MuxMessage, Muxer},
};
use crate::connections::CHANNEL_CAPACITY;
use bytes::Bytes;
pub struct UdpConnection {
pub handle: SocketHandle,
stream_id: u32,
tx_to_net: mpsc::Sender<MuxMessage>,
rx_from_net: mpsc::UnboundedReceiver<Bytes>, // <-- ТУТ
// UPLOAD: Ограниченный канал для передачи данных наружу
tx: mpsc::Sender<Bytes>,
// DOWNLOAD: Канал для приема данных из сети
rx: mpsc::Receiver<Bytes>,
client_endpoint: Option<IpEndpoint>,
last_activity: Instant,
token: CancellationToken,
}
const UDP_TIMEOUT: Duration = Duration::from_secs(60);
impl UdpConnection {
pub fn new(handle: SocketHandle, target_addr: TargetAddress, muxer: Muxer) -> Self {
let stream_id = muxer.next_id();
let token = CancellationToken::new();
let task_token = token.clone();
/// Возвращает:
/// 1. Экземпляр UdpConnection
/// 2. Receiver (для чтения данных, отправляемых ИЗ smoltcp В сеть)
/// 3. Sender (для записи данных ИЗ сети В smoltcp)
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 (tx_to_net, mut rx_from_smol) = mpsc::channel::<MuxMessage>(CHANNEL_CAPACITY);
// ИСПРАВЛЕНИЕ: Внутренний безлимитный канал
let (internal_tx, internal_rx) = mpsc::unbounded_channel::<Bytes>();
let m_clone = muxer.clone();
tokio::spawn(async move {
// Даем Муксеру ограниченный канал (тип Sender), чтобы компилятор пропустил
let (v_tx, mut v_rx) = mpsc::channel::<Bytes>(1024);
m_clone.register_stream(stream_id, v_tx);
let _ = m_clone
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::UdpConnect,
data: Bytes::from(target_addr.to_string()),
})
.await;
let to_proxy = async {
while let Some(msg) = rx_from_smol.recv().await {
if m_clone.send_to_netwrok(msg).await.is_err() {
break;
}
}
};
// Мгновенно выкачиваем данные из Муксера в безлимитный внутренний канал
let from_proxy = async {
while let Some(data) = v_rx.recv().await {
if internal_tx.send(data).is_err() {
break;
}
}
};
tokio::select! {
_ = to_proxy => {}
_ = from_proxy => {} // Не забудь добавить этот блок в select!
_ = task_token.cancelled() => {}
}
let _ = m_clone
.send_to_netwrok(MuxMessage {
stream_id,
frame_type: FrameType::Close,
data: Bytes::new(),
})
.await;
m_clone.remove_stream(stream_id);
});
Self {
let conn = Self {
handle,
stream_id,
tx_to_net,
rx_from_net: internal_rx, // Сохраняем безлимитный ресивер
tx: tx_to_net,
rx: rx_from_net,
client_endpoint: None,
last_activity: Instant::now(),
token,
}
};
(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.handle, "UDP Session closed due to timeout");
self.token.cancel();
socket.close();
return false;
}
// 1. UPLOAD: Читаем из smoltcp и отправляем в виртуальный канал
if socket.can_recv() {
while let Ok((data, metadata)) = socket.recv() {
self.client_endpoint = Some(metadata.endpoint);
let msg = MuxMessage {
stream_id: self.stream_id,
frame_type: FrameType::UdpData,
data: Bytes::copy_from_slice(data),
};
if self.tx_to_net.try_send(msg).is_ok() {
// Копируем данные в Bytes и пытаемся протолкнуть.
// Если канал забит (Full), пакет дропается — это штатное поведение UDP.
if self.tx.try_send(Bytes::copy_from_slice(data)).is_ok() {
self.last_activity = Instant::now();
}
}
}
// 2. DOWNLOAD: Читаем из виртуального канала и пишем в smoltcp
if socket.can_send() {
if let Some(endpoint) = self.client_endpoint {
while let Ok(data) = self.rx_from_net.try_recv() {
while let Ok(data) = self.rx.try_recv() {
// Сигнал закрытия стрима
if data.is_empty() {
self.token.cancel();
break;
socket.close();
return false;
}
match socket.send_slice(&data, endpoint) {
@@ -132,6 +75,8 @@ impl UdpConnection {
self.last_activity = Instant::now();
}
Err(_) => {
// Если сокет smoltcp переполнен, прерываем цикл.
// Пакет теряется, что опять же является нормой для UDP.
break;
}
}
client/src/tun/connection_manager.rs
+159 -7
View File
@@ -1,16 +1,25 @@
use netrunner_core::{protocol::codec::socks::TargetAddress, proxy::connection::muxer::Muxer};
use bytes::Bytes;
use netrunner_core::{
protocol::codec::{frame::FrameType, socks::TargetAddress},
proxy::connection::{
TCP_BUF_SIZE,
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::Instant as StdInstant};
use std::{collections::HashMap, time::Duration, time::Instant as StdInstant};
use tokio::sync::mpsc;
use crate::connections::{
dns::DnsHandler, ip_store::FakeIpStore, tcp_connection::TcpConnection,
CHANNEL_CAPACITY, dns::DnsHandler, ip_store::FakeIpStore, tcp_connection::TcpConnection,
udp_connection::UdpConnection,
};
pub struct ConnectionManager {
last_activity: HashMap<SocketHandle, StdInstant>,
active_tcp_sessions: HashMap<SocketHandle, TcpConnection>,
@@ -35,6 +44,7 @@ impl ConnectionManager {
muxer,
}
}
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) {
@@ -65,7 +75,6 @@ impl ConnectionManager {
Some(ep) => ep,
None => {
warn!(handle=?socket, "Attempted to resolve target for an unconnected socket");
return TargetAddress::Domain("disconnected".to_string(), 0);
}
};
@@ -96,6 +105,7 @@ impl ConnectionManager {
}
}
}
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) {
@@ -134,8 +144,90 @@ impl ConnectionManager {
}
info!(%handle, "New TCP session established for target: {:?}", target);
let conn = TcpConnection::new(handle, target, self.muxer.clone());
// 1. Создаем соединение и получаем каналы
let (conn, mut rx_from_smol, tx_to_smol, handshake_tx) = TcpConnection::new(handle);
self.active_tcp_sessions.insert(handle, conn);
// 2. Выносим логику мультиплексирования в фоновую задачу
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>(TCP_BUF_SIZE);
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;
}
let first_payload =
tokio::time::timeout(Duration::from_secs(10), v_rx.recv()).await;
match first_payload {
Ok(Some(data)) => {
if data.len() >= 2 && data[1] == 0x00 {
let _ = handshake_tx.send(());
} else {
netrunner_logger::warn!(stream_id, "Server rejected TCP connection");
muxer.remove_stream(stream_id);
return;
}
}
_ => {
netrunner_logger::error!(stream_id, "Timeout waiting for proxy response");
muxer.remove_stream(stream_id);
return;
}
}
let to_proxy = async {
while let Some(data) = rx_from_smol.recv().await {
let msg = MuxMessage {
stream_id,
frame_type: FrameType::Data,
data,
};
if muxer.send_to_netwrok(msg).await.is_err() {
break;
}
}
};
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() {
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.active_tcp_sessions.get_mut(&handle) {
@@ -149,6 +241,7 @@ impl ConnectionManager {
socket.close();
}
}
fn handle_udp(&mut self, handle: SocketHandle, socket: &mut udp::Socket) {
self.last_activity.insert(handle, StdInstant::now());
@@ -187,8 +280,65 @@ impl ConnectionManager {
netrunner_logger::info!(%handle, target = %target, "New UDP proxied session established");
let conn = UdpConnection::new(handle, target, self.muxer.clone());
// 1. Создаем UDP соединение и получаем каналы
let (conn, mut rx_from_smol, tx_to_smol) = UdpConnection::new(handle);
self.active_udp_sessions.insert(handle, conn);
// 2. Фоновая задача для Muxer'а
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 {
let msg = MuxMessage {
stream_id,
frame_type: FrameType::UdpData,
data,
};
if muxer.send_to_netwrok(msg).await.is_err() {
break;
}
}
};
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() {
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.active_udp_sessions.get_mut(&handle) {
@@ -223,11 +373,11 @@ impl ConnectionManager {
tcp::SocketBuffer::new(vec![0; buf_size]),
);
// Nagle отключен - это правильно для уменьшения задержек
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;
@@ -357,6 +507,7 @@ impl ConnectionManager {
self.failed_until.remove(&handle);
}
}
pub fn setup_sockets(n_icmp: usize) -> SocketSet<'static> {
let mut sockets = SocketSet::new(Vec::with_capacity(48));
@@ -366,6 +517,7 @@ impl ConnectionManager {
sockets
}
pub fn log_status(&self, socket_set: &SocketSet) {
let mut established = 0;
let mut total_tcp = 0;
client/src/tun/device.rs
+107
View File
@@ -6,9 +6,11 @@ use std::{
Arc, LazyLock, Mutex,
atomic::{AtomicBool, Ordering},
},
time::Instant as StdInstant, // Добавлено для расчёта скорости
};
use bytes::BytesMut;
use netrunner_logger::info; // Предполагается, что у тебя есть этот логгер
use smoltcp::{
phy::{self, Device, DeviceCapabilities},
time::Instant,
@@ -18,6 +20,17 @@ use tokio::sync::mpsc;
const TOKEN_BUFFER_LIST_MAX_SIZE: usize = 64;
static TOKEN_BUFFER_LIST: LazyLock<Mutex<Vec<BytesMut>>> = LazyLock::new(|| Mutex::new(Vec::new()));
// Структура, которую мы будем возвращать пользователю
#[derive(Debug, Clone, Copy)]
pub struct TrafficStats {
pub rx_bytes: u64,
pub tx_bytes: u64,
pub rx_packets: u64,
pub tx_packets: u64,
pub rx_speed_mb_s: f64,
pub tx_speed_mb_s: f64,
}
pub struct TokenBuffer {
buffer: BytesMut,
}
@@ -65,6 +78,22 @@ pub struct VirtTunDevice {
rx_queue: mpsc::UnboundedReceiver<TokenBuffer>,
tx_queue: mpsc::UnboundedSender<TokenBuffer>,
rx_avail: Arc<AtomicBool>,
// === Статистика трафика ===
rx_bytes: u64,
tx_bytes: u64,
rx_packets: u64,
tx_packets: u64,
// Внутренние переменные для расчёта скорости
last_speed_calc: StdInstant,
last_rx_bytes: u64,
last_tx_bytes: u64,
cached_rx_speed: f64,
cached_tx_speed: f64,
// Для периодического логирования
last_log_time: StdInstant,
}
impl VirtTunDevice {
@@ -80,11 +109,26 @@ impl VirtTunDevice {
let (from_smoltcp_tx, from_smoltcp_rx) = mpsc::unbounded_channel();
let rx_avail = Arc::new(AtomicBool::new(false));
let now = StdInstant::now();
let device = Self {
capabilities,
rx_queue: to_smoltcp_rx,
tx_queue: from_smoltcp_tx,
rx_avail: rx_avail.clone(),
rx_bytes: 0,
tx_bytes: 0,
rx_packets: 0,
tx_packets: 0,
last_speed_calc: now,
last_rx_bytes: 0,
last_tx_bytes: 0,
cached_rx_speed: 0.0,
cached_tx_speed: 0.0,
last_log_time: now,
};
(device, to_smoltcp_tx, from_smoltcp_rx, rx_avail)
@@ -94,6 +138,54 @@ impl VirtTunDevice {
pub fn mark_rx_available(&self) {
self.rx_avail.store(true, Ordering::Release);
}
pub fn get_stats(&mut self) -> TrafficStats {
let now = StdInstant::now();
let elapsed_speed = now.duration_since(self.last_speed_calc).as_secs_f64();
if elapsed_speed >= 1.0 {
let rx_diff = self.rx_bytes.saturating_sub(self.last_rx_bytes);
let tx_diff = self.tx_bytes.saturating_sub(self.last_tx_bytes);
// Переводим в МегаБайты в секунду (MB/s)
self.cached_rx_speed = (rx_diff as f64 / 1_048_576.0) / elapsed_speed;
self.cached_tx_speed = (tx_diff as f64 / 1_048_576.0) / elapsed_speed;
self.last_rx_bytes = self.rx_bytes;
self.last_tx_bytes = self.tx_bytes;
self.last_speed_calc = now;
}
TrafficStats {
rx_bytes: self.rx_bytes,
tx_bytes: self.tx_bytes,
rx_packets: self.rx_packets,
tx_packets: self.tx_packets,
rx_speed_mb_s: self.cached_rx_speed,
tx_speed_mb_s: self.cached_tx_speed,
}
}
/// Внутренний метод проверки таймера для вывода логов
fn check_and_log_stats(&mut self) {
let now = StdInstant::now();
// Логируем каждые 5 секунд
if now.duration_since(self.last_log_time).as_secs() >= 5 {
let stats = self.get_stats(); // Заодно обновляем скорости
info!(
"TunDevice Traffic: RX: {:.2} MB ({} pkts) | TX: {:.2} MB ({} pkts) | Speed: ↓{:.2} MB/s, ↑{:.2} MB/s",
stats.rx_bytes as f64 / 1_048_576.0,
stats.rx_packets,
stats.tx_bytes as f64 / 1_048_576.0,
stats.tx_packets,
stats.rx_speed_mb_s,
stats.tx_speed_mb_s
);
self.last_log_time = now;
}
}
}
impl Device for VirtTunDevice {
@@ -101,7 +193,15 @@ impl Device for VirtTunDevice {
type TxToken<'a> = VirtTxToken<'a>;
fn receive(&mut self, _timestamp: Instant) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
// Проверяем, не пора ли записать стату в лог. receive() вызывается очень часто,
// поэтому это отличное место для heartbeat таймера.
self.check_and_log_stats();
if let Ok(buffer) = self.rx_queue.try_recv() {
// Учет входящего трафика (DOWNLOAD)
self.rx_bytes += buffer.len() as u64;
self.rx_packets += 1;
let rx = Self::RxToken {
buffer,
phantom_device: PhantomData,
@@ -147,8 +247,15 @@ impl phy::TxToken for VirtTxToken<'_> {
unsafe {
buffer.set_len(len);
}
let result = f(&mut buffer);
// Учет исходящего трафика (UPLOAD)
self.0.tx_bytes += len as u64;
self.0.tx_packets += 1;
let _ = self.0.tx_queue.send(buffer);
result
}
}