use std::collections::{HashMap, HashSet}; use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; use std::time::Instant; use base64::Engine as _; use sha2::{Digest, Sha256}; use zeroize::Zeroizing; use super::{ProfileKey, TOKEN_BYTES, TokenHash}; use crate::config::{WebLimitsConfig, WebRuntimeConfig, WebRuntimeProfile}; use crate::maestro::generation::RuntimeGeneration; use crate::web::session::WebSession; /// One issued bootstrap and optional idempotent session-creation replay state. pub(super) struct Bootstrap { /// Credential and replay-state expiry deadline. pub(super) expires_at: Instant, /// Stable ordering point used for bounded eviction. pub(super) issued_at: Instant, /// Issuing address charged for the unused-bootstrap quota. pub(super) issuance_ip: IpAddr, /// Immutable profile selected during capability validation. pub(super) profile: Arc, /// Digest of the accepted HELLO body for idempotent retry matching. pub(super) body_digest: TokenHash, /// Zeroizing copy returned only for an exact session-creation retry. pub(super) session_token: Zeroizing, /// Created session retained while retry replay remains valid. pub(super) session: Option>, /// Distinguishes unused issuance quota from completed creation replay state. pub(super) used: bool, } /// Bounded replay marker for one explicitly or naturally closed session token. pub(super) struct ClosedToken { /// Deadline after which the token hash may be forgotten. pub(super) expires_at: Instant, /// Canonical host that owned the session. pub(super) host: String, } /// Token-bucket state for one process-wide creation class. #[derive(Default)] pub(super) struct RateState { tokens: f64, last: Option, } struct StreamPortState { active: HashSet, next: u16, } /// Process-wide WEB registries and quota accounting protected by one short lock. #[derive(Default)] pub(super) struct ManagerState { /// Bootstrap credentials indexed by their SHA-256 token hash. pub(super) bootstraps: HashMap, /// Unused bootstrap ownership counts by forwarded client address. pub(super) bootstraps_per_ip: HashMap, /// Live sessions indexed by bearer-token hash. pub(super) sessions: HashMap>, /// Recently closed token hashes retained for idempotent DELETE semantics. pub(super) closed_tokens: HashMap, /// Live session counts by forwarded client address. pub(super) sessions_per_ip: HashMap, /// Live session counts by stable profile key. pub(super) sessions_per_profile: HashMap, /// Live relay-task counts by stable profile key. pub(super) streams_per_profile: HashMap, /// Process-wide live relay-task count. pub(super) streams_live: usize, stream_ports: HashMap<(IpAddr, SocketAddr), StreamPortState>, /// Total process-wide queued byte reservation. pub(super) pending_bytes: usize, /// Total process-wide queued item reservation. pub(super) pending_items: usize, /// Portion of queued bytes charged to the control reserve. pub(super) pending_control_bytes: usize, /// Portion of queued items charged to the control reserve. pub(super) pending_control_items: usize, /// Bootstrap issuance rate limiter. pub(super) bootstrap_rate: RateState, /// Session creation rate limiter. pub(super) session_rate: RateState, /// Logical-stream creation rate limiter. pub(super) stream_rate: RateState, /// Process shutdown admission latch. pub(super) closed: bool, } /// Generates one collision-checked credential and its stable hash key. pub(super) fn new_unique_token( generation: &RuntimeGeneration, state: &ManagerState, ) -> Option<(String, TokenHash)> { for _ in 0..8 { let mut raw = [0u8; TOKEN_BYTES]; generation.rng.fill(&mut raw); let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw); let hash = Sha256::digest(raw).into(); if !state.bootstraps.contains_key(&hash) && !state.sessions.contains_key(&hash) && !state.closed_tokens.contains_key(&hash) { return Some((token, hash)); } } None } /// Returns the precomputed capability as the stable process profile key. pub(super) fn profile_key(profile: &WebRuntimeProfile) -> ProfileKey { profile.capability } /// Re-resolves an issued profile against the active generation without weakening identity. pub(super) fn matching_profile( runtime: &WebRuntimeConfig, expected: &WebRuntimeProfile, ) -> Option> { runtime .profiles .iter() .find(|profile| { profile.host == expected.host && profile.public_addr == expected.public_addr && profile.user == expected.user && profile.secret_mode == expected.secret_mode && profile.carrier == expected.carrier && profile.capability == expected.capability }) .cloned() } /// Applies one token-bucket admission decision at a caller-supplied monotonic time. pub(super) fn allow_rate( state: &mut RateState, now: Instant, per_minute: u32, burst: u32, ) -> bool { let burst = f64::from(burst); if let Some(last) = state.last { let elapsed = now.saturating_duration_since(last).as_secs_f64(); state.tokens = (state.tokens + elapsed * f64::from(per_minute) / 60.0).min(burst); } else { state.tokens = burst; } state.last = Some(now); if state.tokens < 1.0 { return false; } state.tokens -= 1.0; true } /// Evicts the oldest unused bootstrap while preserving used retry state. pub(super) fn evict_oldest_unused_bootstrap(state: &mut ManagerState) -> bool { let Some(hash) = state .bootstraps .iter() .filter(|(_, bootstrap)| !bootstrap.used) .min_by_key(|(_, bootstrap)| bootstrap.issued_at) .map(|(hash, _)| *hash) else { return false; }; remove_bootstrap_locked(state, hash); true } /// Removes expired bootstrap and closed-token entries while the manager lock is held. pub(super) fn remove_expired_locked(state: &mut ManagerState, now: Instant) { let expired = state .bootstraps .iter() .filter_map(|(hash, bootstrap)| (now > bootstrap.expires_at).then_some(*hash)) .collect::>(); for hash in expired { remove_bootstrap_locked(state, hash); } state .closed_tokens .retain(|_, closed| now <= closed.expires_at); } /// Removes one bootstrap and releases its per-address issuance quota when unused. pub(super) fn remove_bootstrap_locked(state: &mut ManagerState, hash: TokenHash) { let Some(bootstrap) = state.bootstraps.remove(&hash) else { return; }; if !bootstrap.used { decrement_map(&mut state.bootstraps_per_ip, &bootstrap.issuance_ip); } } /// Decrements one counted owner and removes its map entry at zero. pub(super) fn decrement_map(values: &mut HashMap, key: &Q) where K: std::borrow::Borrow + std::hash::Hash + Eq, Q: std::hash::Hash + Eq + ?Sized, { let remove = if let Some(value) = values.get_mut(key) { *value = value.saturating_sub(1); *value == 0 } else { false }; if remove { values.remove(key); } } /// Computes the process-wide item reserve required for session control progress. pub(super) fn control_item_reserve(limits: &WebLimitsConfig) -> usize { limits.max_sessions_global.saturating_mul( 16usize.saturating_add(limits.max_streams_per_session.saturating_mul(3)), ) } /// Allocates a non-zero source port unique among live streams for one KDF route. pub(super) fn allocate_stream_port( state: &mut ManagerState, client_ip: IpAddr, public_addr: SocketAddr, ) -> Option { let ports = state .stream_ports .entry((client_ip, public_addr)) .or_insert_with(|| StreamPortState { active: HashSet::new(), next: 1, }); for _ in 0..u16::MAX { let candidate = ports.next; ports.next = ports.next.checked_add(1).unwrap_or(1); if ports.active.insert(candidate) { return Some(candidate); } } None } /// Releases one source port and reclaims empty per-route allocator state. pub(super) fn release_stream_port( state: &mut ManagerState, client_ip: IpAddr, public_addr: SocketAddr, peer_port: u16, ) -> bool { let key = (client_ip, public_addr); let Some(ports) = state.stream_ports.get_mut(&key) else { return false; }; let removed = ports.active.remove(&peer_port); if ports.active.is_empty() { state.stream_ports.remove(&key); } removed } #[cfg(test)] mod tests { use super::*; #[test] fn synthetic_ports_are_unique_per_live_route_and_state_is_reclaimed() { let mut state = ManagerState::default(); let client_ip = "192.0.2.10".parse().unwrap(); let public_addr = "203.0.113.10:443".parse().unwrap(); let first = allocate_stream_port(&mut state, client_ip, public_addr).unwrap(); let second = allocate_stream_port(&mut state, client_ip, public_addr).unwrap(); assert_ne!(first, second); assert!(release_stream_port( &mut state, client_ip, public_addr, first, )); assert!(release_stream_port( &mut state, client_ip, public_addr, second, )); assert!(state.stream_ports.is_empty()); } }