buf size changes, update leg select, update client socket, set lowat

This commit is contained in:
Трапезников Кирилл Иванович
2026-04-10 20:52:41 +10:00
parent eb23825bf5
commit c0c5943da7
10 changed files with 133 additions and 136 deletions
+2 -2
View File
@@ -49,13 +49,13 @@ impl NetworkConfig {
client_muxer_capacity: 8,
client_tun_capacity: 16,
client_stream_capacity: 16,
client_virtual_stream_capacity: 8,
client_virtual_stream_capacity: 128,
server_muxer_capacity: 64, // Кардинально режем
server_stream_capacity: 32,
// 🔥 Расширяем TCP окна под BBR (Разблокируем Gigabit на дальние дистанции)
tcp_rx_heavy: 64 * 1024, //64KB
tcp_rx_heavy: 512 * 1024, //512
tcp_tx_heavy: 1 * 1024 * 1024, // 1 MB
tcp_rx_light: 16 * 1024,
+46 -41
View File
@@ -13,6 +13,7 @@ use bytes::{Bytes, BytesMut};
use dashmap::DashMap;
use netrunner_logger::{debug, error, info, warn};
use rand::Rng;
use socket2::{Domain, Protocol, Socket, Type};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{
@@ -84,33 +85,51 @@ impl ClientHandler {
session_id: &str,
) -> Result<(), String> {
let leg_name = if leg_id % 2 == 0 { "TCP-Leg" } else { "UDP-Leg" };
info!("Establishing dedicated {} (ID: {}) to {}...", leg_name, leg_id, remote_proxy_addr);
let mut addrs = tokio::net::lookup_host(remote_proxy_addr)
.await
.map_err(|e| format!("DNS resolution failed: {}", e))?;
let addr = addrs.next().ok_or_else(|| format!("No IPs found for {}", remote_proxy_addr))?;
let socket = if addr.is_ipv4() {
tokio::net::TcpSocket::new_v4().map_err(|e| e.to_string())?
} else {
tokio::net::TcpSocket::new_v6().map_err(|e| e.to_string())?
};
// --- Шаг 1: Создание сырого сокета через socket2 ---
let domain = if addr.is_ipv4() { Domain::IPV4 } else { Domain::IPV6 };
let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))
.map_err(|e| e.to_string())?;
let stream = socket.connect(addr).await.map_err(|e| format!("Connect failed: {}", e))?;
let _ = stream.set_nodelay(true);
// --- Шаг 2: Настройка кроссплатформенных опций ---
socket.set_nonblocking(true).map_err(|e| e.to_string())?;
socket.set_nodelay(true).map_err(|e| e.to_string())?;
// 🔥 Шаг 3: Установка TCP_NOTSENT_LOWAT через libc (Linux, Android, iOS, macOS)
#[cfg(any(target_os = "linux", target_os = "android", target_os = "ios", target_os = "macos"))]
unsafe {
use std::os::fd::AsRawFd;
let lowat: libc::c_int = 16384; // 16 КБ
let ret = libc::setsockopt(
socket.as_raw_fd(),
libc::IPPROTO_TCP,
libc::TCP_NOTSENT_LOWAT,
&lowat as *const _ as *const libc::c_void,
std::mem::size_of::<libc::c_int>() as libc::socklen_t,
);
if ret != 0 {
warn!("Could not set TCP_NOTSENT_LOWAT: {}", std::io::Error::last_os_error());
}
}
// --- Шаг 4: Подключение и конвертация в TcpStream для Tokio ---
let _ = socket.connect(&addr.into());
let std_stream: std::net::TcpStream = socket.into();
let stream = TcpStream::from_std(std_stream).map_err(|e| e.to_string())?;
let mut conn = Connection::new(stream);
// --- 1. TLS Handshake Phase ---
let mut session_keys = SessionKeys::new(true);
let ch = TlsBridge::wrap_client_hello(&BrowserProfile::CHROME_131, "ubuntu.com", &session_keys);
conn.outbound.write_all(&ch).await.map_err(|e| e.to_string())?;
loop {
// Проверяем, есть ли уже ServerHello в буфере
match TlsBridge::unpack_handshake(&mut conn.read_buf) {
Ok(Some(msg)) => {
session_keys.update_keys(msg.random(), msg.extensions(), false)
@@ -125,29 +144,22 @@ impl ClientHandler {
}
}
info!("{} TLS Handshake complete.", leg_name);
// --- 2. Data Phase Initialization ---
let (tx_key, tx_iv, rx_key, rx_iv) = session_keys.get_aead_parameters();
let mut cipher = ChaChaCipher::new();
cipher.set_keys(tx_key, tx_iv, rx_key, rx_iv);
// Забираем остатки из хендшейка в новый кодек
let codec = Codec::new(cipher, session_keys.get_auth_key());
let (rx_codec, mut tx_codec) = codec.split();
// --- 3. Encrypted Handshake ---
let handshake_payload = Bytes::from(format!("{}:{}", session_id, leg_id));
let encrypted_handshake = tx_codec.encode_frame(0, FrameType::Handshake, handshake_payload)
.map_err(|e| format!("Failed to encrypt Handshake: {:?}", e))?;
conn.outbound.write_all(&encrypted_handshake).await.map_err(|e| e.to_string())?;
let (control_tx, control_rx) = mpsc::channel(NetworkConfig::global().client_muxer_capacity);
let (data_tx, data_rx) = mpsc::channel(NetworkConfig::global().client_muxer_capacity);
muxer.add_leg(leg_id, control_tx, data_tx);
let handler = Arc::new(StreamHandler::new(muxer.clone(), None));
let engine = TunnelEngine {
leg_id,
inbound: conn.inbound,
@@ -160,7 +172,6 @@ impl ClientHandler {
handler,
muxer: muxer.clone(),
};
engine.run().await.map_err(|e| e.to_string())?;
Err(format!("{} Engine stopped", leg_name))
}
@@ -171,10 +182,9 @@ impl ClientHandler {
tx_to_engine: mpsc::Sender<RawCastFrame>,
) -> Result<(), String> {
let session_id = SessionManager::generate_id();
info!("🔑 Generated Master Session ID: {}", session_id);
let muxer = Arc::new(Muxer::new(true, session_id.clone()));
let registry: Arc<DashMap<u32, StreamContext>> = Arc::new(DashMap::new());
let local_to_global: Arc<DashMap<u32, u32>> = Arc::new(DashMap::new());
let registry: Arc<DashMap<u32, (u64, Ipv4Addr, u16, LocalProtocol)>> = Arc::new(DashMap::new());
let local_to_global: Arc<DashMap<u64, u32>> = Arc::new(DashMap::new());
for id in 0..MAX_TUNNEL_LEGS {
let addr = remote_proxy_addr.to_string();
@@ -204,13 +214,8 @@ impl ClientHandler {
let muxer_inner = muxer.clone();
tokio::spawn(async move {
while let Some(raw_frame) = rx_from_engine.recv().await {
let dst_ip = raw_frame.dst_ip;
let dst_port = raw_frame.dst_port;
let protocol = raw_frame.protocol;
let is_udp = protocol == LocalProtocol::Udp;
if let Ok(nrxp_frame) = RawCastAdapter::to_nrxp(raw_frame) {
let local_socket_id = nrxp_frame.header.stream_id;
if let Ok(nrxp_frame) = RawCastAdapter::to_nrxp(raw_frame.clone()) {
let local_socket_id = raw_frame.socket_id;
let f_type = nrxp_frame.header.frame_type;
let payload = nrxp_frame.payload;
@@ -218,19 +223,19 @@ impl ClientHandler {
FrameType::Connect | FrameType::UdpConnect => {
let global_stream_id = muxer_inner.next_stream_id();
local_to_global.insert(local_socket_id, global_stream_id);
registry.insert(global_stream_id, (local_socket_id, dst_ip, dst_port, protocol));
registry.insert(global_stream_id, (local_socket_id, raw_frame.dst_ip, raw_frame.dst_port, raw_frame.protocol));
let (v_tx, mut v_rx) = mpsc::channel(NetworkConfig::global().client_stream_capacity);
muxer_inner.register_stream(global_stream_id, v_tx);
let tx_to_tun = tx_to_engine.clone();
let reg = registry.clone();
tokio::spawn(async move {
while let Some(back_payload) = v_rx.recv().await {
if let Some((orig_local_id, ip, port, proto)) = reg.get(&global_stream_id).map(|r| *r) {
if let Some(r) = reg.get(&global_stream_id) {
let (orig_local_id, ip, port, proto) = *r;
let out_f_type = if proto == LocalProtocol::Udp { FrameType::UdpData } else { FrameType::Data };
let mock_nrxp = Frame::new(orig_local_id, out_f_type, back_payload);
let mock_nrxp = Frame::new(orig_local_id as u32, out_f_type, back_payload);
if let Ok(raw) = RawCastAdapter::from_nrxp(mock_nrxp, ip, port, proto == LocalProtocol::Udp) {
let _ = tx_to_tun.send(raw).await;
}
@@ -240,12 +245,13 @@ impl ClientHandler {
let _ = muxer_inner.send_control(global_stream_id, f_type, payload).await;
}
FrameType::Data | FrameType::UdpData => {
if let Some(id) = local_to_global.get(&local_socket_id).map(|r| *r) {
let _ = muxer_inner.send_data_safe(id, payload, is_udp).await;
if let Some(id) = local_to_global.get(&local_socket_id) {
let _ = muxer_inner.send_data_safe(*id, payload, raw_frame.protocol == LocalProtocol::Udp).await;
}
}
FrameType::Close => {
if let Some((_, global_stream_id)) = local_to_global.remove(&local_socket_id) {
if let Some(kv) = local_to_global.remove(&local_socket_id) {
let global_stream_id = kv.1;
let _ = muxer_inner.send_control(global_stream_id, FrameType::Close, Bytes::new()).await;
muxer_inner.remove_stream(global_stream_id);
registry.remove(&global_stream_id);
@@ -260,7 +266,6 @@ impl ClientHandler {
Ok(())
}
}
pub struct ServerHandler {
pub(crate) conn: Connection,
pub(crate) session_manager: Arc<SessionManager>,
-1
View File
@@ -203,7 +203,6 @@ impl TunnelEngine {
error!(stream_id, error = %e, "Failed to write encrypted data to network");
e.to_string()
})?;
tokio::task::yield_now().await;
}
debug!(stream_id, "Outbound packet sent successfully");
+35 -33
View File
@@ -29,6 +29,21 @@ struct MuxLeg {
stats: Arc<LegStats>,
}
impl MuxLeg {
/// Возвращает коэффициент загруженности очереди (0.0 - пуста, 1.0 - полна)
fn congestion_factor(&self) -> f64 {
let cap = self.data_tx.capacity() as f64;
let max = crate::net::NetworkConfig::global().client_muxer_capacity as f64;
// Защита от деления на ноль, если конфигурация задана криво
if max <= 0.0 {
return 0.0;
}
1.0 - (cap / max)
}
}
struct IdGenerator {
counter: AtomicU32,
}
@@ -57,7 +72,6 @@ pub struct Muxer {
legs: Arc<DashMap<u32, MuxLeg>>,
streams: Arc<DashMap<u32, (Sender<Bytes>, Arc<StreamStats>)>>,
stream_bindings: Arc<DashMap<u32, u32>>,
// 🚨 Добавляем сюда:
pending_pings: Arc<DashMap<u32, Instant>>,
id_gen: Arc<IdGenerator>,
session_id: Arc<String>,
@@ -100,8 +114,7 @@ impl Muxer {
pub fn remove_leg(&self, leg_id: u32) {
self.legs.remove(&leg_id);
// Удаляем все привязки стримов к этой ноге, чтобы они перебалансировались
self.stream_bindings
.retain(|_, target_leg| *target_leg != leg_id);
self.stream_bindings.retain(|_, target_leg| *target_leg != leg_id);
info!(leg_id, "MUXER: Leg removed and bindings cleared");
}
@@ -142,14 +155,15 @@ impl Muxer {
.collect();
}
// Сортируем по RTT (Health Check записывает сюда данные)
candidates.sort_by_key(|(_, leg)| {
let rtt = leg.stats.rtt_ms.load(Ordering::Relaxed);
if rtt == 0 {
9999
} else {
rtt
}
// 🔥 Умная балансировка: RTT + Штраф за забитость канала
candidates.sort_by(|(_, leg_a), (_, leg_b)| {
let rtt_a = leg_a.stats.rtt_ms.load(Ordering::Relaxed) as f64;
let rtt_b = leg_b.stats.rtt_ms.load(Ordering::Relaxed) as f64;
let score_a = rtt_a + (leg_a.congestion_factor() * 500.0);
let score_b = rtt_b + (leg_b.congestion_factor() * 500.0);
score_a.partial_cmp(&score_b).unwrap_or(std::cmp::Ordering::Equal)
});
let pool_size = std::cmp::min(candidates.len(), MUXER_POOL_SIZE);
@@ -175,7 +189,6 @@ impl Muxer {
let rtt = start_time.elapsed().as_millis() as u32;
if let Some(leg) = self.legs.get(&leg_id) {
// Обновляем атомик, который потом прочитает принт топологии
leg.stats.rtt_ms.store(rtt, Ordering::Relaxed);
trace!(leg_id, rtt, "💓 [Muxer] RTT updated for leg");
}
@@ -190,9 +203,6 @@ impl Muxer {
.select_leg(&message.frame_type, stream_id)
.ok_or_else(|| "MUXER: No active legs available".to_string())?;
// Используем обычный .await для всех типов кадров.
// Это гарантирует, что пакет попадет в очередь ноги, но может притормозить цикл Engine,
// если очередь физического соединения переполнена.
let target_tx = match message.frame_type {
FrameType::Connect
| FrameType::Close
@@ -201,15 +211,14 @@ impl Muxer {
_ => leg.data_tx.clone(),
};
if let Err(_) = target_tx.send(message).await {
if target_tx.send(message).await.is_err() {
self.remove_leg(leg_id);
return Err(format!("MUXER: Leg {} died during send", leg_id));
}
// Обновляем статистику
leg.stats.tx_bytes.fetch_add(size, Ordering::Relaxed);
if let Some(stream_ref) = self.streams.get(&stream_id) {
stream_ref.1.tx_bytes.fetch_add(size, Ordering::Relaxed);
stream_ref.value().1.tx_bytes.fetch_add(size, Ordering::Relaxed);
}
Ok(())
@@ -250,32 +259,25 @@ impl Muxer {
}
pub fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
self.streams
.insert(stream_id, (tx, Arc::new(StreamStats::default())));
self.streams.insert(stream_id, (tx, Arc::new(StreamStats::default())));
}
pub fn remove_stream(&self, stream_id: u32) {
self.streams.remove(&stream_id);
self.stream_bindings.remove(&stream_id); // Чистим привязку
self.stream_bindings.remove(&stream_id);
trace!(stream_id, "MUXER: Stream and bindings removed");
}
pub async fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
let stream_opt = self
.streams
.get(&stream_id)
.map(|s| (s.0.clone(), s.1.clone()));
let stream_opt = self.streams.get(&stream_id).map(|s| (s.value().0.clone(), s.value().1.clone()));
if let Some((tx, stats)) = stream_opt {
let size = data.len() as u64;
match tx.send(data).await {
Ok(_) => {
stats.rx_bytes.fetch_add(size, Ordering::Relaxed);
}
Err(_) => {
self.remove_stream(stream_id);
}
if tx.send(data).await.is_ok() {
stats.rx_bytes.fetch_add(size, Ordering::Relaxed);
} else {
self.remove_stream(stream_id);
}
}
}
@@ -427,4 +429,4 @@ impl Muxer {
}
println!();
}
}
}