Statistics on ME + Dynamic backpressure + KDF with SOCKS

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/transport/middle_proxy/handshake.rs

@@ -14,6 +14,7 @@ use tokio::net::{TcpStream, TcpSocket};
 use tokio::time::timeout;
 use tracing::{debug, info, warn};
 
+use crate::config::MeSocksKdfPolicy;
 use crate::crypto::{SecureRandom, build_middleproxy_prekey, derive_middleproxy_keys, sha256};
 use crate::error::{ProxyError, Result};
 use crate::network::IpFamily;
@@ -117,6 +118,13 @@ impl MePool {
         Some(bound)
     }
 
+    fn is_socks_route(upstream_egress: Option<UpstreamEgressInfo>) -> bool {
+        matches!(
+            upstream_egress.map(|info| info.route_kind),
+            Some(UpstreamRouteKind::Socks4 | UpstreamRouteKind::Socks5)
+        )
+    }
+
     /// TCP connect with timeout + return RTT in milliseconds.
     pub(crate) async fn connect_tcp(
         &self,
@@ -125,14 +133,7 @@ impl MePool {
         let start = Instant::now();
         let (stream, upstream_egress) = if let Some(upstream) = &self.upstream {
             let dc_idx = self.resolve_dc_idx_for_endpoint(addr).await;
-            let (stream, egress) = timeout(
-                Duration::from_secs(ME_CONNECT_TIMEOUT_SECS),
-                upstream.connect_with_details(addr, dc_idx, None),
-            )
-            .await
-            .map_err(|_| ProxyError::ConnectionTimeout {
-                addr: addr.to_string(),
-            })??;
+            let (stream, egress) = upstream.connect_with_details(addr, dc_idx, None).await?;
             (stream, Some(egress))
         } else {
             let connect_fut = async {
@@ -226,9 +227,29 @@ impl MePool {
         } else {
             IpFamily::V6
         };
+        let is_socks_route = Self::is_socks_route(upstream_egress);
         let socks_bound_addr = Self::select_socks_bound_addr(family, upstream_egress);
         let reflected = if let Some(bound) = socks_bound_addr {
             Some(bound)
+        } else if is_socks_route {
+            match self.socks_kdf_policy() {
+                MeSocksKdfPolicy::Strict => {
+                    self.stats.increment_me_socks_kdf_strict_reject();
+                    return Err(ProxyError::InvalidHandshake(
+                        "SOCKS route returned no valid BND.ADDR for ME KDF (strict policy)"
+                            .to_string(),
+                    ));
+                }
+                MeSocksKdfPolicy::Compat => {
+                    self.stats.increment_me_socks_kdf_compat_fallback();
+                    if self.nat_probe {
+                        let bind_ip = Self::direct_bind_ip_for_stun(family, upstream_egress);
+                        self.maybe_reflect_public_addr(family, bind_ip).await
+                    } else {
+                        None
+                    }
+                }
+            }
         } else if self.nat_probe {
             let bind_ip = Self::direct_bind_ip_for_stun(family, upstream_egress);
             self.maybe_reflect_public_addr(family, bind_ip).await

src/transport/middle_proxy/ping.rs

@@ -7,6 +7,7 @@ use tokio::net::UdpSocket;
 use crate::config::{UpstreamConfig, UpstreamType};
 use crate::crypto::SecureRandom;
 use crate::error::ProxyError;
+use crate::transport::{UpstreamEgressInfo, UpstreamRouteKind};
 
 use super::MePool;
 
@@ -20,6 +21,7 @@ pub enum MePingFamily {
 pub struct MePingSample {
     pub dc: i32,
     pub addr: SocketAddr,
+    pub route: Option<String>,
     pub connect_ms: Option<f64>,
     pub handshake_ms: Option<f64>,
     pub error: Option<String>,
@@ -84,6 +86,34 @@ fn pick_target_for_family(reports: &[MePingReport], family: MePingFamily) -> Opt
     })
 }
 
+fn route_from_egress(egress: Option<UpstreamEgressInfo>) -> Option<String> {
+    let info = egress?;
+    match info.route_kind {
+        UpstreamRouteKind::Direct => {
+            let src_ip = info
+                .direct_bind_ip
+                .or_else(|| info.local_addr.map(|addr| addr.ip()));
+            let ip = src_ip?;
+            let mut parts = Vec::new();
+            if let Some(dev) = detect_interface_for_ip(ip) {
+                parts.push(format!("dev={dev}"));
+            }
+            parts.push(format!("src={ip}"));
+            Some(format!("direct {}", parts.join(" ")))
+        }
+        UpstreamRouteKind::Socks4 => Some(
+            info.socks_bound_addr
+                .map(|addr| format!("socks4 bnd={addr}"))
+                .unwrap_or_else(|| "socks4".to_string()),
+        ),
+        UpstreamRouteKind::Socks5 => Some(
+            info.socks_bound_addr
+                .map(|addr| format!("socks5 bnd={addr}"))
+                .unwrap_or_else(|| "socks5".to_string()),
+        ),
+    }
+}
+
 #[cfg(unix)]
 fn detect_interface_for_ip(ip: IpAddr) -> Option<String> {
     use nix::ifaddrs::getifaddrs;
@@ -160,6 +190,15 @@ pub async fn format_me_route(
     v4_ok: bool,
     v6_ok: bool,
 ) -> String {
+    if let Some(route) = reports
+        .iter()
+        .flat_map(|report| report.samples.iter())
+        .find(|sample| sample.error.is_none() && sample.handshake_ms.is_some())
+        .and_then(|sample| sample.route.clone())
+    {
+        return route;
+    }
+
     let enabled_upstreams: Vec<_> = upstreams.iter().filter(|u| u.enabled).collect();
     if enabled_upstreams.is_empty() {
         return detect_direct_route_details(reports, prefer_ipv6, v4_ok, v6_ok)
@@ -222,6 +261,7 @@ mod tests {
         let s = sample(MePingSample {
             dc: 4,
             addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), 8888),
+            route: Some("direct src=1.2.3.4".to_string()),
             connect_ms: Some(12.3),
             handshake_ms: Some(34.7),
             error: None,
@@ -238,6 +278,7 @@ mod tests {
         let s = sample(MePingSample {
             dc: -5,
             addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(5, 6, 7, 8)), 80),
+            route: Some("socks5".to_string()),
             connect_ms: Some(10.0),
             handshake_ms: None,
             error: Some("handshake timeout".to_string()),
@@ -278,10 +319,12 @@ pub async fn run_me_ping(pool: &Arc<MePool>, rng: &SecureRandom) -> Vec<MePingRe
             let mut connect_ms = None;
             let mut handshake_ms = None;
             let mut error = None;
+            let mut route = None;
 
             match pool.connect_tcp(addr).await {
                 Ok((stream, conn_rtt, upstream_egress)) => {
                     connect_ms = Some(conn_rtt);
+                    route = route_from_egress(upstream_egress);
                     match pool.handshake_only(stream, addr, upstream_egress, rng).await {
                         Ok(hs) => {
                             handshake_ms = Some(hs.handshake_ms);
@@ -302,6 +345,7 @@ pub async fn run_me_ping(pool: &Arc<MePool>, rng: &SecureRandom) -> Vec<MePingRe
             samples.push(MePingSample {
                 dc,
                 addr,
+                route,
                 connect_ms,
                 handshake_ms,
                 error,

src/transport/middle_proxy/pool.rs

@@ -1,12 +1,13 @@
 use std::collections::{HashMap, HashSet};
 use std::net::{IpAddr, Ipv6Addr, SocketAddr};
 use std::sync::Arc;
-use std::sync::atomic::{AtomicBool, AtomicI32, AtomicU32, AtomicU64, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicI32, AtomicU8, AtomicU32, AtomicU64, AtomicUsize, Ordering};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 
 use tokio::sync::{Mutex, Notify, RwLock, mpsc};
 use tokio_util::sync::CancellationToken;
 
+use crate::config::MeSocksKdfPolicy;
 use crate::crypto::SecureRandom;
 use crate::network::IpFamily;
 use crate::network::probe::NetworkDecision;
@@ -82,6 +83,7 @@ pub struct MePool {
     pub(super) me_hardswap_warmup_delay_max_ms: AtomicU64,
     pub(super) me_hardswap_warmup_extra_passes: AtomicU32,
     pub(super) me_hardswap_warmup_pass_backoff_base_ms: AtomicU64,
+    pub(super) me_socks_kdf_policy: AtomicU8,
     pool_size: usize,
 }
 
@@ -145,9 +147,19 @@ impl MePool {
         me_hardswap_warmup_delay_max_ms: u64,
         me_hardswap_warmup_extra_passes: u8,
         me_hardswap_warmup_pass_backoff_base_ms: u64,
+        me_socks_kdf_policy: MeSocksKdfPolicy,
+        me_route_backpressure_base_timeout_ms: u64,
+        me_route_backpressure_high_timeout_ms: u64,
+        me_route_backpressure_high_watermark_pct: u8,
     ) -> Arc<Self> {
+        let registry = Arc::new(ConnRegistry::new());
+        registry.update_route_backpressure_policy(
+            me_route_backpressure_base_timeout_ms,
+            me_route_backpressure_high_timeout_ms,
+            me_route_backpressure_high_watermark_pct,
+        );
         Arc::new(Self {
-            registry: Arc::new(ConnRegistry::new()),
+            registry,
             writers: Arc::new(RwLock::new(Vec::new())),
             rr: AtomicU64::new(0),
             decision,
@@ -204,6 +216,7 @@ impl MePool {
             me_hardswap_warmup_pass_backoff_base_ms: AtomicU64::new(
                 me_hardswap_warmup_pass_backoff_base_ms,
             ),
+            me_socks_kdf_policy: AtomicU8::new(me_socks_kdf_policy.as_u8()),
         })
     }
 
@@ -260,6 +273,26 @@ impl MePool {
         &self.registry
     }
 
+    pub fn update_runtime_transport_policy(
+        &self,
+        socks_kdf_policy: MeSocksKdfPolicy,
+        route_backpressure_base_timeout_ms: u64,
+        route_backpressure_high_timeout_ms: u64,
+        route_backpressure_high_watermark_pct: u8,
+    ) {
+        self.me_socks_kdf_policy
+            .store(socks_kdf_policy.as_u8(), Ordering::Relaxed);
+        self.registry.update_route_backpressure_policy(
+            route_backpressure_base_timeout_ms,
+            route_backpressure_high_timeout_ms,
+            route_backpressure_high_watermark_pct,
+        );
+    }
+
+    pub(super) fn socks_kdf_policy(&self) -> MeSocksKdfPolicy {
+        MeSocksKdfPolicy::from_u8(self.me_socks_kdf_policy.load(Ordering::Relaxed))
+    }
+
     pub(super) fn writers_arc(&self) -> Arc<RwLock<Vec<MeWriter>>> {
         self.writers.clone()
     }

src/transport/middle_proxy/reader.rs

@@ -124,7 +124,14 @@ pub(crate) async fn reader_loop(
                     match routed {
                         RouteResult::NoConn => stats.increment_me_route_drop_no_conn(),
                         RouteResult::ChannelClosed => stats.increment_me_route_drop_channel_closed(),
-                        RouteResult::QueueFull => stats.increment_me_route_drop_queue_full(),
+                        RouteResult::QueueFullBase => {
+                            stats.increment_me_route_drop_queue_full();
+                            stats.increment_me_route_drop_queue_full_base();
+                        }
+                        RouteResult::QueueFullHigh => {
+                            stats.increment_me_route_drop_queue_full();
+                            stats.increment_me_route_drop_queue_full_high();
+                        }
                         RouteResult::Routed => {}
                     }
                     reg.unregister(cid).await;
@@ -140,7 +147,14 @@ pub(crate) async fn reader_loop(
                     match routed {
                         RouteResult::NoConn => stats.increment_me_route_drop_no_conn(),
                         RouteResult::ChannelClosed => stats.increment_me_route_drop_channel_closed(),
-                        RouteResult::QueueFull => stats.increment_me_route_drop_queue_full(),
+                        RouteResult::QueueFullBase => {
+                            stats.increment_me_route_drop_queue_full();
+                            stats.increment_me_route_drop_queue_full_base();
+                        }
+                        RouteResult::QueueFullHigh => {
+                            stats.increment_me_route_drop_queue_full();
+                            stats.increment_me_route_drop_queue_full_high();
+                        }
                         RouteResult::Routed => {}
                     }
                     reg.unregister(cid).await;

src/transport/middle_proxy/registry.rs

@@ -1,6 +1,6 @@
 use std::collections::{HashMap, HashSet};
 use std::net::SocketAddr;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
 use std::time::Duration;
 
 use tokio::sync::{mpsc, RwLock};
@@ -10,14 +10,17 @@ use super::codec::WriterCommand;
 use super::MeResponse;
 
 const ROUTE_CHANNEL_CAPACITY: usize = 4096;
-const ROUTE_BACKPRESSURE_TIMEOUT: Duration = Duration::from_millis(25);
+const ROUTE_BACKPRESSURE_BASE_TIMEOUT_MS: u64 = 25;
+const ROUTE_BACKPRESSURE_HIGH_TIMEOUT_MS: u64 = 120;
+const ROUTE_BACKPRESSURE_HIGH_WATERMARK_PCT: u8 = 80;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum RouteResult {
     Routed,
     NoConn,
     ChannelClosed,
-    QueueFull,
+    QueueFullBase,
+    QueueFullHigh,
 }
 
 #[derive(Clone)]
@@ -65,6 +68,9 @@ impl RegistryInner {
 pub struct ConnRegistry {
     inner: RwLock<RegistryInner>,
     next_id: AtomicU64,
+    route_backpressure_base_timeout_ms: AtomicU64,
+    route_backpressure_high_timeout_ms: AtomicU64,
+    route_backpressure_high_watermark_pct: AtomicU8,
 }
 
 impl ConnRegistry {
@@ -73,9 +79,35 @@ impl ConnRegistry {
         Self {
             inner: RwLock::new(RegistryInner::new()),
             next_id: AtomicU64::new(start),
+            route_backpressure_base_timeout_ms: AtomicU64::new(
+                ROUTE_BACKPRESSURE_BASE_TIMEOUT_MS,
+            ),
+            route_backpressure_high_timeout_ms: AtomicU64::new(
+                ROUTE_BACKPRESSURE_HIGH_TIMEOUT_MS,
+            ),
+            route_backpressure_high_watermark_pct: AtomicU8::new(
+                ROUTE_BACKPRESSURE_HIGH_WATERMARK_PCT,
+            ),
         }
     }
 
+    pub fn update_route_backpressure_policy(
+        &self,
+        base_timeout_ms: u64,
+        high_timeout_ms: u64,
+        high_watermark_pct: u8,
+    ) {
+        let base = base_timeout_ms.max(1);
+        let high = high_timeout_ms.max(base);
+        let watermark = high_watermark_pct.clamp(1, 100);
+        self.route_backpressure_base_timeout_ms
+            .store(base, Ordering::Relaxed);
+        self.route_backpressure_high_timeout_ms
+            .store(high, Ordering::Relaxed);
+        self.route_backpressure_high_watermark_pct
+            .store(watermark, Ordering::Relaxed);
+    }
+
     pub async fn register(&self) -> (u64, mpsc::Receiver<MeResponse>) {
         let id = self.next_id.fetch_add(1, Ordering::Relaxed);
         let (tx, rx) = mpsc::channel(ROUTE_CHANNEL_CAPACITY);
@@ -112,10 +144,40 @@ impl ConnRegistry {
             Err(TrySendError::Closed(_)) => RouteResult::ChannelClosed,
             Err(TrySendError::Full(resp)) => {
                 // Absorb short bursts without dropping/closing the session immediately.
-                match tokio::time::timeout(ROUTE_BACKPRESSURE_TIMEOUT, tx.send(resp)).await {
+                let base_timeout_ms =
+                    self.route_backpressure_base_timeout_ms.load(Ordering::Relaxed).max(1);
+                let high_timeout_ms = self
+                    .route_backpressure_high_timeout_ms
+                    .load(Ordering::Relaxed)
+                    .max(base_timeout_ms);
+                let high_watermark_pct = self
+                    .route_backpressure_high_watermark_pct
+                    .load(Ordering::Relaxed)
+                    .clamp(1, 100);
+                let used = ROUTE_CHANNEL_CAPACITY.saturating_sub(tx.capacity());
+                let used_pct = if ROUTE_CHANNEL_CAPACITY == 0 {
+                    100
+                } else {
+                    (used.saturating_mul(100) / ROUTE_CHANNEL_CAPACITY) as u8
+                };
+                let high_profile = used_pct >= high_watermark_pct;
+                let timeout_ms = if high_profile {
+                    high_timeout_ms
+                } else {
+                    base_timeout_ms
+                };
+                let timeout_dur = Duration::from_millis(timeout_ms);
+
+                match tokio::time::timeout(timeout_dur, tx.send(resp)).await {
                     Ok(Ok(())) => RouteResult::Routed,
                     Ok(Err(_)) => RouteResult::ChannelClosed,
-                    Err(_) => RouteResult::QueueFull,
+                    Err(_) => {
+                        if high_profile {
+                            RouteResult::QueueFullHigh
+                        } else {
+                            RouteResult::QueueFullBase
+                        }
+                    }
                 }
             }
         }