WEB

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
Co-Authored-By: John Preston <17900494+john-preston@users.noreply.github.com>

src/web/manager/state.rs

@@ -0,0 +1,293 @@
+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 {
+    /// Generation that issued the bootstrap.
+    pub(super) generation_id: u64,
+    /// Credential and replay-state expiry deadline.
+    pub(super) expires_at: Instant,
+    /// Stable ordering point used for bounded eviction.
+    pub(super) issued_at: Instant,
+    /// Forwarded client address that owns this credential.
+    pub(super) issuance_ip: IpAddr,
+    /// Immutable profile selected during capability validation.
+    pub(super) profile: Arc<WebRuntimeProfile>,
+    /// 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<String>,
+    /// Created session retained while retry replay remains valid.
+    pub(super) session: Option<Arc<WebSession>>,
+    /// 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<Instant>,
+}
+
+struct StreamPortState {
+    active: HashSet<u16>,
+    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<TokenHash, Bootstrap>,
+    /// Unused bootstrap ownership counts by forwarded client address.
+    pub(super) bootstraps_per_ip: HashMap<IpAddr, usize>,
+    /// Live sessions indexed by bearer-token hash.
+    pub(super) sessions: HashMap<TokenHash, Arc<WebSession>>,
+    /// Recently closed token hashes retained for idempotent DELETE semantics.
+    pub(super) closed_tokens: HashMap<TokenHash, ClosedToken>,
+    /// Live session counts by forwarded client address.
+    pub(super) sessions_per_ip: HashMap<IpAddr, usize>,
+    /// Live session counts by stable profile key.
+    pub(super) sessions_per_profile: HashMap<ProfileKey, usize>,
+    /// Live relay-task counts by stable profile key.
+    pub(super) streams_per_profile: HashMap<ProfileKey, usize>,
+    /// 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<Arc<WebRuntimeProfile>> {
+    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.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::<Vec<_>>();
+    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<K, Q>(values: &mut HashMap<K, usize>, key: &Q)
+where
+    K: std::borrow::Borrow<Q> + 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<u16> {
+    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());
+    }
+}