From 9f1d04fcafbf8912f84241735e9a06a01033fb16 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 4 Jun 2026 16:25:51 -0600 Subject: [PATCH] fix: harden tunnel performance under load Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 1 + Cargo.toml | 1 + src/cli.rs | 52 ++++ src/framing.rs | 1 - src/main.rs | 84 +++++- src/tunnel.rs | 693 +++++++++++++++++++++++++++++++++++++++---------- 6 files changed, 695 insertions(+), 137 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 88795dd..d31dc6a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1293,6 +1293,7 @@ dependencies = [ "http-body-util", "hyper", "hyper-util", + "libc", "rand", "rcgen", "rustls", diff --git a/Cargo.toml b/Cargo.toml index c0840eb..47f0faf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ sha2 = "0.10" bytes = "1" futures = "0.3" url = "2" +libc = "0.2" [dev-dependencies] tempfile = "3" diff --git a/src/cli.rs b/src/cli.rs index 0b50a36..7e16d14 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -55,6 +55,34 @@ pub enum Commands { /// Skip TLS certificate verification (insecure — for DPI/intercepted environments only) #[arg(long, default_value = "false")] insecure_skip_tls_verify: bool, + + /// Target process file descriptor limit to request on Unix + #[arg(long, default_value_t = 1_048_576)] + max_open_files: u64, + + /// Maximum concurrent authenticated HTTPS tunnel sessions + #[arg(long, default_value_t = 4096)] + max_tunnel_sessions: usize, + + /// Maximum concurrent SOCKS5 client connections + #[arg(long, default_value_t = 65_536)] + max_socks_connections: usize, + + /// Maximum concurrent streams per tunnel + #[arg(long, default_value_t = 65_536)] + max_streams_per_tunnel: usize, + + /// TLS/HTTP/SOCKS handshake timeout in milliseconds + #[arg(long, default_value_t = 10_000)] + handshake_timeout_ms: u64, + + /// DNS/target connect timeout in milliseconds + #[arg(long, default_value_t = 10_000)] + connect_timeout_ms: u64, + + /// Timeout waiting for tunneled CONNECT replies in milliseconds + #[arg(long, default_value_t = 15_000)] + stream_open_timeout_ms: u64, }, /// Connect to the HTTPS tunnel listener @@ -102,6 +130,30 @@ pub enum Commands { /// Skip TLS certificate verification (insecure — for DPI/intercepted environments only) #[arg(long, default_value = "false")] insecure_skip_tls_verify: bool, + + /// Target process file descriptor limit to request on Unix + #[arg(long, default_value_t = 1_048_576)] + max_open_files: u64, + + /// Maximum concurrent SOCKS5 client connections + #[arg(long, default_value_t = 65_536)] + max_socks_connections: usize, + + /// Maximum concurrent streams per tunnel + #[arg(long, default_value_t = 65_536)] + max_streams_per_tunnel: usize, + + /// TLS/HTTP/SOCKS handshake timeout in milliseconds + #[arg(long, default_value_t = 10_000)] + handshake_timeout_ms: u64, + + /// DNS/target connect timeout in milliseconds + #[arg(long, default_value_t = 10_000)] + connect_timeout_ms: u64, + + /// Timeout waiting for tunneled CONNECT replies in milliseconds + #[arg(long, default_value_t = 15_000)] + stream_open_timeout_ms: u64, }, /// Generate local certificate and auth material diff --git a/src/framing.rs b/src/framing.rs index 008b47d..459b75e 100644 --- a/src/framing.rs +++ b/src/framing.rs @@ -268,7 +268,6 @@ impl FrameWriter { ) -> Result<(), FrameError> { let bytes = frame.serialize(); writer.write_all(&bytes).await?; - writer.flush().await?; Ok(()) } } diff --git a/src/main.rs b/src/main.rs index 7c78da0..656f687 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,12 +13,15 @@ mod tunnel; use std::path::Path; use std::process; use std::sync::Arc; +use std::time::Duration; use clap::Parser; use crate::cli::Cli; use crate::connkey::ConnectionKey; -use crate::tunnel::{ClientTlsMaterial, ConnectorConfig, ListenerConfig, ServerTlsMaterial}; +use crate::tunnel::{ + ClientTlsMaterial, ConnectorConfig, ListenerConfig, PerformanceConfig, ServerTlsMaterial, +}; fn main() { // Install the default crypto provider (ring) for rustls @@ -47,6 +50,13 @@ fn main() { auth_token, auth_token_file, insecure_skip_tls_verify, + max_open_files, + max_tunnel_sessions, + max_socks_connections, + max_streams_per_tunnel, + handshake_timeout_ms, + connect_timeout_ms, + stream_open_timeout_ms, } => run_listen( listen, advertise, @@ -58,6 +68,13 @@ fn main() { auth_token, auth_token_file, insecure_skip_tls_verify, + max_open_files, + max_tunnel_sessions, + max_socks_connections, + max_streams_per_tunnel, + handshake_timeout_ms, + connect_timeout_ms, + stream_open_timeout_ms, ), crate::cli::Commands::Connect { connection_key_arg, @@ -70,6 +87,12 @@ fn main() { auth_token, auth_token_file, insecure_skip_tls_verify, + max_open_files, + max_socks_connections, + max_streams_per_tunnel, + handshake_timeout_ms, + connect_timeout_ms, + stream_open_timeout_ms, } => run_connect( connection_key_arg, target, @@ -81,6 +104,12 @@ fn main() { auth_token, auth_token_file, insecure_skip_tls_verify, + max_open_files, + max_socks_connections, + max_streams_per_tunnel, + handshake_timeout_ms, + connect_timeout_ms, + stream_open_timeout_ms, ), crate::cli::Commands::Generate { out, @@ -112,7 +141,24 @@ fn run_listen( auth_token: String, auth_token_file: Option, insecure_skip_tls_verify: bool, + max_open_files: u64, + max_tunnel_sessions: usize, + max_socks_connections: usize, + max_streams_per_tunnel: usize, + handshake_timeout_ms: u64, + connect_timeout_ms: u64, + stream_open_timeout_ms: u64, ) -> i32 { + configure_performance( + max_open_files, + max_tunnel_sessions, + max_socks_connections, + max_streams_per_tunnel, + handshake_timeout_ms, + connect_timeout_ms, + stream_open_timeout_ms, + ); + let (host, port) = match cli::parse_host_port(&listen) { Ok(v) => v, Err(e) => { @@ -310,7 +356,23 @@ fn run_connect( auth_token: String, auth_token_file: Option, insecure_skip_tls_verify: bool, + max_open_files: u64, + max_socks_connections: usize, + max_streams_per_tunnel: usize, + handshake_timeout_ms: u64, + connect_timeout_ms: u64, + stream_open_timeout_ms: u64, ) -> i32 { + configure_performance( + max_open_files, + 1, + max_socks_connections, + max_streams_per_tunnel, + handshake_timeout_ms, + connect_timeout_ms, + stream_open_timeout_ms, + ); + let raw_connection_key = connection_key_arg.or(connection_key); let decoded_key = if let Some(raw) = raw_connection_key { match ConnectionKey::decode(&raw) { @@ -465,6 +527,26 @@ fn run_connect( 0 } +fn configure_performance( + max_open_files: u64, + max_tunnel_sessions: usize, + max_socks_connections: usize, + max_streams_per_tunnel: usize, + handshake_timeout_ms: u64, + connect_timeout_ms: u64, + stream_open_timeout_ms: u64, +) { + tunnel::set_performance_config(PerformanceConfig { + max_open_files, + max_tunnel_sessions, + max_socks_connections, + max_streams_per_tunnel, + handshake_timeout: Duration::from_millis(handshake_timeout_ms), + connect_timeout: Duration::from_millis(connect_timeout_ms), + stream_open_timeout: Duration::from_millis(stream_open_timeout_ms), + }); +} + fn select_connect_target<'a>( target: Option<&'a str>, decoded_key: Option<&'a ConnectionKey>, diff --git a/src/tunnel.rs b/src/tunnel.rs index 620fed8..7494cca 100644 --- a/src/tunnel.rs +++ b/src/tunnel.rs @@ -18,16 +18,18 @@ /// 2. Client certificate validation (server side, mTLS) /// 3. Server certificate validation (client side) /// 4. Application auth token validation -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::net::SocketAddr; use std::path::Path; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::time::Duration; use bytes::Bytes; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::{mpsc, oneshot, watch}; +use tokio::sync::{Semaphore, mpsc, oneshot, watch}; use crate::errors::{AuthError, HostError, TunnelError}; use crate::framing::{ @@ -131,6 +133,96 @@ pub struct ConnectorConfig { pub insecure_skip_tls_verify: bool, } +/// Runtime performance limits and timeouts. +#[derive(Debug, Clone)] +pub struct PerformanceConfig { + pub max_open_files: u64, + pub max_tunnel_sessions: usize, + pub max_socks_connections: usize, + pub max_streams_per_tunnel: usize, + pub handshake_timeout: Duration, + pub connect_timeout: Duration, + pub stream_open_timeout: Duration, +} + +impl Default for PerformanceConfig { + fn default() -> Self { + Self { + max_open_files: 1_048_576, + max_tunnel_sessions: 4096, + max_socks_connections: 65_536, + max_streams_per_tunnel: 65_536, + handshake_timeout: Duration::from_secs(10), + connect_timeout: Duration::from_secs(10), + stream_open_timeout: Duration::from_secs(15), + } + } +} + +static PERFORMANCE_CONFIG: OnceLock = OnceLock::new(); + +pub fn set_performance_config(config: PerformanceConfig) { + let _ = PERFORMANCE_CONFIG.set(config); +} + +fn performance_config() -> PerformanceConfig { + PERFORMANCE_CONFIG.get().cloned().unwrap_or_default() +} + +#[cfg(unix)] +pub fn raise_nofile_limit(target: u64) -> std::io::Result<()> { + let mut limits = libc::rlimit { + rlim_cur: 0, + rlim_max: 0, + }; + let rc = unsafe { libc::getrlimit(libc::RLIMIT_NOFILE, &mut limits) }; + if rc != 0 { + return Err(std::io::Error::last_os_error()); + } + + let target = target as libc::rlim_t; + let new_max = limits.rlim_max.max(target); + let new_cur = limits.rlim_cur.max(target).min(new_max); + let updated = libc::rlimit { + rlim_cur: new_cur, + rlim_max: new_max, + }; + let rc = unsafe { libc::setrlimit(libc::RLIMIT_NOFILE, &updated) }; + if rc != 0 { + return Err(std::io::Error::last_os_error()); + } + Ok(()) +} + +#[cfg(not(unix))] +pub fn raise_nofile_limit(_target: u64) -> std::io::Result<()> { + Ok(()) +} + +fn is_fd_exhaustion(error: &std::io::Error) -> bool { + error.raw_os_error() == Some(24) +} + +async fn accept_error_backoff(error: &std::io::Error, delay: &mut Duration, label: &str) { + if is_fd_exhaustion(error) { + tracing::error!( + "{} accept hit file descriptor exhaustion: {} — backing off for {:?}", + label, + error, + delay + ); + } else { + tracing::error!( + "{} accept error: {} — backing off for {:?}", + label, + error, + delay + ); + } + tokio::time::sleep(*delay).await; + *delay = (*delay * 2).min(Duration::from_secs(1)); +} + #[derive(Debug, Clone)] pub enum ServerTlsMaterial { Paths { @@ -396,6 +488,14 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { let addr = config.bind_addr; let tls_material = config.tls.clone(); let auth_token = config.auth_token.clone(); + let performance = performance_config(); + if let Err(e) = raise_nofile_limit(performance.max_open_files) { + tracing::warn!( + "Failed to raise RLIMIT_NOFILE to {}: {}", + performance.max_open_files, + e + ); + } let tls_config = tls_material .build_config(config.insecure_skip_tls_verify) @@ -434,6 +534,16 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { addr, Redacted::new(&*auth_token) ); + tracing::info!( + "Performance config: max_open_files={}, max_tunnel_sessions={}, max_socks_connections={}, max_streams_per_tunnel={}, handshake_timeout={:?}, connect_timeout={:?}, stream_open_timeout={:?}", + performance.max_open_files, + performance.max_tunnel_sessions, + performance.max_socks_connections, + performance.max_streams_per_tunnel, + performance.handshake_timeout, + performance.connect_timeout, + performance.stream_open_timeout + ); tracing::info!( "TLS material fingerprints: ca={}, server={}", tls::cert_fingerprint_sha256(&ca_certs[0]), @@ -455,6 +565,8 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { // Listen for shutdown signal let shutdown = crate::signal::shutdown_signal(); tokio::pin!(shutdown); + let session_permits = Arc::new(Semaphore::new(performance.max_tunnel_sessions)); + let mut accept_delay = Duration::from_millis(10); loop { tokio::select! { @@ -465,6 +577,18 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { result = listener.accept() => { match result { Ok((stream, peer_addr)) => { + accept_delay = Duration::from_millis(10); + if let Err(e) = stream.set_nodelay(true) { + tracing::debug!("Failed to set TCP_NODELAY for {}: {}", peer_addr, e); + } + let Ok(session_permit) = session_permits.clone().try_acquire_owned() else { + tracing::warn!( + "Connection from {} rejected: max tunnel sessions ({}) reached", + peer_addr, + performance.max_tunnel_sessions + ); + continue; + }; tracing::info!("New connection from {}", peer_addr); let acceptor = tls_acceptor.clone(); let token = auth_token.clone(); @@ -476,13 +600,14 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { let insecure = config.insecure_skip_tls_verify; tokio::spawn(async move { + let _session_permit = session_permit; if let Err(e) = handle_listener_connection(stream, peer_addr, acceptor, token, session_mux_sender, insecure).await { tracing::warn!("Connection from {} failed: {}", peer_addr, e); } }); } Err(e) => { - tracing::error!("Accept error: {}", e); + accept_error_backoff(&e, &mut accept_delay, "HTTPS").await; } } } @@ -509,11 +634,15 @@ async fn handle_listener_connection( mux_sender: Option>>>, insecure_skip_tls_verify: bool, ) -> Result<(), TunnelError> { + let performance = performance_config(); // TLS accept - let tls_stream = acceptor.accept(stream).await.map_err(|e| { - tracing::warn!("TLS accept failed for {}: {}", peer_addr, e); - TunnelError::Tls(crate::errors::TlsError::VerificationFailed(e.to_string())) - })?; + let tls_stream = tokio::time::timeout(performance.handshake_timeout, acceptor.accept(stream)) + .await + .map_err(|_| TunnelError::Protocol("TLS accept timed out".into()))? + .map_err(|e| { + tracing::warn!("TLS accept failed for {}: {}", peer_addr, e); + TunnelError::Tls(crate::errors::TlsError::VerificationFailed(e.to_string())) + })?; if insecure_skip_tls_verify { tracing::info!( @@ -540,7 +669,12 @@ async fn handle_listener_connection( Box::pin(tls_stream); // Phase 1: HTTP auth handshake - let (method, path, headers) = parse_http_request(&mut *tls_stream).await?; + let (method, path, headers) = tokio::time::timeout( + performance.handshake_timeout, + parse_http_request(&mut *tls_stream), + ) + .await + .map_err(|_| TunnelError::Protocol("HTTP auth handshake timed out".into()))??; if method == "POST" && path == "/tunnel" { let received_token = headers @@ -585,8 +719,12 @@ async fn handle_listener_connection( /// Entry for an active stream on the listener side. struct ServerStreamEntry { - /// Write half of the target TCP connection - target_write: tokio::net::tcp::OwnedWriteHalf, + target_tx: mpsc::Sender, +} + +enum ServerStreamUpdate { + Opened(u64, ServerStreamEntry), + Failed(u64), } /// Handle the binary framing protocol on the server side. @@ -601,12 +739,14 @@ where let mut stream = stream; let mut reader = FrameReader::new(); let (outbound_tx, mut outbound_rx) = mpsc::channel::(256); + let (stream_tx, mut stream_rx) = mpsc::channel::(256); let mux = mux_sender .as_ref() .map(|_| Arc::new(StreamMux::new(outbound_tx.clone()))); // Active streams: stream_id → target write half let mut streams: HashMap = HashMap::new(); + let mut pending_streams: HashSet = HashSet::new(); tracing::info!("Framing session started with {}", peer_addr); if let (Some(sender), Some(mux)) = (&mux_sender, &mux) { @@ -623,6 +763,18 @@ where } } + Some(update) = stream_rx.recv() => { + match update { + ServerStreamUpdate::Opened(stream_id, entry) => { + pending_streams.remove(&stream_id); + streams.insert(stream_id, entry); + } + ServerStreamUpdate::Failed(stream_id) => { + pending_streams.remove(&stream_id); + } + } + } + // Read inbound frames from the tunnel (CONNECT, DATA from client, CLOSE) result = reader.read_frame(&mut stream) => { match result { @@ -639,7 +791,9 @@ where &frame, &mut stream, &outbound_tx, + &stream_tx, &mut streams, + &mut pending_streams, peer_addr, ).await } @@ -648,7 +802,9 @@ where &frame, &mut stream, &outbound_tx, + &stream_tx, &mut streams, + &mut pending_streams, peer_addr, ).await }; @@ -678,8 +834,7 @@ where } // Clean up all streams - for (sid, mut entry) in streams.drain() { - let _ = entry.target_write.shutdown().await; + for (sid, _entry) in streams.drain() { tracing::debug!("Closed stream {} on shutdown", sid); } if let Some(sender) = mux_sender { @@ -695,7 +850,9 @@ async fn handle_inbound_frame( frame: &Frame, _stream: &mut S, outbound_tx: &mpsc::Sender, + stream_tx: &mpsc::Sender, streams: &mut HashMap, + pending_streams: &mut HashSet, peer_addr: SocketAddr, ) -> Result<(), TunnelError> where @@ -703,7 +860,8 @@ where { match frame.frame_type { FRAME_CONNECT => { - if streams.contains_key(&frame.stream_id) { + if streams.contains_key(&frame.stream_id) || pending_streams.contains(&frame.stream_id) + { let _ = outbound_tx .send(Frame::error(frame.stream_id, "duplicate stream ID")) .await; @@ -721,100 +879,29 @@ where target_str ); - // Resolve and connect to target from the listener side - let target_addr = match socks5::resolve_target(&target).await { - Ok(addr) => addr, - Err(e) => { - tracing::warn!( - "Stream {} resolve failed for {}: {}", - frame.stream_id, - target_str, - e - ); - let _ = outbound_tx - .send(Frame::connect_reply(frame.stream_id, CONNECT_FAILED)) - .await; - let _ = outbound_tx - .send(Frame::error(frame.stream_id, format!("resolve: {}", e))) - .await; - return Ok(()); - } - }; - - match TcpStream::connect(target_addr).await { - Ok(target_stream) => { - tracing::info!( - "Stream {} connected to target {} (listener side)", - frame.stream_id, - target_addr - ); - - let (mut target_read, target_write) = target_stream.into_split(); - let sid = frame.stream_id; - let outbound = outbound_tx.clone(); - - // Spawn target → tunnel direction - tokio::spawn(async move { - let mut buf = [0u8; 32 * 1024]; - loop { - match target_read.read(&mut buf).await { - Ok(0) => break, // Target closed (EOF) - Ok(n) => { - if outbound.send(Frame::data(sid, &buf[..n])).await.is_err() { - break; - } - } - Err(e) => { - tracing::debug!("Stream {} target read error: {}", sid, e); - let _ = outbound.send(Frame::error(sid, e.to_string())).await; - break; - } - } - } - tracing::info!("Stream {} target → tunnel EOF", sid); - let _ = outbound.send(Frame::close(sid)).await; - }); - - streams.insert(frame.stream_id, ServerStreamEntry { target_write }); - let _ = outbound_tx - .send(Frame::connect_reply(frame.stream_id, CONNECT_OK)) - .await; - } - Err(e) => { - let err_msg = if e.kind() == std::io::ErrorKind::ConnectionRefused { - format!("connection refused to {}", target_addr) - } else { - format!("connect to {}: {}", target_addr, e) - }; - tracing::warn!( - "Stream {} connect to {} failed: {}", - frame.stream_id, - target_addr, - err_msg - ); - let _ = outbound_tx - .send(Frame::connect_reply(frame.stream_id, CONNECT_FAILED)) - .await; - let _ = outbound_tx - .send(Frame::error(frame.stream_id, err_msg)) - .await; - } - } + pending_streams.insert(frame.stream_id); + spawn_target_connect( + frame.stream_id, + target, + target_str, + peer_addr, + outbound_tx.clone(), + stream_tx.clone(), + ); } FRAME_DATA => { // Forward data to the target if let Some(entry) = streams.get_mut(&frame.stream_id) - && let Err(e) = entry.target_write.write_all(&frame.payload).await + && let Err(e) = entry.target_tx.try_send(frame.payload.clone()) { - tracing::debug!("Stream {} target write error: {}", frame.stream_id, e); + tracing::debug!("Stream {} target write queue error: {}", frame.stream_id, e); streams.remove(&frame.stream_id); } } FRAME_CLOSE => { - if let Some(mut entry) = streams.remove(&frame.stream_id) { - let _ = entry.target_write.shutdown().await; + if streams.remove(&frame.stream_id).is_some() { tracing::info!( "Stream {} closed by client (sender-side EOF)", frame.stream_id @@ -825,9 +912,7 @@ where FRAME_ERROR => { let msg = String::from_utf8_lossy(&frame.payload); tracing::warn!("Stream {} error from client: {}", frame.stream_id, msg); - if let Some(mut entry) = streams.remove(&frame.stream_id) { - let _ = entry.target_write.shutdown().await; - } + streams.remove(&frame.stream_id); } _ => { @@ -843,6 +928,163 @@ where Ok(()) } +fn spawn_target_connect( + stream_id: u64, + target: socks5::Socks5Target, + target_str: String, + peer_addr: SocketAddr, + outbound_tx: mpsc::Sender, + stream_tx: mpsc::Sender, +) { + tokio::spawn(async move { + let performance = performance_config(); + let target_addr = match tokio::time::timeout( + performance.connect_timeout, + socks5::resolve_target(&target), + ) + .await + { + Ok(Ok(addr)) => addr, + Ok(Err(e)) => { + tracing::warn!( + "Stream {} resolve failed for {}: {}", + stream_id, + target_str, + e + ); + let _ = outbound_tx + .send(Frame::connect_reply(stream_id, CONNECT_FAILED)) + .await; + let _ = outbound_tx + .send(Frame::error(stream_id, format!("resolve: {}", e))) + .await; + let _ = stream_tx.send(ServerStreamUpdate::Failed(stream_id)).await; + return; + } + Err(e) => { + tracing::warn!( + "Stream {} resolve timed out for {} after {:?}: {}", + stream_id, + target_str, + performance.connect_timeout, + e + ); + let _ = outbound_tx + .send(Frame::connect_reply(stream_id, CONNECT_FAILED)) + .await; + let _ = outbound_tx + .send(Frame::error(stream_id, "resolve timed out")) + .await; + let _ = stream_tx.send(ServerStreamUpdate::Failed(stream_id)).await; + return; + } + }; + + match tokio::time::timeout(performance.connect_timeout, TcpStream::connect(target_addr)) + .await + { + Ok(Ok(target_stream)) => { + if let Err(e) = target_stream.set_nodelay(true) { + tracing::debug!( + "Stream {} failed to set TCP_NODELAY for target {}: {}", + stream_id, + target_addr, + e + ); + } + tracing::info!( + "Stream {} connected from {} to target {}", + stream_id, + peer_addr, + target_addr + ); + + let (mut target_read, mut target_write) = target_stream.into_split(); + let outbound = outbound_tx.clone(); + let (target_tx, mut target_rx) = mpsc::channel::(64); + + tokio::spawn(async move { + let mut buf = [0u8; 32 * 1024]; + loop { + match target_read.read(&mut buf).await { + Ok(0) => break, + Ok(n) => { + if outbound + .send(Frame::data(stream_id, &buf[..n])) + .await + .is_err() + { + break; + } + } + Err(e) => { + tracing::debug!("Stream {} target read error: {}", stream_id, e); + let _ = outbound.send(Frame::error(stream_id, e.to_string())).await; + break; + } + } + } + tracing::info!("Stream {} target → tunnel EOF", stream_id); + let _ = outbound.send(Frame::close(stream_id)).await; + }); + + tokio::spawn(async move { + while let Some(data) = target_rx.recv().await { + if let Err(e) = target_write.write_all(&data).await { + tracing::debug!("Stream {} target write error: {}", stream_id, e); + break; + } + } + let _ = target_write.shutdown().await; + }); + + if stream_tx + .send(ServerStreamUpdate::Opened( + stream_id, + ServerStreamEntry { target_tx }, + )) + .await + .is_ok() + { + let _ = outbound_tx + .send(Frame::connect_reply(stream_id, CONNECT_OK)) + .await; + } + } + Ok(Err(e)) => { + let err_msg = if e.kind() == std::io::ErrorKind::ConnectionRefused { + format!("connection refused to {}", target_addr) + } else { + format!("connect to {}: {}", target_addr, e) + }; + tracing::warn!( + "Stream {} connect to {} failed: {}", + stream_id, + target_addr, + err_msg + ); + let _ = outbound_tx + .send(Frame::connect_reply(stream_id, CONNECT_FAILED)) + .await; + let _ = outbound_tx.send(Frame::error(stream_id, err_msg)).await; + let _ = stream_tx.send(ServerStreamUpdate::Failed(stream_id)).await; + } + Err(_) => { + let err_msg = format!( + "connect to {} timed out after {:?}", + target_addr, performance.connect_timeout + ); + tracing::warn!("Stream {} {}", stream_id, err_msg); + let _ = outbound_tx + .send(Frame::connect_reply(stream_id, CONNECT_FAILED)) + .await; + let _ = outbound_tx.send(Frame::error(stream_id, err_msg)).await; + let _ = stream_tx.send(ServerStreamUpdate::Failed(stream_id)).await; + } + } + }); +} + // --------------------------------------------------------------------------- // Connector: StreamMux for multiplexed streams // --------------------------------------------------------------------------- @@ -853,6 +1095,7 @@ struct ConnectorStreamEntry { data_tx: mpsc::Sender, /// One-shot channel for CONNECT_REPLY (Option so we can take it for send) reply_tx: Option>, + _permit: tokio::sync::OwnedSemaphorePermit, } /// Stream multiplexer for the connector side. @@ -862,6 +1105,7 @@ pub struct StreamMux { outbound_tx: mpsc::Sender, /// Active streams registry streams: Arc>>, + stream_permits: Arc, /// Monotonically increasing stream ID counter (odd IDs for client-initiated) next_id: Arc, /// Auth state @@ -874,6 +1118,7 @@ impl StreamMux { Self { outbound_tx, streams: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + stream_permits: Arc::new(Semaphore::new(performance_config().max_streams_per_tunnel)), next_id: Arc::new(AtomicU64::new(1)), auth_complete: Arc::new(AtomicBool::new(true)), } @@ -888,6 +1133,11 @@ impl StreamMux { if !self.auth_complete.load(Ordering::SeqCst) { return Err(TunnelError::Auth(AuthError::NotAuthenticated)); } + let permit = self + .stream_permits + .clone() + .try_acquire_owned() + .map_err(|_| TunnelError::Protocol("max streams per tunnel reached".into()))?; let stream_id = self.next_id.fetch_add(2, Ordering::Relaxed); // odd IDs let (reply_tx, reply_rx) = oneshot::channel(); @@ -898,18 +1148,32 @@ impl StreamMux { ConnectorStreamEntry { data_tx, reply_tx: Some(reply_tx), + _permit: permit, }, ); let target_bytes = encode_target(target); - self.outbound_tx + if self + .outbound_tx .send(Frame::connect(stream_id, target_bytes)) .await - .map_err(|_| TunnelError::Protocol("outbound channel closed".into()))?; + .is_err() + { + self.streams.lock().await.remove(&stream_id); + return Err(TunnelError::Protocol("outbound channel closed".into())); + } - let status = reply_rx + let status = tokio::time::timeout(performance_config().stream_open_timeout, reply_rx) .await - .map_err(|_| TunnelError::Protocol("connect reply channel closed".into()))?; + .map_err(|_| TunnelError::Protocol("CONNECT reply timed out".into()))? + .map_err(|_| TunnelError::Protocol("connect reply channel closed".into())); + let status = match status { + Ok(status) => status, + Err(e) => { + self.streams.lock().await.remove(&stream_id); + return Err(e); + } + }; if status != CONNECT_OK { self.streams.lock().await.remove(&stream_id); @@ -940,13 +1204,19 @@ impl StreamMux { } } FRAME_DATA => { - let mut streams = self.streams.lock().await; - if let Some(entry) = streams.get(&frame.stream_id) - && entry.data_tx.send(frame.payload.clone()).await.is_err() + let data_tx = { + let streams = self.streams.lock().await; + streams + .get(&frame.stream_id) + .map(|entry| entry.data_tx.clone()) + }; + if let Some(data_tx) = data_tx + && data_tx.send(frame.payload.clone()).await.is_err() { - // SOCKS5 handler dropped — clean up - streams.remove(&frame.stream_id); - let _ = self.outbound_tx.send(Frame::close(frame.stream_id)).await; + let mut streams = self.streams.lock().await; + if streams.remove(&frame.stream_id).is_some() { + let _ = self.outbound_tx.send(Frame::close(frame.stream_id)).await; + } } } FRAME_CLOSE => { @@ -997,6 +1267,14 @@ pub async fn run_connector_with_socks( config: ConnectorConfig, socks_addr: SocketAddr, ) -> Result<(), TunnelError> { + let performance = performance_config(); + if let Err(e) = raise_nofile_limit(performance.max_open_files) { + tracing::warn!( + "Failed to raise RLIMIT_NOFILE to {}: {}", + performance.max_open_files, + e + ); + } let client_config = config .tls .build_config(config.insecure_skip_tls_verify) @@ -1017,6 +1295,15 @@ pub async fn run_connector_with_socks( socks_addr, Redacted::new(&*config.auth_token) ); + tracing::info!( + "Performance config: max_open_files={}, max_socks_connections={}, max_streams_per_tunnel={}, handshake_timeout={:?}, connect_timeout={:?}, stream_open_timeout={:?}", + performance.max_open_files, + performance.max_socks_connections, + performance.max_streams_per_tunnel, + performance.handshake_timeout, + performance.connect_timeout, + performance.stream_open_timeout + ); tracing::info!( "TLS material fingerprints: ca={}, client={}", tls::cert_fingerprint_sha256(&ca_certs[0]), @@ -1087,9 +1374,24 @@ async fn run_reconnect_loop( ); // Connect with fresh TLS (re-runs all security gates) - let stream = match TcpStream::connect(target_addr).await { - Ok(s) => s, - Err(e) => { + let performance = performance_config(); + let stream = match tokio::time::timeout( + performance.connect_timeout, + TcpStream::connect(target_addr), + ) + .await + { + Ok(Ok(s)) => { + if let Err(e) = s.set_nodelay(true) { + tracing::debug!( + "Failed to set TCP_NODELAY for tunnel {}: {}", + target_addr, + e + ); + } + s + } + Ok(Err(e)) => { tracing::warn!( "Tunnel connect failed to {}: {} — retrying in {:?}", target_addr, @@ -1100,12 +1402,28 @@ async fn run_reconnect_loop( reconnect_delay = (reconnect_delay * 2).min(max_delay); continue; } + Err(_) => { + tracing::warn!( + "Tunnel connect to {} timed out after {:?} — retrying in {:?}", + target_addr, + performance.connect_timeout, + reconnect_delay + ); + tokio::time::sleep(reconnect_delay).await; + reconnect_delay = (reconnect_delay * 2).min(max_delay); + continue; + } }; let tls_connector = tokio_rustls::TlsConnector::from(client_config.clone()); - let tls_stream = match tls_connector.connect(server_name.clone(), stream).await { - Ok(s) => s, - Err(e) => { + let tls_stream = match tokio::time::timeout( + performance.handshake_timeout, + tls_connector.connect(server_name.clone(), stream), + ) + .await + { + Ok(Ok(s)) => s, + Ok(Err(e)) => { if let Some(hint) = tls_handshake_hint(&e.to_string()) { tracing::warn!("{}", hint); } @@ -1120,6 +1438,18 @@ async fn run_reconnect_loop( reconnect_delay = (reconnect_delay * 2).min(max_delay); continue; } + Err(_) => { + tracing::warn!( + "TLS handshake to {}:{} timed out after {:?} — retrying in {:?}", + config.target_host, + config.target_port, + performance.handshake_timeout, + reconnect_delay + ); + tokio::time::sleep(reconnect_delay).await; + reconnect_delay = (reconnect_delay * 2).min(max_delay); + continue; + } }; tracing::info!( @@ -1221,7 +1551,12 @@ async fn authenticate_tunnel( stream.write_all(request.as_bytes()).await?; // Read HTTP response - let status = read_http_response_status(&mut *stream).await?; + let status = tokio::time::timeout( + performance_config().handshake_timeout, + read_http_response_status(&mut *stream), + ) + .await + .map_err(|_| TunnelError::Protocol("HTTP auth response timed out".into()))??; if status != 200 { tracing::warn!( @@ -1264,8 +1599,10 @@ async fn run_framing_client( S: AsyncReadExt + AsyncWriteExt + Unpin, { let mut reader = FrameReader::new(); - let (target_tx, mut target_rx) = mpsc::channel::(256); + let (target_frame_tx, mut target_frame_rx) = mpsc::channel::(256); + let (stream_tx, mut stream_rx) = mpsc::channel::(256); let mut target_streams: HashMap = HashMap::new(); + let mut pending_streams: HashSet = HashSet::new(); loop { tokio::select! { @@ -1277,13 +1614,25 @@ async fn run_framing_client( } } - Some(frame) = target_rx.recv() => { + Some(frame) = target_frame_rx.recv() => { if let Err(e) = FrameWriter::write_frame(&mut stream, &frame).await { tracing::debug!("Client target write error for stream {}: {}", frame.stream_id, e); break; } } + Some(update) = stream_rx.recv() => { + match update { + ServerStreamUpdate::Opened(stream_id, entry) => { + pending_streams.remove(&stream_id); + target_streams.insert(stream_id, entry); + } + ServerStreamUpdate::Failed(stream_id) => { + pending_streams.remove(&stream_id); + } + } + } + // Read inbound frames and dispatch result = reader.read_frame(&mut stream) => { match result { @@ -1294,8 +1643,10 @@ async fn run_framing_client( if let Err(e) = handle_inbound_frame( &frame, &mut stream, - &target_tx, + &target_frame_tx, + &stream_tx, &mut target_streams, + &mut pending_streams, SocketAddr::from(([0, 0, 0, 0], 0)), ).await { tracing::warn!( @@ -1304,7 +1655,7 @@ async fn run_framing_client( peer_name, e ); - let _ = target_tx + let _ = target_frame_tx .send(Frame::error(frame.stream_id, e.to_string())) .await; } @@ -1325,8 +1676,7 @@ async fn run_framing_client( } } - for (sid, mut entry) in target_streams.drain() { - let _ = entry.target_write.shutdown().await; + for (sid, _entry) in target_streams.drain() { tracing::debug!("Closed connector-side target stream {} on shutdown", sid); } @@ -1343,6 +1693,7 @@ async fn run_socks5_proxy( socks_addr: SocketAddr, mux_receiver: watch::Receiver>>, ) -> Result<(), TunnelError> { + let performance = performance_config(); let socks_listener = TcpListener::bind(socks_addr).await.map_err(|e| { if e.kind() == std::io::ErrorKind::AddrInUse { TunnelError::PortInUse(socks_addr.port()) @@ -1358,6 +1709,8 @@ async fn run_socks5_proxy( let shutdown = crate::signal::shutdown_signal(); tokio::pin!(shutdown); + let socks_permits = Arc::new(Semaphore::new(performance.max_socks_connections)); + let mut accept_delay = Duration::from_millis(10); loop { tokio::select! { @@ -1371,11 +1724,24 @@ async fn run_socks5_proxy( result = socks_listener.accept() => { match result { Ok((stream, peer_addr)) => { + accept_delay = Duration::from_millis(10); + if let Err(e) = stream.set_nodelay(true) { + tracing::debug!("Failed to set TCP_NODELAY for SOCKS5 {}: {}", peer_addr, e); + } + let Ok(socks_permit) = socks_permits.clone().try_acquire_owned() else { + tracing::warn!( + "SOCKS5 connection from {} rejected: max SOCKS5 connections ({}) reached", + peer_addr, + performance.max_socks_connections + ); + continue; + }; let mux_opt = mux_receiver.borrow().clone(); match mux_opt { Some(ref m) => { let m = m.clone(); tokio::spawn(async move { + let _socks_permit = socks_permit; if let Err(e) = handle_socks5_connection(stream, peer_addr, m).await { tracing::warn!("SOCKS5 connection from {} failed: {}", peer_addr, e); } @@ -1390,7 +1756,7 @@ async fn run_socks5_proxy( } } Err(e) => { - tracing::error!("SOCKS5 accept error: {}", e); + accept_error_backoff(&e, &mut accept_delay, "SOCKS5").await; } } } @@ -1422,7 +1788,12 @@ async fn handle_socks5_connection( // --- SOCKS5 Handshake --- let (mut socks_read, mut socks_write) = stream.into_split(); - let target = socks5_handshake(&mut socks_read, &mut socks_write).await?; + let target = tokio::time::timeout( + performance_config().handshake_timeout, + socks5_handshake(&mut socks_read, &mut socks_write), + ) + .await + .map_err(|_| socks5::Socks5Error::Protocol("SOCKS5 handshake timed out".into()))??; let target_str = socks5::target_to_string(&target); tracing::info!( @@ -1638,23 +2009,39 @@ pub async fn connect_tunnel(config: ConnectorConfig) -> Result<(), TunnelError> let target_addr = resolve_target(&config.target_host, config.target_port).await?; let server_name = tls::server_name_from_host(&config.target_host).map_err(TunnelError::Tls)?; - let stream = TcpStream::connect(target_addr).await.map_err(|e| { + let stream = tokio::time::timeout( + performance_config().connect_timeout, + TcpStream::connect(target_addr), + ) + .await + .map_err(|_| TunnelError::Protocol("connect timed out".into()))? + .map_err(|e| { if e.kind() == std::io::ErrorKind::ConnectionRefused { TunnelError::ConnectionRefused(target_addr.to_string()) } else { TunnelError::Io(e) } })?; + if let Err(e) = stream.set_nodelay(true) { + tracing::debug!( + "Failed to set TCP_NODELAY for tunnel {}: {}", + target_addr, + e + ); + } - let tls_stream = tls_connector - .connect(server_name.clone(), stream) - .await - .map_err(|e| { - TunnelError::Tls(crate::errors::TlsError::VerificationFailed(format!( - "TLS handshake to {}:{} failed: {}", - config.target_host, config.target_port, e - ))) - })?; + let tls_stream = tokio::time::timeout( + performance_config().handshake_timeout, + tls_connector.connect(server_name.clone(), stream), + ) + .await + .map_err(|_| TunnelError::Protocol("TLS handshake timed out".into()))? + .map_err(|e| { + TunnelError::Tls(crate::errors::TlsError::VerificationFailed(format!( + "TLS handshake to {}:{} failed: {}", + config.target_host, config.target_port, e + ))) + })?; // Authenticate let mut stream: std::pin::Pin>> = @@ -1665,7 +2052,12 @@ pub async fn connect_tunnel(config: ConnectorConfig) -> Result<(), TunnelError> config.target_host, config.auth_token ); stream.write_all(request.as_bytes()).await?; - let status = read_http_response_status(&mut *stream).await?; + let status = tokio::time::timeout( + performance_config().handshake_timeout, + read_http_response_status(&mut *stream), + ) + .await + .map_err(|_| TunnelError::Protocol("HTTP auth response timed out".into()))??; if status != 200 { return Err(TunnelError::Auth(AuthError::Invalid)); @@ -1763,6 +2155,37 @@ mod tests { assert!(hint.contains("does not match the target host")); } + #[test] + fn fd_exhaustion_detection_matches_emfile() { + let emfile = std::io::Error::from_raw_os_error(24); + assert!(is_fd_exhaustion(&emfile)); + + let other = std::io::Error::from(std::io::ErrorKind::ConnectionRefused); + assert!(!is_fd_exhaustion(&other)); + } + + #[tokio::test] + async fn accept_backoff_increases_and_caps() { + let err = std::io::Error::from_raw_os_error(24); + let mut delay = std::time::Duration::from_millis(1); + + accept_error_backoff(&err, &mut delay, "test").await; + assert_eq!(delay, std::time::Duration::from_millis(2)); + + let mut delay = std::time::Duration::from_secs(1); + accept_error_backoff(&err, &mut delay, "test").await; + assert_eq!(delay, std::time::Duration::from_secs(1)); + } + + #[test] + fn performance_defaults_are_high_volume_oriented() { + let config = PerformanceConfig::default(); + assert!(config.max_open_files >= 1_048_576); + assert!(config.max_socks_connections >= 65_536); + assert!(config.max_streams_per_tunnel >= 65_536); + assert!(config.handshake_timeout <= std::time::Duration::from_secs(10)); + } + // ========================================================================= // Auth / mTLS tests // =========================================================================