Harden overload auth scans and masking safeguards

src/proxy/handshake.rs

@@ -55,6 +55,7 @@ const STICKY_HINT_MAX_ENTRIES: usize = 65_536;
 const CANDIDATE_HINT_TRACK_CAP: usize = 64;
 const OVERLOAD_CANDIDATE_BUDGET_HINTED: usize = 16;
 const OVERLOAD_CANDIDATE_BUDGET_UNHINTED: usize = 8;
+const OVERLOAD_FULL_SCAN_USER_THRESHOLD: usize = CANDIDATE_HINT_TRACK_CAP;
 const RECENT_USER_RING_SCAN_LIMIT: usize = 32;
 
 type HmacSha256 = Hmac<Sha256>;
@@ -242,6 +243,9 @@ fn budget_for_validation(total_users: usize, overload: bool, has_hint: bool) ->
     if !overload {
         return total_users;
     }
+    if total_users <= OVERLOAD_FULL_SCAN_USER_THRESHOLD {
+        return total_users;
+    }
     let cap = if has_hint {
         OVERLOAD_CANDIDATE_BUDGET_HINTED
     } else {
@@ -250,6 +254,28 @@ fn budget_for_validation(total_users: usize, overload: bool, has_hint: bool) ->
     total_users.min(cap.max(1))
 }
 
