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/admission.rs

@@ -0,0 +1,130 @@
+use std::net::{IpAddr, SocketAddr};
+use std::sync::atomic::Ordering;
+use std::time::Instant;
+
+use super::state::{
+    allocate_stream_port, allow_rate, decrement_map, release_stream_port,
+};
+use super::{ProfileKey, WebProcessRuntime};
+
+impl WebProcessRuntime {
+    /// Reserves one process-wide and per-profile live logical-stream slot.
+    pub(crate) fn try_acquire_stream(
+        &self,
+        profile_key: ProfileKey,
+        max_streams: usize,
+        client_ip: IpAddr,
+        public_addr: SocketAddr,
+    ) -> Option<u16> {
+        let now = Instant::now();
+        let mut state = self.state.lock();
+        if state.closed
+            || state.streams_live >= self.limits.max_streams_global
+            || state
+                .streams_per_profile
+                .get(&profile_key)
+                .copied()
+                .unwrap_or(0)
+                >= max_streams
+            || !allow_rate(
+                &mut state.stream_rate,
+                now,
+                self.limits.new_streams_per_minute,
+                self.limits.new_streams_burst,
+            )
+        {
+            self.streams_rejected.fetch_add(1, Ordering::Relaxed);
+            self.limit_hits.fetch_add(1, Ordering::Relaxed);
+            return None;
+        }
+        let Some(peer_port) = allocate_stream_port(&mut state, client_ip, public_addr) else {
+            self.streams_rejected.fetch_add(1, Ordering::Relaxed);
+            self.limit_hits.fetch_add(1, Ordering::Relaxed);
+            return None;
+        };
+        state.streams_live += 1;
+        *state
+            .streams_per_profile
+            .entry(profile_key)
+            .or_insert(0) += 1;
+        self.streams_opened.fetch_add(1, Ordering::Relaxed);
+        Some(peer_port)
+    }
+
+    /// Releases one live logical-stream slot after its relay task exits.
+    pub(crate) fn release_stream(
+        &self,
+        profile_key: ProfileKey,
+        client_ip: IpAddr,
+        public_addr: SocketAddr,
+        peer_port: u16,
+    ) {
+        let mut state = self.state.lock();
+        if !release_stream_port(&mut state, client_ip, public_addr, peer_port) {
+            return;
+        }
+        state.streams_live = state.streams_live.saturating_sub(1);
+        decrement_map(&mut state.streams_per_profile, &profile_key);
+    }
+
+    /// Records a logical stream rejected outside manager quota acquisition.
+    pub(crate) fn record_stream_rejected(&self) {
+        self.streams_rejected.fetch_add(1, Ordering::Relaxed);
+        self.record_limit_hit();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arc_swap::ArcSwap;
+
+    use super::*;
+    use crate::config::ProxyConfig;
+    use crate::maestro::generation::test_runtime_generation;
+    use crate::web::session::QUEUE_ITEM_COST;
+
+    #[tokio::test]
+    async fn global_downlink_budget_preserves_one_maximum_uplink_batch() {
+        let generation = test_runtime_generation(1, ProxyConfig::default());
+        let runtime = WebProcessRuntime::start(Arc::new(ArcSwap::from(generation)));
+        let control_items = super::super::state::control_item_reserve(&runtime.limits);
+        let data_bytes = runtime
+            .limits
+            .pending_bytes_global
+            .saturating_sub(runtime.limits.control_bytes_global);
+        let data_items = runtime
+            .limits
+            .pending_items_global
+            .saturating_sub(control_items);
+        let uplink_bytes = runtime
+            .limits
+            .max_body_bytes
+            .saturating_add(runtime.limits.max_frames_per_body * QUEUE_ITEM_COST);
+        let downlink_bytes = data_bytes - uplink_bytes;
+        let downlink_items = data_items - runtime.limits.max_frames_per_body;
+
+        assert!(runtime.try_reserve_pending(
+            downlink_bytes,
+            downlink_items,
+            false,
+            true,
+        ));
+        assert!(runtime.try_reserve_pending(
+            uplink_bytes,
+            runtime.limits.max_frames_per_body,
+            false,
+            false,
+        ));
+        assert!(!runtime.try_reserve_pending(1, 1, false, true));
+
+        runtime.release_pending(downlink_bytes, downlink_items, false);
+        runtime.release_pending(
+            uplink_bytes,
+            runtime.limits.max_frames_per_body,
+            false,
+        );
+        runtime.shutdown().await;
+    }
+}

src/web/manager/lifecycle.rs

@@ -0,0 +1,149 @@
+use std::net::IpAddr;
+use std::sync::atomic::Ordering;
+use std::time::{Duration, Instant};
+
+use tracing::info;
+
+use super::{ProfileKey, TokenHash, WebProcessRuntime};
+use super::state::{
+    ClosedToken, decrement_map, remove_bootstrap_locked, remove_expired_locked,
+};
+
+impl WebProcessRuntime {
+    /// Removes one closed session and retains a bounded host-bound replay marker.
+    pub(crate) fn session_finished(
+        &self,
+        hash: TokenHash,
+        client_ip: IpAddr,
+        profile_key: ProfileKey,
+        profile_host: &str,
+    ) {
+        let mut state = self.state.lock();
+        if state.sessions.remove(&hash).is_none() {
+            return;
+        }
+        decrement_map(&mut state.sessions_per_ip, &client_ip);
+        decrement_map(&mut state.sessions_per_profile, &profile_key);
+        let expiry = Instant::now()
+            + Duration::from_secs(
+                self.active_runtime
+                    .load()
+                    .config()
+                    .web
+                    .timeouts
+                    .bootstrap_lifetime_secs,
+            );
+        state.closed_tokens.insert(
+            hash,
+            ClosedToken {
+                expires_at: expiry,
+                host: profile_host.to_string(),
+            },
+        );
+        while state.closed_tokens.len() > self.limits.max_sessions_global.saturating_mul(16) {
+            let Some(oldest) = state
+                .closed_tokens
+                .iter()
+                .min_by_key(|(_, closed)| closed.expires_at)
+                .map(|(hash, _)| *hash)
+            else {
+                break;
+            };
+            state.closed_tokens.remove(&oldest);
+        }
+        let bootstrap_hashes = state
+            .bootstraps
+            .iter()
+            .filter_map(|(bootstrap_hash, bootstrap)| {
+                bootstrap
+                    .session
+                    .as_ref()
+                    .is_some_and(|session| session.token_hash() == hash)
+                    .then_some(*bootstrap_hash)
+            })
+            .collect::<Vec<_>>();
+        for bootstrap_hash in bootstrap_hashes {
+            remove_bootstrap_locked(&mut state, bootstrap_hash);
+        }
+        self.sessions_closed.fetch_add(1, Ordering::Relaxed);
+    }
+
+    /// Stops issuance, closes all sessions, and joins bounded child work.
+    pub(crate) async fn shutdown(&self) {
+        self.shutdown.cancel();
+        let sessions = {
+            let mut state = self.state.lock();
+            state.closed = true;
+            state.bootstraps.clear();
+            state.bootstraps_per_ip.clear();
+            state.sessions.values().cloned().collect::<Vec<_>>()
+        };
+        for session in &sessions {
+            session.close();
+        }
+        let timeout_secs = self
+            .active_runtime
+            .load()
+            .config()
+            .web
+            .timeouts
+            .shutdown_secs;
+        let waits = async {
+            for session in sessions {
+                session.wait().await;
+            }
+        };
+        let _ = tokio::time::timeout(Duration::from_secs(timeout_secs), waits).await;
+        self.tasks.close();
+        let _ = tokio::time::timeout(Duration::from_secs(timeout_secs), self.tasks.wait()).await;
+        let (sessions_live, streams_live, pending_bytes, pending_items) = {
+            let state = self.state.lock();
+            (
+                state.sessions.len(),
+                state.streams_live,
+                state.pending_bytes,
+                state.pending_items,
+            )
+        };
+        info!(
+            target: "telemt::web",
+            sessions_created = self.sessions_created.load(Ordering::Relaxed),
+            sessions_closed = self.sessions_closed.load(Ordering::Relaxed),
+            sessions_live,
+            streams_opened = self.streams_opened.load(Ordering::Relaxed),
+            streams_rejected = self.streams_rejected.load(Ordering::Relaxed),
+            streams_live,
+            pending_bytes,
+            pending_items,
+            bytes_up = self.bytes_up.load(Ordering::Relaxed),
+            bytes_down = self.bytes_down.load(Ordering::Relaxed),
+            limit_hits = self.limit_hits.load(Ordering::Relaxed),
+            "WEB runtime stopped"
+        );
+    }
+
+    /// Expires credentials and closes idle sessions without holding locks across callbacks.
+    pub(super) fn cleanup(&self) {
+        let generation_id = self.active_runtime.load().id;
+        let now = Instant::now();
+        let sessions = {
+            let mut state = self.state.lock();
+            remove_expired_locked(&mut state, now);
+            let stale_bootstraps = state
+                .bootstraps
+                .iter()
+                .filter_map(|(hash, bootstrap)| {
+                    (bootstrap.generation_id != generation_id && !bootstrap.used)
+                        .then_some(*hash)
+                })
+                .collect::<Vec<_>>();
+            for hash in stale_bootstraps {
+                remove_bootstrap_locked(&mut state, hash);
+            }
+            state.sessions.values().cloned().collect::<Vec<_>>()
+        };
+        for session in sessions.into_iter().filter(|session| session.is_idle(now)) {
+            session.close();
+        }
+    }
+}

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());
+    }
+}