use std::collections::BTreeMap; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Instant; use parking_lot::Mutex; use sha2::{Digest, Sha256}; use crate::config::{ProxyConfig, WebRuntimeVhost, WebWildcardLimitsConfig}; /// Request-owned table; a cold allocation is bounded by the HTTP handler permit. pub(crate) enum CapabilityTable { /// Positive cache hit, including its retained-memory lease. Warm(Arc), /// Complete temporary derivation, not yet eligible for retention. Cold(Box<[[u8; 32]]>), } impl CapabilityTable { /// Returns the dense table used for both containment and full authentication scans. pub(crate) fn capabilities(&self) -> &[[u8; 32]] { match self { Self::Warm(entry) => &entry.capabilities, Self::Cold(capabilities) => capabilities, } } } /// A cache entry retains no runtime generation, template, policy, or decoy snapshot. pub(crate) struct CacheEntry { snapshot: u64, group: usize, sequence: u64, capabilities: Box<[[u8; 32]]>, charged: usize, live_bytes: Arc, } impl Drop for CacheEntry { fn drop(&mut self) { self.live_bytes.fetch_sub(self.charged, Ordering::AcqRel); } } #[derive(Default)] struct State { hosts: BTreeMap>, fifo: BTreeMap, known: BTreeMap<[u8; 32], usize>, next_sequence: u64, last: Option, tokens: f64, #[cfg(test)] evaluations: usize, } /// Process-owned positive cache and atomic HMAC token bucket, independent of quotas. pub(crate) struct WildcardState { limits: WebWildcardLimitsConfig, state: Mutex, live_bytes: Arc, } impl WildcardState { /// Initializes an empty process-scoped cache without allocating hostname inventories. pub(crate) fn new(limits: WebWildcardLimitsConfig) -> Self { Self { limits, state: Mutex::new(State::default()), live_bytes: Arc::default(), } } /// Returns a hit or reserves the entire cold cost before any HMAC evaluation. pub(crate) fn prepare( &self, config: &ProxyConfig, vhost: &WebRuntimeVhost, host: &str, ) -> Option { let group = vhost.wildcard.as_ref()?; let snapshot = config.web.runtime.as_ref()?.snapshot_id; { let mut state = self.state.lock(); if let Some(entry) = state .hosts .get(host) .filter(|entry| entry.snapshot == snapshot && entry.group == group.group_id) { return Some(CapabilityTable::Warm(Arc::clone(entry))); } let count = group.profiles.len(); if count > self.limits.max_profiles_per_request || !reserve(&mut state, Instant::now(), count, &self.limits) { return None; } } // No registry lock spans derivation or immutable authentication lookup. config .wildcard_capabilities(vhost, host) .ok() .map(CapabilityTable::Cold) } /// Tests only positive-cache membership; this never authorizes a profile. pub(crate) fn contains(&self, candidate: &[u8; 32]) -> bool { let digest: [u8; 32] = Sha256::digest(candidate).into(); self.state.lock().known.contains_key(&digest) } /// Retains a complete table only after successful bootstrap issuance. pub(crate) fn publish( &self, snapshot: u64, vhost: &WebRuntimeVhost, host: &str, table: CapabilityTable, is_current: impl FnOnce() -> bool, ) { let CapabilityTable::Cold(capabilities) = table else { return; }; let Some(group) = vhost.wildcard.as_ref() else { return; }; if self.limits.cache_hosts == 0 { return; } // Includes dense bytes, digest-index nodes, two owned host keys, tree nodes, // Arc headers and allocator overhead; detached borrowers keep the full charge. let Some(charged) = capabilities .len() .checked_mul(256) .and_then(|value| value.checked_add(2048 + 2 * host.len())) .filter(|value| *value <= self.limits.cache_bytes) else { return; }; let mut state = self.state.lock(); // Snapshot IDs are unique, not activation-ordered. Recheck under the cache // lock so a late old request cannot replace an already published current entry. if !is_current() { return; } if state .hosts .get(host) .is_some_and(|entry| entry.snapshot == snapshot && entry.group == group.group_id) { return; } if let Some(entry) = state.hosts.remove(host) { detach(&mut state, &entry); } loop { if state.hosts.len() < self.limits.cache_hosts && self.live_bytes.load(Ordering::Acquire) <= self.limits.cache_bytes - charged { break; } let Some((_, oldest)) = state.fifo.first_key_value() else { return; }; let oldest = oldest.clone(); if let Some(entry) = state.hosts.remove(&oldest) { detach(&mut state, &entry); } } let Some(next) = state.next_sequence.checked_add(1) else { return; }; let sequence = state.next_sequence; state.next_sequence = next; self.live_bytes.fetch_add(charged, Ordering::AcqRel); for capability in &capabilities { *state .known .entry(Sha256::digest(capability).into()) .or_default() += 1; } state.fifo.insert(sequence, host.to_string()); state.hosts.insert( host.to_string(), Arc::new(CacheEntry { snapshot, group: group.group_id, sequence, capabilities, charged, live_bytes: Arc::clone(&self.live_bytes), }), ); } } fn detach(state: &mut State, entry: &CacheEntry) { state.fifo.remove(&entry.sequence); for capability in &entry.capabilities { let digest: [u8; 32] = Sha256::digest(capability).into(); if let Some(count) = state.known.get_mut(&digest) { *count -= 1; if *count == 0 { state.known.remove(&digest); } } } // Empty BTreeMaps may retain root allocations. Release them explicitly so // a fully detached cache has no index memory outside live entry leases. if state.hosts.is_empty() { state.hosts = BTreeMap::new(); } if state.fifo.is_empty() { state.fifo = BTreeMap::new(); } if state.known.is_empty() { state.known = BTreeMap::new(); } } fn reserve(state: &mut State, now: Instant, cost: usize, limits: &WebWildcardLimitsConfig) -> bool { let burst = f64::from(limits.cold_hmac_burst); state.tokens = state.last.map_or(burst, |last| { (state.tokens + now.saturating_duration_since(last).as_secs_f64() * f64::from(limits.cold_hmac_per_minute) / 60.0) .min(burst) }); state.last = Some(now); if state.tokens < cost as f64 { return false; } state.tokens -= cost as f64; #[cfg(test)] { state.evaluations += cost; } true } #[cfg(test)] impl WildcardState { /// Exposes deterministic resource evidence without production metrics or hot-path counters. pub(crate) fn test_stats(&self) -> (usize, usize, usize) { let state = self.state.lock(); ( state.hosts.len(), self.live_bytes.load(Ordering::Acquire), state.evaluations, ) } } // Cache accounting and token-bucket invariants are tested without network timing. #[cfg(test)] mod tests;