+// Rotate partial overload scans across larger snapshots so one truncated
+// validation window does not permanently starve the same cold users.
+fn candidate_scan_start_offset_in(
+    shared: &ProxySharedState,
+    peer_ip: IpAddr,
+    total_users: usize,
+    candidate_budget: usize,
+) -> usize {
+    if total_users == 0 || candidate_budget >= total_users {
+        return 0;
+    }
+
+    let seq = shared
+        .handshake
+        .auth_candidate_scan_seq
+        .fetch_add(1, Ordering::Relaxed);
+    let mut hasher = shared.handshake.auth_probe_eviction_hasher.build_hasher();
+    peer_ip.hash(&mut hasher);
+    seq.hash(&mut hasher);
+    hasher.finish() as usize % total_users
+}
+
 fn parse_tls_auth_material(
     handshake: &[u8],
     ignore_time_skew: bool,
@@ -1312,7 +1338,14 @@ where
         }
 
         if !matched && !budget_exhausted {
-            for idx in 0..snapshot.entries().len() {
+            let fallback_start = candidate_scan_start_offset_in(
+                shared,
+                peer.ip(),
+                snapshot.entries().len(),
+                candidate_budget,
+            );
+            for offset in 0..snapshot.entries().len() {
+                let idx = (fallback_start + offset) % snapshot.entries().len();
                 let Some(user_id) = u32::try_from(idx).ok() else {
                     break;
                 };
@@ -1679,7 +1712,14 @@ where
         }
 
         if !matched && !budget_exhausted {
-            for idx in 0..snapshot.entries().len() {
+            let fallback_start = candidate_scan_start_offset_in(
+                shared,
+                peer.ip(),
+                snapshot.entries().len(),
+                candidate_budget,
+            );
+            for offset in 0..snapshot.entries().len() {
+                let idx = (fallback_start + offset) % snapshot.entries().len();
                 let Some(user_id) = u32::try_from(idx).ok() else {
                     break;
                 };

src/proxy/masking.rs

@@ -506,6 +506,40 @@ fn is_mask_target_local_listener_with_interfaces(
     local_addr: SocketAddr,
     resolved_override: Option<SocketAddr>,
     interface_ips: &[IpAddr],
+) -> bool {
+    let resolved_candidates = resolved_override
+        .as_ref()
+        .map(std::slice::from_ref)
+        .unwrap_or(&[]);
+    is_mask_target_local_listener_candidates_with_interfaces(
+        mask_host,
+        mask_port,
+        local_addr,
+        resolved_candidates,
+        interface_ips,
+    )
+}
+
+fn mask_ip_targets_local_listener(
+    mask_ip: IpAddr,
+    local_ip: IpAddr,
+    interface_ips: &[IpAddr],
+) -> bool {
+    let mask_ip = canonical_ip(mask_ip);
+    if mask_ip == local_ip {
+        return true;
+    }
+
+    local_ip.is_unspecified()
+        && (mask_ip.is_loopback() || mask_ip.is_unspecified() || interface_ips.contains(&mask_ip))
+}
+
+fn is_mask_target_local_listener_candidates_with_interfaces(
+    mask_host: &str,
+    mask_port: u16,
+    local_addr: SocketAddr,
+    resolved_candidates: &[SocketAddr],
+    interface_ips: &[IpAddr],
 ) -> bool {
     if mask_port != local_addr.port() {
         return false;
@@ -514,31 +548,14 @@ fn is_mask_target_local_listener_with_interfaces(
     let local_ip = canonical_ip(local_addr.ip());
     let literal_mask_ip = parse_mask_host_ip_literal(mask_host).map(canonical_ip);
 
-    if let Some(addr) = resolved_override {
-        let resolved_ip = canonical_ip(addr.ip());
-        if resolved_ip == local_ip {
-            return true;
-        }
-
-        if local_ip.is_unspecified()
-            && (resolved_ip.is_loopback()
-                || resolved_ip.is_unspecified()
-                || interface_ips.contains(&resolved_ip))
-        {
+    for addr in resolved_candidates {
+        if mask_ip_targets_local_listener(addr.ip(), local_ip, interface_ips) {
             return true;
         }
     }
 
     if let Some(mask_ip) = literal_mask_ip {
-        if mask_ip == local_ip {
-            return true;
-        }
-
-        if local_ip.is_unspecified()
-            && (mask_ip.is_loopback()
-                || mask_ip.is_unspecified()
-                || interface_ips.contains(&mask_ip))
-        {
+        if mask_ip_targets_local_listener(mask_ip, local_ip, interface_ips) {
             return true;
         }
     }
@@ -572,21 +589,67 @@ async fn is_mask_target_local_listener_async(
     mask_port: u16,
     local_addr: SocketAddr,
     resolved_override: Option<SocketAddr>,
+) -> bool {
+    let resolved_candidates = resolved_override
+        .as_ref()
+        .map(std::slice::from_ref)
+        .unwrap_or(&[]);
+    is_mask_target_local_listener_candidates_async(
+        mask_host,
+        mask_port,
+        local_addr,
+        resolved_candidates,
+    )
+    .await
+}
+
+async fn is_mask_target_local_listener_candidates_async(
+    mask_host: &str,
+    mask_port: u16,
+    local_addr: SocketAddr,
+    resolved_candidates: &[SocketAddr],
 ) -> bool {
     if mask_port != local_addr.port() {
         return false;
     }
 
     let interfaces = local_interface_ips_async().await;
-    is_mask_target_local_listener_with_interfaces(
+    is_mask_target_local_listener_candidates_with_interfaces(
         mask_host,
         mask_port,
         local_addr,
-        resolved_override,
+        resolved_candidates,
         &interfaces,
     )
 }
 
+// Resolve hostnames through the same OS DNS path `TcpStream::connect` uses so
+// self-target rejection also catches loopback and local-interface hostnames.
+async fn resolve_mask_target_candidates(
+    mask_host: &str,
+    mask_port: u16,
+    resolved_override: Option<SocketAddr>,
+) -> Vec<SocketAddr> {
+    if let Some(addr) = resolved_override {
+        return vec![addr];
+    }
+
+    if parse_mask_host_ip_literal(mask_host).is_some() {
+        return Vec::new();
+    }
+
+    let mut resolved = Vec::new();
+    if let Ok(addrs) = tokio::net::lookup_host((mask_host, mask_port)).await {
+        for addr in addrs {
+            if !resolved.contains(&addr) {
+                resolved.push(addr);
+            }
+        }
+    }
+
+    resolved
+}
+
 fn masking_beobachten_ttl(config: &ProxyConfig) -> Duration {
     let minutes = config.general.beobachten_minutes;
     let clamped = minutes.clamp(1, 24 * 60);
@@ -731,8 +794,15 @@ pub async fn handle_bad_client<R, W>(
     // Self-referential masking can create recursive proxy loops under
     // misconfiguration and leak distinguishable load spikes to adversaries.
     let resolved_mask_addr = resolve_socket_addr(mask_host, mask_port);
-    if is_mask_target_local_listener_async(mask_host, mask_port, local_addr, resolved_mask_addr)
-        .await
+    let resolved_mask_candidates =
+        resolve_mask_target_candidates(mask_host, mask_port, resolved_mask_addr).await;
+    if is_mask_target_local_listener_candidates_async(
+        mask_host,
+        mask_port,
+        local_addr,
+        &resolved_mask_candidates,
+    )
+    .await
     {
         let outcome_started = Instant::now();
         debug!(

src/proxy/shared_state.rs

@@ -48,6 +48,7 @@ pub(crate) struct HandshakeSharedState {
     pub(crate) sticky_user_by_sni_hash: DashMap<u64, u32>,
     pub(crate) recent_user_ring: Box<[AtomicU32]>,
     pub(crate) recent_user_ring_seq: AtomicU64,
+    pub(crate) auth_candidate_scan_seq: AtomicU64,
     pub(crate) auth_expensive_checks_total: AtomicU64,
     pub(crate) auth_budget_exhausted_total: AtomicU64,
 }
@@ -86,6 +87,7 @@ impl ProxySharedState {
                     .collect::<Vec<_>>()
                     .into_boxed_slice(),
                 recent_user_ring_seq: AtomicU64::new(0),
+                auth_candidate_scan_seq: AtomicU64::new(0),
                 auth_expensive_checks_total: AtomicU64::new(0),
                 auth_budget_exhausted_total: AtomicU64::new(0),
             },

src/proxy/tests/handshake_security_tests.rs

@@ -1146,9 +1146,9 @@ async fn tls_overload_budget_limits_candidate_scan_depth() {
     let mut config = ProxyConfig::default();
     config.access.users.clear();
     config.access.ignore_time_skew = true;
-    for idx in 0..32u8 {
+    for idx in 0..96u8 {
         config.access.users.insert(
-            format!("user-{idx}"),
+            format!("user-{idx:02}"),
             format!("{:032x}", u128::from(idx) + 1),
         );
     }
@@ -1203,6 +1203,64 @@ async fn tls_overload_budget_limits_candidate_scan_depth() {
     );
 }
 
+#[tokio::test]
+async fn tls_overload_full_scans_small_runtime_snapshot_to_preserve_cold_user_auth() {
+    let mut config = ProxyConfig::default();
+    config.access.users.clear();
+    config.access.ignore_time_skew = true;
+    for idx in 0..32u8 {
+        config.access.users.insert(
+            format!("user-{idx:02}"),
+            format!("{:032x}", u128::from(idx) + 1),
+        );
+    }
+    config.rebuild_runtime_user_auth().unwrap();
+
+    let replay_checker = ReplayChecker::new(128, Duration::from_secs(60));
+    let rng = SecureRandom::new();
+    let shared = ProxySharedState::new();
+    let now = Instant::now();
+    {
+        let mut saturation = shared.handshake.auth_probe_saturation.lock().unwrap();
+        *saturation = Some(AuthProbeSaturationState {
+            fail_streak: AUTH_PROBE_BACKOFF_START_FAILS,
+            blocked_until: now + Duration::from_millis(200),
+            last_seen: now,
+        });
+    }
+
+    let peer: SocketAddr = "198.51.100.214:44326".parse().unwrap();
+    let mut secret = [0u8; 16];
+    secret[15] = 32;
+    let handshake = make_valid_tls_handshake(&secret, 0);
+
+    let result = handle_tls_handshake_with_shared(
+        &handshake,
+        tokio::io::empty(),
+        tokio::io::sink(),
+        peer,
+        &config,
+        &replay_checker,
+        &rng,
+        None,
+        shared.as_ref(),
+    )
+    .await;
+
+    assert!(
+        matches!(result, HandshakeResult::Success(_)),
+        "overload mode must still authenticate valid cold users when runtime snapshot stays small"
+    );
+    assert_eq!(
+        shared
+            .handshake
+            .auth_expensive_checks_total
+            .load(Ordering::Relaxed),
+        32,
+        "small saturated snapshots must remain fully scannable"
+    );
+}
+
 #[tokio::test]
 async fn mtproto_runtime_snapshot_prefers_preferred_user_hint() {
     let mut config = ProxyConfig::default();
@@ -1255,6 +1313,66 @@ async fn mtproto_runtime_snapshot_prefers_preferred_user_hint() {
     );
 }
 
+#[tokio::test]
+async fn mtproto_overload_full_scans_small_runtime_snapshot_to_preserve_cold_user_auth() {
+    let mut config = ProxyConfig::default();
+    config.general.modes.secure = true;
+    config.access.users.clear();
+    config.access.ignore_time_skew = true;
+    for idx in 0..32u8 {
+        config.access.users.insert(
+            format!("user-{idx:02}"),
+            format!("{:032x}", u128::from(idx) + 1),
+        );
+    }
+    config.rebuild_runtime_user_auth().unwrap();
+
+    let shared = ProxySharedState::new();
+    let now = Instant::now();
+    {
+        let mut saturation = shared.handshake.auth_probe_saturation.lock().unwrap();
+        *saturation = Some(AuthProbeSaturationState {
+            fail_streak: AUTH_PROBE_BACKOFF_START_FAILS,
+            blocked_until: now + Duration::from_millis(200),
+            last_seen: now,
+        });
+    }
+
+    let replay_checker = ReplayChecker::new(128, Duration::from_secs(60));
+    let handshake = make_valid_mtproto_handshake(
+        "00000000000000000000000000000020",
+        ProtoTag::Secure,
+        2,
+    );
+    let peer: SocketAddr = "198.51.100.215:44326".parse().unwrap();
+
+    let result = handle_mtproto_handshake_with_shared(
+        &handshake,
+        tokio::io::empty(),
+        tokio::io::sink(),
+        peer,
+        &config,
+        &replay_checker,
+        false,
+        None,
+        shared.as_ref(),
+    )
+    .await;
+
+    assert!(
+        matches!(result, HandshakeResult::Success(_)),
+        "overload mode must still authenticate valid direct MTProto users when runtime snapshot stays small"
+    );
+    assert_eq!(
+        shared
+            .handshake
+            .auth_expensive_checks_total
+            .load(Ordering::Relaxed),
+        32,
+        "small saturated MTProto snapshots must remain fully scannable"
+    );
+}
+
 #[tokio::test]
 async fn alpn_enforce_rejects_unsupported_client_alpn() {
     let secret = [0x33u8; 16];

src/proxy/tests/masking_self_target_loop_security_tests.rs

@@ -88,6 +88,45 @@ async fn self_target_fallback_refuses_recursive_loopback_connect() {
     );
 }
 
+#[tokio::test]
+async fn self_target_fallback_refuses_recursive_hostname_connect() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let local_addr = listener.local_addr().unwrap();
+    let accept_task = tokio::spawn(async move {
+        timeout(Duration::from_millis(120), listener.accept())
+            .await
+            .is_ok()
+    });
+
+    let mut config = ProxyConfig::default();
+    config.general.beobachten = false;
+    config.censorship.mask = true;
+    config.censorship.mask_unix_sock = None;
+    config.censorship.mask_host = Some("localhost".to_string());
+    config.censorship.mask_port = local_addr.port();
+    config.censorship.mask_proxy_protocol = 0;
+
+    let peer: SocketAddr = "203.0.113.99:55099".parse().unwrap();
+    let beobachten = BeobachtenStore::new();
+
+    handle_bad_client(
+        tokio::io::empty(),
+        tokio::io::sink(),
+        b"GET /",
+        peer,
+        local_addr,
+        &config,
+        &beobachten,
+    )
+    .await;
+
+    let accepted = accept_task.await.unwrap();
+    assert!(
+        !accepted,
+        "hostname self-target masking must fail closed without connecting to local listener"
+    );
+}
+
 #[tokio::test]
 async fn same_ip_different_port_still_forwards_to_mask_backend() {
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();