Merge pull request #332 from telemt/bump

Update Cargo.toml

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "telemt"
-version = "3.3.1"
+version = "3.3.3"
 edition = "2024"
 
 [dependencies]

src/config/defaults.rs

@@ -15,6 +15,7 @@ const DEFAULT_ME_ADAPTIVE_FLOOR_RECOVER_GRACE_SECS: u64 = 180;
 const DEFAULT_USER_MAX_UNIQUE_IPS_WINDOW_SECS: u64 = 30;
 const DEFAULT_UPSTREAM_CONNECT_RETRY_ATTEMPTS: u32 = 2;
 const DEFAULT_UPSTREAM_UNHEALTHY_FAIL_THRESHOLD: u32 = 5;
+const DEFAULT_UPSTREAM_CONNECT_BUDGET_MS: u64 = 3000;
 const DEFAULT_LISTEN_ADDR_IPV6: &str = "::";
 const DEFAULT_ACCESS_USER: &str = "default";
 const DEFAULT_ACCESS_SECRET: &str = "00000000000000000000000000000000";
@@ -113,6 +114,10 @@ pub(crate) fn default_api_minimal_runtime_cache_ttl_ms() -> u64 {
     1000
 }
 
+pub(crate) fn default_proxy_protocol_header_timeout_ms() -> u64 {
+    500
+}
+
 pub(crate) fn default_prefer_4() -> u8 {
     4
 }
@@ -253,6 +258,10 @@ pub(crate) fn default_upstream_unhealthy_fail_threshold() -> u32 {
     DEFAULT_UPSTREAM_UNHEALTHY_FAIL_THRESHOLD
 }
 
+pub(crate) fn default_upstream_connect_budget_ms() -> u64 {
+    DEFAULT_UPSTREAM_CONNECT_BUDGET_MS
+}
+
 pub(crate) fn default_upstream_connect_failfast_hard_errors() -> bool {
     false
 }

src/config/load.rs

@@ -265,6 +265,12 @@ impl ProxyConfig {
             ));
         }
 
