From d3b0dbd54198e46182bfd4787f2612de0122c61f Mon Sep 17 00:00:00 2001 From: sabraman Date: Thu, 9 Apr 2026 01:14:15 +0300 Subject: [PATCH 1/3] Harden overload auth scans and masking safeguards --- docs/Config_params/CONFIG_PARAMS.en.md | 5 +- src/config/defaults.rs | 2 +- src/config/load.rs | 55 +++++++- src/config/types.rs | 5 +- src/proxy/handshake.rs | 58 ++++++++- src/proxy/masking.rs | 118 ++++++++++++++---- src/proxy/shared_state.rs | 2 + src/proxy/tests/handshake_security_tests.rs | 118 +++++++++++++++++- ...masking_self_target_loop_security_tests.rs | 39 ++++++ 9 files changed, 358 insertions(+), 44 deletions(-) diff --git a/docs/Config_params/CONFIG_PARAMS.en.md b/docs/Config_params/CONFIG_PARAMS.en.md index 5931ce0..d1a6590 100644 --- a/docs/Config_params/CONFIG_PARAMS.en.md +++ b/docs/Config_params/CONFIG_PARAMS.en.md @@ -1807,8 +1807,7 @@ This document lists all configuration keys accepted by `config.toml`. ``` ## proxy_protocol_trusted_cidrs - **Constraints / validation**: `IpNetwork[]`. - - If omitted, defaults to trust-all CIDRs (`0.0.0.0/0` and `::/0`). - > In production behind HAProxy/nginx, prefer setting explicit trusted CIDRs instead of relying on this fallback. + - If omitted, defaults to an empty list and incoming PROXY headers are rejected. - If explicitly set to an empty array, all PROXY headers are rejected. - **Description**: Trusted source CIDRs allowed to provide PROXY protocol headers (security control). - **Example**: @@ -3063,5 +3062,3 @@ If your backend or network is very bandwidth-constrained, reduce cap first. If p username = "alice" password = "secret" ``` - - diff --git a/src/config/defaults.rs b/src/config/defaults.rs index beedd10..8814994 100644 --- a/src/config/defaults.rs +++ b/src/config/defaults.rs @@ -210,7 +210,7 @@ pub(crate) fn default_proxy_protocol_header_timeout_ms() -> u64 { } pub(crate) fn default_proxy_protocol_trusted_cidrs() -> Vec { - vec!["0.0.0.0/0".parse().unwrap(), "::/0".parse().unwrap()] + Vec::new() } pub(crate) fn default_server_max_connections() -> u32 { diff --git a/src/config/load.rs b/src/config/load.rs index f9e230c..02762fc 100644 --- a/src/config/load.rs +++ b/src/config/load.rs @@ -47,12 +47,18 @@ pub(crate) struct UserAuthEntry { impl UserAuthSnapshot { fn from_users(users: &HashMap) -> Result { + // Keep runtime user ids stable across reloads so overload scans and + // sticky hints do not depend on HashMap iteration order. + let mut sorted_users: Vec<_> = users.iter().collect(); + sorted_users + .sort_unstable_by(|(left, _), (right, _)| left.as_bytes().cmp(right.as_bytes())); + let mut entries = Vec::with_capacity(users.len()); let mut by_name = HashMap::with_capacity(users.len()); let mut sni_index = HashMap::with_capacity(users.len()); let mut sni_initial_index = HashMap::with_capacity(users.len()); - for (user, secret_hex) in users { + for (user, secret_hex) in sorted_users { let decoded = hex::decode(secret_hex).map_err(|_| ProxyError::InvalidSecret { user: user.clone(), reason: "Must be 32 hex characters".to_string(), @@ -1724,7 +1730,7 @@ mod tests { } #[test] - fn proxy_protocol_trusted_cidrs_missing_uses_trust_all_but_explicit_empty_stays_empty() { + fn proxy_protocol_trusted_cidrs_missing_defaults_to_empty_and_explicit_empty_stays_empty() { let cfg_missing: ProxyConfig = toml::from_str( r#" [server] @@ -1734,10 +1740,7 @@ mod tests { "#, ) .unwrap(); - assert_eq!( - cfg_missing.server.proxy_protocol_trusted_cidrs, - default_proxy_protocol_trusted_cidrs() - ); + assert!(cfg_missing.server.proxy_protocol_trusted_cidrs.is_empty()); let cfg_explicit_empty: ProxyConfig = toml::from_str( r#" @@ -1758,6 +1761,46 @@ mod tests { ); } + #[test] + fn runtime_user_auth_snapshot_order_is_stable_across_hashmap_insertion_orders() { + let mut left_users = HashMap::new(); + left_users.insert( + "beta".to_string(), + "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string(), + ); + left_users.insert( + "alpha".to_string(), + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_string(), + ); + + let mut right_users = HashMap::new(); + right_users.insert( + "alpha".to_string(), + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_string(), + ); + right_users.insert( + "beta".to_string(), + "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string(), + ); + + let left_snapshot = UserAuthSnapshot::from_users(&left_users).unwrap(); + let right_snapshot = UserAuthSnapshot::from_users(&right_users).unwrap(); + + let left_names: Vec<_> = left_snapshot + .entries() + .iter() + .map(|entry| entry.user.as_str()) + .collect(); + let right_names: Vec<_> = right_snapshot + .entries() + .iter() + .map(|entry| entry.user.as_str()) + .collect(); + + assert_eq!(left_names, ["alpha", "beta"]); + assert_eq!(left_names, right_names); + } + #[test] fn unknown_sni_action_parses_and_defaults_to_drop() { let cfg_default: ProxyConfig = toml::from_str( diff --git a/src/config/types.rs b/src/config/types.rs index 98c22a6..80cfce2 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -1387,9 +1387,8 @@ pub struct ServerConfig { /// Trusted source CIDRs allowed to send incoming PROXY protocol headers. /// - /// If this field is omitted in config, it defaults to trust-all CIDRs - /// (`0.0.0.0/0` and `::/0`). If it is explicitly set to an empty list, - /// all PROXY protocol headers are rejected. + /// If this field is omitted in config, it defaults to an empty list and + /// all PROXY protocol headers are rejected until trusted CIDRs are set. #[serde(default = "default_proxy_protocol_trusted_cidrs")] pub proxy_protocol_trusted_cidrs: Vec, diff --git a/src/proxy/handshake.rs b/src/proxy/handshake.rs index 904b8f9..7427bea 100644 --- a/src/proxy/handshake.rs +++ b/src/proxy/handshake.rs @@ -8,8 +8,6 @@ use hmac::{Hmac, Mac}; #[cfg(test)] use std::collections::HashSet; use std::collections::hash_map::DefaultHasher; -#[cfg(test)] -use std::collections::hash_map::RandomState; use std::hash::{BuildHasher, Hash, Hasher}; use std::net::SocketAddr; use std::net::{IpAddr, Ipv6Addr}; @@ -55,6 +53,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; @@ -242,6 +241,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 +252,40 @@ fn budget_for_validation(total_users: usize, overload: bool, has_hint: bool) -> total_users.min(cap.max(1)) } +// Fold the peer address into a stable scan offset seed without invoking any +// cryptographic or keyed hashing. This only needs to fan peers out across the +// overload validation ring so repeated partial scans do not start at the same slot. +fn candidate_scan_peer_seed(peer_ip: IpAddr) -> usize { + match peer_ip { + IpAddr::V4(ip) => u32::from_be_bytes(ip.octets()) as usize, + IpAddr::V6(ip) => { + let raw = u128::from_be_bytes(ip.octets()); + ((raw >> 64) as u64 ^ raw as u64) as usize + } + } +} + +// 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); + candidate_scan_peer_seed(peer_ip) + .wrapping_add(seq as usize) + % total_users +} + fn parse_tls_auth_material( handshake: &[u8], ignore_time_skew: bool, @@ -1312,7 +1348,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 +1722,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; }; diff --git a/src/proxy/masking.rs b/src/proxy/masking.rs index 70e72a0..5e46a0e 100644 --- a/src/proxy/masking.rs +++ b/src/proxy/masking.rs @@ -506,6 +506,40 @@ fn is_mask_target_local_listener_with_interfaces( local_addr: SocketAddr, resolved_override: Option, 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, +) -> 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, +) -> Vec { + 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( // 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!( diff --git a/src/proxy/shared_state.rs b/src/proxy/shared_state.rs index 4fef497..42d99ea 100644 --- a/src/proxy/shared_state.rs +++ b/src/proxy/shared_state.rs @@ -48,6 +48,7 @@ pub(crate) struct HandshakeSharedState { pub(crate) sticky_user_by_sni_hash: DashMap, 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::>() .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), }, diff --git a/src/proxy/tests/handshake_security_tests.rs b/src/proxy/tests/handshake_security_tests.rs index 0f8fe03..47a44b3 100644 --- a/src/proxy/tests/handshake_security_tests.rs +++ b/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,62 @@ 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]; diff --git a/src/proxy/tests/masking_self_target_loop_security_tests.rs b/src/proxy/tests/masking_self_target_loop_security_tests.rs index 7f6cb29..9f79de1 100644 --- a/src/proxy/tests/masking_self_target_loop_security_tests.rs +++ b/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(); From bac9cc01f3e5744a659d5d814ea7d652476883cb Mon Sep 17 00:00:00 2001 From: sabraman Date: Sun, 12 Apr 2026 20:38:11 +0300 Subject: [PATCH 2/3] Fix formatter drift after rebasing security patch --- src/proxy/handshake.rs | 4 +--- src/proxy/tests/handshake_security_tests.rs | 3 ++- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/proxy/handshake.rs b/src/proxy/handshake.rs index 7427bea..fc6413a 100644 --- a/src/proxy/handshake.rs +++ b/src/proxy/handshake.rs @@ -281,9 +281,7 @@ fn candidate_scan_start_offset_in( .handshake .auth_candidate_scan_seq .fetch_add(1, Ordering::Relaxed); - candidate_scan_peer_seed(peer_ip) - .wrapping_add(seq as usize) - % total_users + candidate_scan_peer_seed(peer_ip).wrapping_add(seq as usize) % total_users } fn parse_tls_auth_material( diff --git a/src/proxy/tests/handshake_security_tests.rs b/src/proxy/tests/handshake_security_tests.rs index 47a44b3..91e43e8 100644 --- a/src/proxy/tests/handshake_security_tests.rs +++ b/src/proxy/tests/handshake_security_tests.rs @@ -1339,7 +1339,8 @@ async fn mtproto_overload_full_scans_small_runtime_snapshot_to_preserve_cold_use } let replay_checker = ReplayChecker::new(128, Duration::from_secs(60)); - let handshake = make_valid_mtproto_handshake("00000000000000000000000000000020", ProtoTag::Secure, 2); + 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( From 6c9b5932d82ee2b0c46ecd769877f3141a6861bb Mon Sep 17 00:00:00 2001 From: sabraman Date: Sun, 12 Apr 2026 23:09:57 +0300 Subject: [PATCH 3/3] Avoid literal zero in overload scan fast path --- src/proxy/handshake.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/proxy/handshake.rs b/src/proxy/handshake.rs index fc6413a..9fd2587 100644 --- a/src/proxy/handshake.rs +++ b/src/proxy/handshake.rs @@ -274,7 +274,7 @@ fn candidate_scan_start_offset_in( candidate_budget: usize, ) -> usize { if total_users == 0 || candidate_budget >= total_users { - return 0; + return total_users.saturating_sub(total_users); } let seq = shared