+        if config.general.upstream_connect_budget_ms == 0 {
+            return Err(ProxyError::Config(
+                "general.upstream_connect_budget_ms must be > 0".to_string(),
+            ));
+        }
+
         if config.general.upstream_unhealthy_fail_threshold == 0 {
             return Err(ProxyError::Config(
                 "general.upstream_unhealthy_fail_threshold must be > 0".to_string(),
@@ -462,6 +468,12 @@ impl ProxyConfig {
             ));
         }
 
+        if config.server.proxy_protocol_header_timeout_ms == 0 {
+            return Err(ProxyError::Config(
+                "server.proxy_protocol_header_timeout_ms must be > 0".to_string(),
+            ));
+        }
+
         if config.general.effective_me_pool_force_close_secs() > 0
             && config.general.effective_me_pool_force_close_secs()
                 < config.general.me_pool_drain_ttl_secs
@@ -548,6 +560,12 @@ impl ProxyConfig {
             config.general.middle_proxy_nat_probe = true;
             warn!("Auto-enabled middle_proxy_nat_probe for middle proxy mode");
         }
+        if config.general.use_middle_proxy && !config.general.me_secret_atomic_snapshot {
+            config.general.me_secret_atomic_snapshot = true;
+            warn!(
+                "Auto-enabled me_secret_atomic_snapshot for middle proxy mode to keep KDF key_selector/secret coherent"
+            );
+        }
 
         validate_network_cfg(&mut config.network)?;
         crate::network::dns_overrides::validate_entries(&config.network.dns_overrides)?;

src/config/types.rs

@@ -532,6 +532,10 @@ pub struct GeneralConfig {
     #[serde(default = "default_upstream_connect_retry_backoff_ms")]
     pub upstream_connect_retry_backoff_ms: u64,
 
+    /// Total wall-clock budget in milliseconds for one upstream connect request across retries.
+    #[serde(default = "default_upstream_connect_budget_ms")]
+    pub upstream_connect_budget_ms: u64,
+
     /// Consecutive failed requests before upstream is marked unhealthy.
     #[serde(default = "default_upstream_unhealthy_fail_threshold")]
     pub upstream_unhealthy_fail_threshold: u32,
@@ -774,6 +778,7 @@ impl Default for GeneralConfig {
             me_adaptive_floor_recover_grace_secs: default_me_adaptive_floor_recover_grace_secs(),
             upstream_connect_retry_attempts: default_upstream_connect_retry_attempts(),
             upstream_connect_retry_backoff_ms: default_upstream_connect_retry_backoff_ms(),
+            upstream_connect_budget_ms: default_upstream_connect_budget_ms(),
             upstream_unhealthy_fail_threshold: default_upstream_unhealthy_fail_threshold(),
             upstream_connect_failfast_hard_errors: default_upstream_connect_failfast_hard_errors(),
             stun_iface_mismatch_ignore: false,
@@ -962,6 +967,10 @@ pub struct ServerConfig {
     #[serde(default)]
     pub proxy_protocol: bool,
 
+    /// Timeout in milliseconds for reading and parsing PROXY protocol headers.
+    #[serde(default = "default_proxy_protocol_header_timeout_ms")]
+    pub proxy_protocol_header_timeout_ms: u64,
+
     #[serde(default)]
     pub metrics_port: Option<u16>,
 
@@ -985,6 +994,7 @@ impl Default for ServerConfig {
             listen_unix_sock_perm: None,
             listen_tcp: None,
             proxy_protocol: false,
+            proxy_protocol_header_timeout_ms: default_proxy_protocol_header_timeout_ms(),
             metrics_port: None,
             metrics_whitelist: default_metrics_whitelist(),
             api: ApiConfig::default(),

src/crypto/random.rs

@@ -21,6 +21,7 @@ struct SecureRandomInner {
     rng: StdRng,
     cipher: AesCtr,
     buffer: Vec<u8>,
+    buffer_start: usize,
 }
 
 impl Drop for SecureRandomInner {
@@ -48,6 +49,7 @@ impl SecureRandom {
                 rng,
                 cipher,
                 buffer: Vec::with_capacity(1024),
+                buffer_start: 0,
             }),
         }
     }
@@ -59,16 +61,29 @@ impl SecureRandom {
 
         let mut written = 0usize;
         while written < out.len() {
+            if inner.buffer_start >= inner.buffer.len() {
+                inner.buffer.clear();
+                inner.buffer_start = 0;
+            }
+
             if inner.buffer.is_empty() {
                 let mut chunk = vec![0u8; CHUNK_SIZE];
                 inner.rng.fill_bytes(&mut chunk);
                 inner.cipher.apply(&mut chunk);
                 inner.buffer.extend_from_slice(&chunk);
+                inner.buffer_start = 0;
             }
 
-            let take = (out.len() - written).min(inner.buffer.len());
-            out[written..written + take].copy_from_slice(&inner.buffer[..take]);
-            inner.buffer.drain(..take);
+            let available = inner.buffer.len().saturating_sub(inner.buffer_start);
+            let take = (out.len() - written).min(available);
+            let start = inner.buffer_start;
+            let end = start + take;
+            out[written..written + take].copy_from_slice(&inner.buffer[start..end]);
+            inner.buffer_start = end;
+            if inner.buffer_start >= inner.buffer.len() {
+                inner.buffer.clear();
+                inner.buffer_start = 0;
+            }
             written += take;
         }
     }

src/main.rs

@@ -464,6 +464,7 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
         config.upstreams.clone(),
         config.general.upstream_connect_retry_attempts,
         config.general.upstream_connect_retry_backoff_ms,
+        config.general.upstream_connect_budget_ms,
         config.general.upstream_unhealthy_fail_threshold,
         config.general.upstream_connect_failfast_hard_errors,
         stats.clone(),
@@ -1339,7 +1340,7 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
     let (admission_tx, admission_rx) = watch::channel(true);
     if config.general.use_middle_proxy {
         if let Some(pool) = me_pool.as_ref() {
-            let initial_open = pool.admission_ready_full_floor().await;
+            let initial_open = pool.admission_ready_conditional_cast().await;
             admission_tx.send_replace(initial_open);
             if initial_open {
                 info!("Conditional-admission gate: open (ME pool ready)");
@@ -1354,7 +1355,7 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
                 let mut open_streak = if initial_open { 1u32 } else { 0u32 };
                 let mut close_streak = if initial_open { 0u32 } else { 1u32 };
                 loop {
-                    let ready = pool_for_gate.admission_ready_full_floor().await;
+                    let ready = pool_for_gate.admission_ready_conditional_cast().await;
                     if ready {
                         open_streak = open_streak.saturating_add(1);
                         close_streak = 0;
@@ -1374,7 +1375,7 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
                             admission_tx_gate.send_replace(false);
                             warn!(
                                 close_streak,
-                                "Conditional-admission gate closed (ME pool below required floor)"
+                                "Conditional-admission gate closed (ME pool has uncovered DC groups)"
                             );
                         }
                     }

src/proxy/client.rs

@@ -97,8 +97,11 @@ where
         .unwrap_or_else(|_| "0.0.0.0:443".parse().unwrap());
 
     if proxy_protocol_enabled {
-        match parse_proxy_protocol(&mut stream, peer).await {
-            Ok(info) => {
+        let proxy_header_timeout = Duration::from_millis(
+            config.server.proxy_protocol_header_timeout_ms.max(1),
+        );
+        match timeout(proxy_header_timeout, parse_proxy_protocol(&mut stream, peer)).await {
+            Ok(Ok(info)) => {
                 debug!(
                     peer = %peer,
                     client = %info.src_addr,
@@ -110,12 +113,18 @@ where
                     local_addr = dst;
                 }
             }
-            Err(e) => {
+            Ok(Err(e)) => {
                 stats.increment_connects_bad();
                 warn!(peer = %peer, error = %e, "Invalid PROXY protocol header");
                 record_beobachten_class(&beobachten, &config, peer.ip(), "other");
                 return Err(e);
             }
+            Err(_) => {
+                stats.increment_connects_bad();
+                warn!(peer = %peer, timeout_ms = proxy_header_timeout.as_millis(), "PROXY protocol header timeout");
+                record_beobachten_class(&beobachten, &config, peer.ip(), "other");
+                return Err(ProxyError::InvalidProxyProtocol);
+            }
         }
     }
 
@@ -161,7 +170,7 @@ where
 
             let (read_half, write_half) = tokio::io::split(stream);
 
-            let (mut tls_reader, tls_writer, _tls_user) = match handle_tls_handshake(
+            let (mut tls_reader, tls_writer, tls_user) = match handle_tls_handshake(
                 &handshake, read_half, write_half, real_peer,
                 &config, &replay_checker, &rng, tls_cache.clone(),
             ).await {
@@ -190,7 +199,7 @@ where
 
             let (crypto_reader, crypto_writer, success) = match handle_mtproto_handshake(
                 &mtproto_handshake, tls_reader, tls_writer, real_peer,
-                &config, &replay_checker, true,
+                &config, &replay_checker, true, Some(tls_user.as_str()),
             ).await {
                 HandshakeResult::Success(result) => result,
                 HandshakeResult::BadClient { reader: _, writer: _ } => {
@@ -234,7 +243,7 @@ where
 
             let (crypto_reader, crypto_writer, success) = match handle_mtproto_handshake(
                 &handshake, read_half, write_half, real_peer,
-                &config, &replay_checker, false,
+                &config, &replay_checker, false, None,
             ).await {
                 HandshakeResult::Success(result) => result,
                 HandshakeResult::BadClient { reader, writer } => {
@@ -415,8 +424,16 @@ impl RunningClientHandler {
         let mut local_addr = self.stream.local_addr().map_err(ProxyError::Io)?;
 
         if self.proxy_protocol_enabled {
-            match parse_proxy_protocol(&mut self.stream, self.peer).await {
-                Ok(info) => {
+            let proxy_header_timeout = Duration::from_millis(
+                self.config.server.proxy_protocol_header_timeout_ms.max(1),
+            );
+            match timeout(
+                proxy_header_timeout,
+                parse_proxy_protocol(&mut self.stream, self.peer),
+            )
+            .await
+            {
+                Ok(Ok(info)) => {
                     debug!(
                         peer = %self.peer,
                         client = %info.src_addr,
@@ -428,7 +445,7 @@ impl RunningClientHandler {
                         local_addr = dst;
                     }
                 }
-                Err(e) => {
+                Ok(Err(e)) => {
                     self.stats.increment_connects_bad();
                     warn!(peer = %self.peer, error = %e, "Invalid PROXY protocol header");
                     record_beobachten_class(
@@ -439,6 +456,21 @@ impl RunningClientHandler {
                     );
                     return Err(e);
                 }
+                Err(_) => {
+                    self.stats.increment_connects_bad();
+                    warn!(
+                        peer = %self.peer,
+                        timeout_ms = proxy_header_timeout.as_millis(),
+                        "PROXY protocol header timeout"
+                    );
+                    record_beobachten_class(
+                        &self.beobachten,
+                        &self.config,
+                        self.peer.ip(),
+                        "other",
+                    );
+                    return Err(ProxyError::InvalidProxyProtocol);
+                }
             }
         }
 
@@ -494,7 +526,7 @@ impl RunningClientHandler {
 
         let (read_half, write_half) = self.stream.into_split();
 
-        let (mut tls_reader, tls_writer, _tls_user) = match handle_tls_handshake(
+        let (mut tls_reader, tls_writer, tls_user) = match handle_tls_handshake(
             &handshake,
             read_half,
             write_half,
@@ -538,6 +570,7 @@ impl RunningClientHandler {
             &config,
             &replay_checker,
             true,
+            Some(tls_user.as_str()),
         )
         .await
         {
@@ -611,6 +644,7 @@ impl RunningClientHandler {
             &config,
             &replay_checker,
             false,
+            None,
         )
         .await
         {

src/proxy/direct_relay.rs

@@ -34,7 +34,7 @@ where
     let user = &success.user;
     let dc_addr = get_dc_addr_static(success.dc_idx, &config)?;
 
-    info!(
+    debug!(
         user = %user,
         peer = %success.peer,
         dc = success.dc_idx,

src/proxy/handshake.rs

@@ -6,7 +6,7 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
-use tracing::{debug, warn, trace, info};
+use tracing::{debug, warn, trace};
 use zeroize::Zeroize;
 
 use crate::crypto::{sha256, AesCtr, SecureRandom};
@@ -19,6 +19,31 @@ use crate::stats::ReplayChecker;
 use crate::config::ProxyConfig;
 use crate::tls_front::{TlsFrontCache, emulator};
 
+fn decode_user_secrets(
+    config: &ProxyConfig,
+    preferred_user: Option<&str>,
+) -> Vec<(String, Vec<u8>)> {
+    let mut secrets = Vec::with_capacity(config.access.users.len());
+
+    if let Some(preferred) = preferred_user
+        && let Some(secret_hex) = config.access.users.get(preferred)
+        && let Ok(bytes) = hex::decode(secret_hex)
+    {
+        secrets.push((preferred.to_string(), bytes));
+    }
+
+    for (name, secret_hex) in &config.access.users {
+        if preferred_user.is_some_and(|preferred| preferred == name.as_str()) {
+            continue;
+        }
+        if let Ok(bytes) = hex::decode(secret_hex) {
+            secrets.push((name.clone(), bytes));
+        }
+    }
+
+    secrets
+}
+
 /// Result of successful handshake
 ///
 /// Key material (`dec_key`, `dec_iv`, `enc_key`, `enc_iv`) is
@@ -82,11 +107,7 @@ where
         return HandshakeResult::BadClient { reader, writer };
     }
 
-    let secrets: Vec<(String, Vec<u8>)> = config.access.users.iter()
-        .filter_map(|(name, hex)| {
-            hex::decode(hex).ok().map(|bytes| (name.clone(), bytes))
-        })
-        .collect();
+    let secrets = decode_user_secrets(config, None);
 
     let validation = match tls::validate_tls_handshake(
         handshake,
@@ -201,7 +222,7 @@ where
         return HandshakeResult::Error(ProxyError::Io(e));
     }
 
-    info!(
+    debug!(
         peer = %peer,
         user = %validation.user,
         "TLS handshake successful"
@@ -223,6 +244,7 @@ pub async fn handle_mtproto_handshake<R, W>(
     config: &ProxyConfig,
     replay_checker: &ReplayChecker,
     is_tls: bool,
+    preferred_user: Option<&str>,
 ) -> HandshakeResult<(CryptoReader<R>, CryptoWriter<W>, HandshakeSuccess), R, W>
 where
     R: AsyncRead + Unpin + Send,
@@ -239,11 +261,9 @@ where
 
     let enc_prekey_iv: Vec<u8> = dec_prekey_iv.iter().rev().copied().collect();
 
-    for (user, secret_hex) in &config.access.users {
-        let secret = match hex::decode(secret_hex) {
-            Ok(s) => s,
-            Err(_) => continue,
-        };
+    let decoded_users = decode_user_secrets(config, preferred_user);
+
+    for (user, secret) in decoded_users {
 
         let dec_prekey = &dec_prekey_iv[..PREKEY_LEN];
         let dec_iv_bytes = &dec_prekey_iv[PREKEY_LEN..];
@@ -311,7 +331,7 @@ where
             is_tls,
         };
 
-        info!(
+        debug!(
             peer = %peer,
             user = %user,
             dc = dc_idx,

src/proxy/middle_relay.rs

@@ -8,7 +8,7 @@ use std::time::{Duration, Instant};
 
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use tokio::sync::{mpsc, oneshot};
-use tracing::{debug, info, trace, warn};
+use tracing::{debug, trace, warn};
 
 use crate::config::ProxyConfig;
 use crate::crypto::SecureRandom;
@@ -210,7 +210,7 @@ where
     let proto_tag = success.proto_tag;
     let pool_generation = me_pool.current_generation();
 
-    info!(
+    debug!(
         user = %user,
         peer = %peer,
         dc = success.dc_idx,

src/stats/mod.rs

@@ -846,16 +846,30 @@ impl Stats {
         if !self.telemetry_user_enabled() {
             return;
         }
-        self.user_stats.entry(user.to_string()).or_default()
-            .connects.fetch_add(1, Ordering::Relaxed);
+        if let Some(stats) = self.user_stats.get(user) {
+            stats.connects.fetch_add(1, Ordering::Relaxed);
+            return;
+        }
+        self.user_stats
+            .entry(user.to_string())
+            .or_default()
+            .connects
+            .fetch_add(1, Ordering::Relaxed);
     }
     
     pub fn increment_user_curr_connects(&self, user: &str) {
         if !self.telemetry_user_enabled() {
             return;
         }
-        self.user_stats.entry(user.to_string()).or_default()
-            .curr_connects.fetch_add(1, Ordering::Relaxed);
+        if let Some(stats) = self.user_stats.get(user) {
+            stats.curr_connects.fetch_add(1, Ordering::Relaxed);
+            return;
+        }
+        self.user_stats
+            .entry(user.to_string())
+            .or_default()
+            .curr_connects
+            .fetch_add(1, Ordering::Relaxed);
     }
     
     pub fn decrement_user_curr_connects(&self, user: &str) {
@@ -889,32 +903,60 @@ impl Stats {
         if !self.telemetry_user_enabled() {
             return;
         }
-        self.user_stats.entry(user.to_string()).or_default()
-            .octets_from_client.fetch_add(bytes, Ordering::Relaxed);
+        if let Some(stats) = self.user_stats.get(user) {
+            stats.octets_from_client.fetch_add(bytes, Ordering::Relaxed);
+            return;
+        }
+        self.user_stats
+            .entry(user.to_string())
+            .or_default()
+            .octets_from_client
+            .fetch_add(bytes, Ordering::Relaxed);
     }
     
     pub fn add_user_octets_to(&self, user: &str, bytes: u64) {
         if !self.telemetry_user_enabled() {
             return;
         }
-        self.user_stats.entry(user.to_string()).or_default()
-            .octets_to_client.fetch_add(bytes, Ordering::Relaxed);
+        if let Some(stats) = self.user_stats.get(user) {
+            stats.octets_to_client.fetch_add(bytes, Ordering::Relaxed);
+            return;
+        }
+        self.user_stats
+            .entry(user.to_string())
+            .or_default()
+            .octets_to_client
+            .fetch_add(bytes, Ordering::Relaxed);
     }
     
     pub fn increment_user_msgs_from(&self, user: &str) {
         if !self.telemetry_user_enabled() {
             return;
         }
-        self.user_stats.entry(user.to_string()).or_default()
-            .msgs_from_client.fetch_add(1, Ordering::Relaxed);
+        if let Some(stats) = self.user_stats.get(user) {
+            stats.msgs_from_client.fetch_add(1, Ordering::Relaxed);
+            return;
+        }
+        self.user_stats
+            .entry(user.to_string())
+            .or_default()
+            .msgs_from_client
+            .fetch_add(1, Ordering::Relaxed);
     }
     
     pub fn increment_user_msgs_to(&self, user: &str) {
         if !self.telemetry_user_enabled() {
             return;
         }
-        self.user_stats.entry(user.to_string()).or_default()
-            .msgs_to_client.fetch_add(1, Ordering::Relaxed);
+        if let Some(stats) = self.user_stats.get(user) {
+            stats.msgs_to_client.fetch_add(1, Ordering::Relaxed);
+            return;
+        }
+        self.user_stats
+            .entry(user.to_string())
+            .or_default()
+            .msgs_to_client
+            .fetch_add(1, Ordering::Relaxed);
     }
     
     pub fn get_user_total_octets(&self, user: &str) -> u64 {

src/transport/middle_proxy/handshake.rs

@@ -387,9 +387,11 @@ impl MePool {
             socks_bound_addr.map(|value| value.ip()),
             client_port_source,
         );
-        let mut kdf_fingerprint_guard = self.kdf_material_fingerprint.lock().await;
-        if let Some((prev_fingerprint, prev_client_port)) =
+        let previous_kdf_fingerprint = {
+            let kdf_fingerprint_guard = self.kdf_material_fingerprint.read().await;
             kdf_fingerprint_guard.get(&peer_addr_nat).copied()
+        };
+        if let Some((prev_fingerprint, prev_client_port)) = previous_kdf_fingerprint
         {
             if prev_fingerprint != kdf_fingerprint {
                 self.stats.increment_me_kdf_drift_total();
@@ -416,6 +418,9 @@ impl MePool {
                 );
             }
         }
+        // Keep fingerprint updates eventually consistent for diagnostics while avoiding
+        // serializing all concurrent handshakes on a single async mutex.
+        let mut kdf_fingerprint_guard = self.kdf_material_fingerprint.write().await;
         kdf_fingerprint_guard.insert(peer_addr_nat, (kdf_fingerprint, client_port_for_kdf));
         drop(kdf_fingerprint_guard);
 

src/transport/middle_proxy/pool.rs

@@ -119,6 +119,8 @@ pub struct MePool {
     pub(super) ping_tracker: Arc<Mutex<HashMap<i64, (std::time::Instant, u64)>>>,
     pub(super) rtt_stats: Arc<Mutex<HashMap<u64, (f64, f64)>>>,
     pub(super) nat_reflection_cache: Arc<Mutex<NatReflectionCache>>,
+    pub(super) nat_reflection_singleflight_v4: Arc<Mutex<()>>,
+    pub(super) nat_reflection_singleflight_v6: Arc<Mutex<()>>,
     pub(super) writer_available: Arc<Notify>,
     pub(super) refill_inflight: Arc<Mutex<HashSet<SocketAddr>>>,
     pub(super) refill_inflight_dc: Arc<Mutex<HashSet<RefillDcKey>>>,
@@ -132,7 +134,7 @@ pub struct MePool {
     pub(super) pending_hardswap_map_hash: AtomicU64,
     pub(super) hardswap: AtomicBool,
     pub(super) endpoint_quarantine: Arc<Mutex<HashMap<SocketAddr, Instant>>>,
-    pub(super) kdf_material_fingerprint: Arc<Mutex<HashMap<SocketAddr, (u64, u16)>>>,
+    pub(super) kdf_material_fingerprint: Arc<RwLock<HashMap<SocketAddr, (u64, u16)>>>,
     pub(super) me_pool_drain_ttl_secs: AtomicU64,
     pub(super) me_pool_force_close_secs: AtomicU64,
     pub(super) me_pool_min_fresh_ratio_permille: AtomicU32,
@@ -323,6 +325,8 @@ impl MePool {
             ping_tracker: Arc::new(Mutex::new(HashMap::new())),
             rtt_stats: Arc::new(Mutex::new(HashMap::new())),
             nat_reflection_cache: Arc::new(Mutex::new(NatReflectionCache::default())),
+            nat_reflection_singleflight_v4: Arc::new(Mutex::new(())),
+            nat_reflection_singleflight_v6: Arc::new(Mutex::new(())),
             writer_available: Arc::new(Notify::new()),
             refill_inflight: Arc::new(Mutex::new(HashSet::new())),
             refill_inflight_dc: Arc::new(Mutex::new(HashSet::new())),
@@ -335,7 +339,7 @@ impl MePool {
             pending_hardswap_map_hash: AtomicU64::new(0),
             hardswap: AtomicBool::new(hardswap),
             endpoint_quarantine: Arc::new(Mutex::new(HashMap::new())),
-            kdf_material_fingerprint: Arc::new(Mutex::new(HashMap::new())),
+            kdf_material_fingerprint: Arc::new(RwLock::new(HashMap::new())),
             me_pool_drain_ttl_secs: AtomicU64::new(me_pool_drain_ttl_secs),
             me_pool_force_close_secs: AtomicU64::new(me_pool_force_close_secs),
             me_pool_min_fresh_ratio_permille: AtomicU32::new(Self::ratio_to_permille(

src/transport/middle_proxy/pool_init.rs

@@ -14,10 +14,12 @@ use super::pool::MePool;
 impl MePool {
     pub async fn init(self: &Arc<Self>, pool_size: usize, rng: &Arc<SecureRandom>) -> Result<()> {
         let family_order = self.family_order();
+        let connect_concurrency = self.me_reconnect_max_concurrent_per_dc.max(1) as usize;
         let ks = self.key_selector().await;
         info!(
             me_servers = self.proxy_map_v4.read().await.len(),
             pool_size,
+            connect_concurrency,
             key_selector = format_args!("0x{ks:08x}"),
             secret_len = self.proxy_secret.read().await.secret.len(),
             "Initializing ME pool"
@@ -41,23 +43,39 @@ impl MePool {
                 })
                 .collect();
             dc_addrs.sort_unstable_by_key(|(dc, _)| *dc);
+            dc_addrs.sort_by_key(|(_, addrs)| (addrs.len() != 1, addrs.len()));
 
-            // Ensure at least one live writer per DC group; run missing DCs in parallel.
+            // Stage 1: build base coverage for conditional-cast.
+            // Single-endpoint DCs are prefilled first; multi-endpoint DCs require one live writer.
             let mut join = tokio::task::JoinSet::new();
             for (dc, addrs) in dc_addrs.iter().cloned() {
                 if addrs.is_empty() {
                     continue;
                 }
+                let target_writers = if addrs.len() == 1 {
+                    self.required_writers_for_dc_with_floor_mode(addrs.len(), false)
+                } else {
+                    1usize
+                };
                 let endpoints: HashSet<SocketAddr> = addrs
                     .iter()
                     .map(|(ip, port)| SocketAddr::new(*ip, *port))
                     .collect();
-                if self.active_writer_count_for_endpoints(&endpoints).await > 0 {
+                if self.active_writer_count_for_endpoints(&endpoints).await >= target_writers {
                     continue;
                 }
                 let pool = Arc::clone(self);
                 let rng_clone = Arc::clone(rng);
-                join.spawn(async move { pool.connect_primary_for_dc(dc, addrs, rng_clone).await });
+                join.spawn(async move {
+                    pool.connect_primary_for_dc(
+                        dc,
+                        addrs,
+                        target_writers,
+                        rng_clone,
+                        connect_concurrency,
+                    )
+                    .await
+                });
             }
             while join.join_next().await.is_some() {}
 
@@ -77,47 +95,35 @@ impl MePool {
                 )));
             }
 
-            // Warm reserve writers asynchronously so startup does not block after first working pool is ready.
+            // Stage 2: continue saturating multi-endpoint DC groups in background.
             let pool = Arc::clone(self);
             let rng_clone = Arc::clone(rng);
             let dc_addrs_bg = dc_addrs.clone();
             tokio::spawn(async move {
-                if pool.me_warmup_stagger_enabled {
-                    for (dc, addrs) in &dc_addrs_bg {
-                        for (ip, port) in addrs {
-                            if pool.connection_count() >= pool_size {
-                                break;
-                            }
-                            let addr = SocketAddr::new(*ip, *port);
-                            let jitter = rand::rng()
-                                .random_range(0..=pool.me_warmup_step_jitter.as_millis() as u64);
-                            let delay_ms = pool.me_warmup_step_delay.as_millis() as u64 + jitter;
-                            tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
-                            if let Err(e) = pool.connect_one(addr, rng_clone.as_ref()).await {
-                                debug!(%addr, dc = %dc, error = %e, "Extra ME connect failed (staggered)");
-                            }
-                        }
-                    }
-                } else {
-                    for (dc, addrs) in &dc_addrs_bg {
-                        for (ip, port) in addrs {
-                            if pool.connection_count() >= pool_size {
-                                break;
-                            }
-                            let addr = SocketAddr::new(*ip, *port);
-                            if let Err(e) = pool.connect_one(addr, rng_clone.as_ref()).await {
-                                debug!(%addr, dc = %dc, error = %e, "Extra ME connect failed");
-                            }
-                        }
-                        if pool.connection_count() >= pool_size {
-                            break;
-                        }
+                let mut join_bg = tokio::task::JoinSet::new();
+                for (dc, addrs) in dc_addrs_bg {
+                    if addrs.len() <= 1 {
+                        continue;
                     }
+                    let target_writers = pool.required_writers_for_dc_with_floor_mode(addrs.len(), false);
+                    let pool_clone = Arc::clone(&pool);
+                    let rng_clone_local = Arc::clone(&rng_clone);
+                    join_bg.spawn(async move {
+                        pool_clone
+                            .connect_primary_for_dc(
+                                dc,
+                                addrs,
+                                target_writers,
+                                rng_clone_local,
+                                connect_concurrency,
+                            )
+                            .await
+                    });
                 }
+                while join_bg.join_next().await.is_some() {}
                 debug!(
-                    target_pool_size = pool_size,
                     current_pool_size = pool.connection_count(),
-                    "Background ME reserve warmup finished"
+                    "Background ME saturation warmup finished"
                 );
             });
 
@@ -140,62 +146,85 @@ impl MePool {
         self: Arc<Self>,
         dc: i32,
         mut addrs: Vec<(IpAddr, u16)>,
+        target_writers: usize,
         rng: Arc<SecureRandom>,
+        connect_concurrency: usize,
     ) -> bool {
         if addrs.is_empty() {
             return false;
         }
+        let target_writers = target_writers.max(1);
         addrs.shuffle(&mut rand::rng());
-        if addrs.len() > 1 {
-            let concurrency = 2usize;
+        let endpoints: Vec<SocketAddr> = addrs
+            .iter()
+            .map(|(ip, port)| SocketAddr::new(*ip, *port))
+            .collect();
+        let endpoint_set: HashSet<SocketAddr> = endpoints.iter().copied().collect();
+
+        loop {
+            let alive = self.active_writer_count_for_endpoints(&endpoint_set).await;
+            if alive >= target_writers {
+                info!(
+                    dc = %dc,
+                    alive,
+                    target_writers,
+                    "ME connected"
+                );
+                return true;
+            }
+
+            let missing = target_writers.saturating_sub(alive).max(1);
+            let concurrency = connect_concurrency.max(1).min(missing);
             let mut join = tokio::task::JoinSet::new();
-            let mut next_idx = 0usize;
+            for _ in 0..concurrency {
+                let pool = Arc::clone(&self);
+                let rng_clone = Arc::clone(&rng);
+                let endpoints_clone = endpoints.clone();
+                join.spawn(async move {
+                    pool.connect_endpoints_round_robin(&endpoints_clone, rng_clone.as_ref())
+                        .await
+                });
+            }
 
-            while next_idx < addrs.len() || !join.is_empty() {
-                while next_idx < addrs.len() && join.len() < concurrency {
-                    let (ip, port) = addrs[next_idx];
-                    next_idx += 1;
-                    let addr = SocketAddr::new(ip, port);
-                    let pool = Arc::clone(&self);
-                    let rng_clone = Arc::clone(&rng);
-                    join.spawn(async move {
-                        (addr, pool.connect_one(addr, rng_clone.as_ref()).await)
-                    });
-                }
-
-                let Some(res) = join.join_next().await else {
-                    break;
-                };
+            let mut progress = false;
+            while let Some(res) = join.join_next().await {
                 match res {
-                    Ok((addr, Ok(()))) => {
-                        info!(%addr, dc = %dc, "ME connected");
-                        join.abort_all();
-                        while join.join_next().await.is_some() {}
-                        return true;
-                    }
-                    Ok((addr, Err(e))) => {
-                        warn!(%addr, dc = %dc, error = %e, "ME connect failed, trying next");
+                    Ok(true) => {
+                        progress = true;
                     }
+                    Ok(false) => {}
                     Err(e) => {
                         warn!(dc = %dc, error = %e, "ME connect task failed");
                     }
                 }
             }
-            warn!(dc = %dc, "All ME servers for DC failed at init");
-            return false;
-        }
 
-        for (ip, port) in addrs {
-            let addr = SocketAddr::new(ip, port);
-            match self.connect_one(addr, rng.as_ref()).await {
-                Ok(()) => {
-                    info!(%addr, dc = %dc, "ME connected");
-                    return true;
-                }
-                Err(e) => warn!(%addr, dc = %dc, error = %e, "ME connect failed, trying next"),
+            let alive_after = self.active_writer_count_for_endpoints(&endpoint_set).await;
+            if alive_after >= target_writers {
+                info!(
+                    dc = %dc,
+                    alive = alive_after,
+                    target_writers,
+                    "ME connected"
+                );
+                return true;
+            }
+            if !progress {
+                warn!(
+                    dc = %dc,
+                    alive = alive_after,
+                    target_writers,
+                    "All ME servers for DC failed at init"
+                );
+                return false;
+            }
+
+            if self.me_warmup_stagger_enabled {
+                let jitter = rand::rng()
+                    .random_range(0..=self.me_warmup_step_jitter.as_millis() as u64);
+                let delay_ms = self.me_warmup_step_delay.as_millis() as u64 + jitter;
+                tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
             }
         }
-        warn!(dc = %dc, "All ME servers for DC failed at init");
-        false
     }
 }

src/transport/middle_proxy/pool_nat.rs

@@ -248,6 +248,43 @@ impl MePool {
             }
         }
 
+        let _singleflight_guard = if use_shared_cache {
+            Some(match family {
+                IpFamily::V4 => self.nat_reflection_singleflight_v4.lock().await,
+                IpFamily::V6 => self.nat_reflection_singleflight_v6.lock().await,
+            })
+        } else {
+            None
+        };
+
+        if use_shared_cache
+            && let Some(until) = *self.stun_backoff_until.read().await
+            && Instant::now() < until
+        {
+            if let Ok(cache) = self.nat_reflection_cache.try_lock() {
+                let slot = match family {
+                    IpFamily::V4 => cache.v4,
+                    IpFamily::V6 => cache.v6,
+                };
+                return slot.map(|(_, addr)| addr);
+            }
+            return None;
+        }
+
+        if use_shared_cache
+            && let Ok(mut cache) = self.nat_reflection_cache.try_lock()
+        {
+            let slot = match family {
+                IpFamily::V4 => &mut cache.v4,
+                IpFamily::V6 => &mut cache.v6,
+            };
+            if let Some((ts, addr)) = slot
+                && ts.elapsed() < STUN_CACHE_TTL
+            {
+                return Some(*addr);
+            }
+        }
+
         let attempt = if use_shared_cache {
             self.nat_probe_attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
         } else {

src/transport/middle_proxy/pool_status.rs

@@ -100,6 +100,68 @@ pub(crate) struct MeApiRuntimeSnapshot {
 }
 
 impl MePool {
+    pub(crate) async fn admission_ready_conditional_cast(&self) -> bool {
+        let mut endpoints_by_dc = BTreeMap::<i16, BTreeSet<SocketAddr>>::new();
+        if self.decision.ipv4_me {
+            let map = self.proxy_map_v4.read().await.clone();
+            for (dc, addrs) in map {
+                let abs_dc = dc.abs();
+                if abs_dc == 0 {
+                    continue;
+                }
+                let Ok(dc_idx) = i16::try_from(abs_dc) else {
+                    continue;
+                };
+                let entry = endpoints_by_dc.entry(dc_idx).or_default();
+                for (ip, port) in addrs {
+                    entry.insert(SocketAddr::new(ip, port));
+                }
+            }
+        }
+        if self.decision.ipv6_me {
+            let map = self.proxy_map_v6.read().await.clone();
+            for (dc, addrs) in map {
+                let abs_dc = dc.abs();
+                if abs_dc == 0 {
+                    continue;
+                }
+                let Ok(dc_idx) = i16::try_from(abs_dc) else {
+                    continue;
+                };
+                let entry = endpoints_by_dc.entry(dc_idx).or_default();
+                for (ip, port) in addrs {
+                    entry.insert(SocketAddr::new(ip, port));
+                }
+            }
+        }
+
+        if endpoints_by_dc.is_empty() {
+            return false;
+        }
+
+        let writers = self.writers.read().await.clone();
+        let mut live_writers_by_endpoint = HashMap::<SocketAddr, usize>::new();
+        for writer in writers {
+            if writer.draining.load(Ordering::Relaxed) {
+                continue;
+            }
+            *live_writers_by_endpoint.entry(writer.addr).or_insert(0) += 1;
+        }
+
+        for endpoints in endpoints_by_dc.values() {
+            let alive: usize = endpoints
+                .iter()
+                .map(|endpoint| live_writers_by_endpoint.get(endpoint).copied().unwrap_or(0))
+                .sum();
+            if alive == 0 {
+                return false;
+            }
+        }
+
+        true
+    }
+
+    #[allow(dead_code)]
     pub(crate) async fn admission_ready_full_floor(&self) -> bool {
         let mut endpoints_by_dc = BTreeMap::<i16, BTreeSet<SocketAddr>>::new();
         if self.decision.ipv4_me {

src/transport/middle_proxy/reader.rs

@@ -124,7 +124,7 @@ pub(crate) async fn reader_loop(
                 let data = Bytes::copy_from_slice(&body[12..]);
                 trace!(cid, flags, len = data.len(), "RPC_PROXY_ANS");
 
-                let routed = reg.route(cid, MeResponse::Data { flags, data }).await;
+                let routed = reg.route_nowait(cid, MeResponse::Data { flags, data }).await;
                 if !matches!(routed, RouteResult::Routed) {
                     match routed {
                         RouteResult::NoConn => stats.increment_me_route_drop_no_conn(),
@@ -147,7 +147,7 @@ pub(crate) async fn reader_loop(
                 let cfm = u32::from_le_bytes(body[8..12].try_into().unwrap());
                 trace!(cid, cfm, "RPC_SIMPLE_ACK");
 
-                let routed = reg.route(cid, MeResponse::Ack(cfm)).await;
+                let routed = reg.route_nowait(cid, MeResponse::Ack(cfm)).await;
                 if !matches!(routed, RouteResult::Routed) {
                     match routed {
                         RouteResult::NoConn => stats.increment_me_route_drop_no_conn(),

src/transport/middle_proxy/registry.rs

@@ -208,6 +208,23 @@ impl ConnRegistry {
         }
     }
 
+    pub async fn route_nowait(&self, id: u64, resp: MeResponse) -> RouteResult {
+        let tx = {
+            let inner = self.inner.read().await;
+            inner.map.get(&id).cloned()
+        };
+
+        let Some(tx) = tx else {
+            return RouteResult::NoConn;
+        };
+
+        match tx.try_send(resp) {
+            Ok(()) => RouteResult::Routed,
+            Err(TrySendError::Closed(_)) => RouteResult::ChannelClosed,
+            Err(TrySendError::Full(_)) => RouteResult::QueueFullBase,
+        }
+    }
+
     pub async fn bind_writer(
         &self,
         conn_id: u64,

src/transport/upstream.rs

@@ -225,6 +225,7 @@ pub struct UpstreamManager {
     upstreams: Arc<RwLock<Vec<UpstreamState>>>,
     connect_retry_attempts: u32,
     connect_retry_backoff: Duration,
+    connect_budget: Duration,
     unhealthy_fail_threshold: u32,
     connect_failfast_hard_errors: bool,
     stats: Arc<Stats>,
@@ -235,6 +236,7 @@ impl UpstreamManager {
         configs: Vec<UpstreamConfig>,
         connect_retry_attempts: u32,
         connect_retry_backoff_ms: u64,
+        connect_budget_ms: u64,
         unhealthy_fail_threshold: u32,
         connect_failfast_hard_errors: bool,
         stats: Arc<Stats>,
@@ -248,6 +250,7 @@ impl UpstreamManager {
             upstreams: Arc::new(RwLock::new(states)),
             connect_retry_attempts: connect_retry_attempts.max(1),
             connect_retry_backoff: Duration::from_millis(connect_retry_backoff_ms),
+            connect_budget: Duration::from_millis(connect_budget_ms.max(1)),
             unhealthy_fail_threshold: unhealthy_fail_threshold.max(1),
             connect_failfast_hard_errors,
             stats,
@@ -593,11 +596,27 @@ impl UpstreamManager {
         let mut last_error: Option<ProxyError> = None;
         let mut attempts_used = 0u32;
         for attempt in 1..=self.connect_retry_attempts {
+            let elapsed = connect_started_at.elapsed();
+            if elapsed >= self.connect_budget {
+                last_error = Some(ProxyError::ConnectionTimeout {
+                    addr: target.to_string(),
+                });
+                break;
+            }
+            let remaining_budget = self.connect_budget.saturating_sub(elapsed);
+            let attempt_timeout = Duration::from_secs(DIRECT_CONNECT_TIMEOUT_SECS)
+                .min(remaining_budget);
+            if attempt_timeout.is_zero() {
+                last_error = Some(ProxyError::ConnectionTimeout {
+                    addr: target.to_string(),
+                });
+                break;
+            }
             attempts_used = attempt;
             self.stats.increment_upstream_connect_attempt_total();
             let start = Instant::now();
             match self
-                .connect_via_upstream(&upstream, target, bind_rr.clone())
+                .connect_via_upstream(&upstream, target, bind_rr.clone(), attempt_timeout)
                 .await
             {
                 Ok((stream, egress)) => {
@@ -707,6 +726,7 @@ impl UpstreamManager {
         config: &UpstreamConfig,
         target: SocketAddr,
         bind_rr: Option<Arc<AtomicUsize>>,
+        connect_timeout: Duration,
     ) -> Result<(TcpStream, UpstreamEgressInfo)> {
         match &config.upstream_type {
             UpstreamType::Direct { interface, bind_addresses } => {
@@ -735,7 +755,6 @@ impl UpstreamManager {
                 let std_stream: std::net::TcpStream = socket.into();
                 let stream = TcpStream::from_std(std_stream)?;
 
-                let connect_timeout = Duration::from_secs(DIRECT_CONNECT_TIMEOUT_SECS);
                 match tokio::time::timeout(connect_timeout, stream.writable()).await {
                     Ok(Ok(())) => {}
                     Ok(Err(e)) => return Err(ProxyError::Io(e)),
@@ -762,7 +781,6 @@ impl UpstreamManager {
                 ))
             },
             UpstreamType::Socks4 { address, interface, user_id } => {
-                let connect_timeout = Duration::from_secs(DIRECT_CONNECT_TIMEOUT_SECS);
                 // Try to parse as SocketAddr first (IP:port), otherwise treat as hostname:port
                 let mut stream = if let Ok(proxy_addr) = address.parse::<SocketAddr>() {
                     // IP:port format - use socket with optional interface binding
@@ -841,7 +859,6 @@ impl UpstreamManager {
                 ))
             },
             UpstreamType::Socks5 { address, interface, username, password } => {
-                let connect_timeout = Duration::from_secs(DIRECT_CONNECT_TIMEOUT_SECS);
                 // Try to parse as SocketAddr first (IP:port), otherwise treat as hostname:port
                 let mut stream = if let Ok(proxy_addr) = address.parse::<SocketAddr>() {
                     // IP:port format - use socket with optional interface binding
@@ -1165,7 +1182,14 @@ impl UpstreamManager {
         target: SocketAddr,
     ) -> Result<f64> {
         let start = Instant::now();
-        let _ = self.connect_via_upstream(config, target, bind_rr).await?;
+        let _ = self
+            .connect_via_upstream(
+                config,
+                target,
+                bind_rr,
+                Duration::from_secs(DC_PING_TIMEOUT_SECS),
+            )
+            .await?;
         Ok(start.elapsed().as_secs_f64() * 1000.0)
     }
 
@@ -1337,7 +1361,12 @@ impl UpstreamManager {
                             let start = Instant::now();
                             let result = tokio::time::timeout(
                                 Duration::from_secs(HEALTH_CHECK_CONNECT_TIMEOUT_SECS),
-                                self.connect_via_upstream(&config, endpoint, Some(bind_rr.clone())),
+                                self.connect_via_upstream(
+                                    &config,
+                                    endpoint,
+                                    Some(bind_rr.clone()),
+                                    Duration::from_secs(HEALTH_CHECK_CONNECT_TIMEOUT_SECS),
+                                ),
                             )
                             .await;