diff --git a/crates/node/src/handlers.rs b/crates/node/src/handlers.rs index c904a17ae7..1da3d6c1e5 100644 --- a/crates/node/src/handlers.rs +++ b/crates/node/src/handlers.rs @@ -22,6 +22,7 @@ mod specialized_node_invite; pub(crate) mod state_delta; mod stream_opened; pub(crate) mod tee_attestation_admission; +pub(crate) mod tee_attestation_throttle; impl Handler for NodeManager { type Result = (); diff --git a/crates/node/src/handlers/network_event/specialized.rs b/crates/node/src/handlers/network_event/specialized.rs index 82c9e35f7a..271a4ad723 100644 --- a/crates/node/src/handlers/network_event/specialized.rs +++ b/crates/node/src/handlers/network_event/specialized.rs @@ -1,9 +1,12 @@ -use actix::{AsyncContext, WrapFuture}; +use actix::{ActorFutureExt, AsyncContext, WrapFuture}; +use calimero_context_config::types::ContextGroupId; use calimero_network_primitives::messages::NetworkEvent; use calimero_network_primitives::specialized_node_invite::SpecializedNodeType; use calimero_node_primitives::sync::BroadcastMessage; +use sha2::{Digest, Sha256}; use tracing::{debug, error, info, warn}; +use crate::handlers::tee_attestation_throttle::Decision; use crate::handlers::{specialized_node_invite, tee_attestation_admission}; use crate::run::NodeMode; use crate::NodeManager; @@ -143,13 +146,114 @@ pub(super) fn handle_specialized_broadcast( "Received TEE attestation announce on namespace topic" ); + // Admission-control gates (TEE-01 / audit #48). The heavy + // `verify_attestation` path (outbound Intel-PCS fetch + DCAP + // verify) runs BEFORE any policy lookup, so an unguarded announce + // lets a malicious mesh peer amplify a 64 KiB gossip frame into a + // CPU verify + outbound PCS request by replaying a real quote under + // fresh nonces. Guard the spawn synchronously here, on the actor + // thread, before any verify work is scheduled. + // `quote_hash` is the key for both the durable `is_quote_hash_used` + // check and the throttle's dedup gate, so it must be computed before + // either can run — it cannot be deferred behind the rate/inflight + // gates. This is acceptable: `quote_bytes` is bounded to gossipsub's + // 64 KiB default at the transport, and SHA-256 over ≤64 KiB is tens + // of microseconds, negligible beside the PCS-fetch + DCAP verify the + // hash exists to gate. + let quote_hash: [u8; 32] = Sha256::digest(quote_bytes).into(); + let group_id = ContextGroupId::from(namespace_id_bytes); + + // In-memory admission gates FIRST — per-group quote dedup + per-peer + // rate limit + global inflight-verify cap — so a flood is rejected on + // cheap in-memory state before we touch the store below. The returned + // permit must outlive the verify, so it is moved into the spawned + // task and held until completion. Key the throttle on + // `group_id.to_bytes()` so it and the durable check share one + // `ContextGroupId`-derived key (no divergence if the repr changes). + let now = std::time::Instant::now(); + let verify_permit = match this.tee_admission_throttle.check( + now, + source, + group_id.to_bytes(), + quote_hash, + ) { + Decision::Proceed(permit) => permit, + Decision::Duplicate => { + debug!( + %source, + quote_hash = %hex::encode(quote_hash), + "Dropping TeeAttestationAnnounce: recently-seen quote (dedup)" + ); + return true; + } + Decision::RateLimited => { + warn!( + %source, + "Dropping TeeAttestationAnnounce: per-peer attestation rate limit exceeded" + ); + return true; + } + Decision::AtCapacity => { + warn!( + %source, + "Dropping TeeAttestationAnnounce: attestation verify capacity saturated" + ); + return true; + } + }; + + // Durable dedup (a store read), pulled earlier from `admit_tee_node`: + // a quote already admitted to this group never needs re-verifying. + // Run it only AFTER the cheap in-memory gates above, so an + // unauthenticated announce flood can't drive a store read per frame. + // A read error is non-fatal — the authoritative check still runs in + // `admit_tee_node` — so we proceed rather than drop a possibly- + // legitimate announce. On a hit we drop here: `verify_permit` is + // released as it goes out of scope, and the dedup entry the throttle + // just recorded keeps suppressing further re-announces of this quote. + // + // Trade-off of throttle-before-store: an already-admitted quote that + // is re-announced burns one per-peer token before we learn it's + // admitted. That is bounded (the dedup entry suppresses the rest of + // the burst) and rare (a node stops announcing once admitted), and it + // is the price of not doing a store read for every flood frame. We + // deliberately do NOT `forget_quote` here — keeping the entry is what + // makes the follow-up re-announces cheap — nor restore the token, + // which would just re-open the store read on the next frame. + match calimero_governance_store::is_quote_hash_used( + &this.datastore, + &group_id, + "e_hash, + ) { + Ok(true) => { + debug!( + %source, + quote_hash = %hex::encode(quote_hash), + "Dropping TeeAttestationAnnounce: quote already admitted to group" + ); + return true; + } + Ok(false) => {} + Err(err) => { + warn!( + %source, + error = %err, + "Failed to read quote-hash usage; proceeding to verify" + ); + } + } + let context_client = this.clients.context.clone(); let quote_bytes = quote_bytes.clone(); let public_key = *public_key; let nonce = *nonce; + let group_key = group_id.to_bytes(); let _ignored = ctx.spawn( async move { - if let Err(err) = tee_attestation_admission::handle_tee_attestation_announce( + // Hold the inflight permit for the lifetime of the verify so + // the global concurrency cap stays accurate; dropped here. + let _verify_permit = verify_permit; + tee_attestation_admission::handle_tee_attestation_announce( &context_client, source, quote_bytes, @@ -158,15 +262,38 @@ pub(super) fn handle_specialized_broadcast( namespace_id_bytes, ) .await - { + } + .into_actor(this) + .map(move |result, actor, _ctx| { + if let Err(err) = result { warn!( %source, error = %err, "Failed to handle TEE attestation announce" ); + // The verify *errored* (a transient infra failure such + // as an Intel-PCS fetch error — an invalid quote instead + // returns Ok above). Forget the dedup entry so a + // re-announce of the same quote bytes (the fleet-join + // re-announce loop reuses the identical payload) can be + // re-verified rather than being suppressed for the full + // dedup TTL — which, at TTL ≫ the 30s admission window, + // would otherwise sink the whole join on one PCS hiccup. + // A genuinely invalid quote stays deduped, and the + // per-peer rate limit + inflight cap still bound any + // replay-driven re-verification. + // + // This `.map` runs only if the spawned future completes. + // The sole way it is dropped first is the actor stopping + // — but the throttle (and its in-memory dedup map) lives + // on this actor, so it is torn down together; a stale + // entry can't outlive its map to suppress a post-restart + // re-announce (a restarted node starts with an empty map). + actor + .tee_admission_throttle + .forget_quote(group_key, quote_hash); } - } - .into_actor(this), + }), ); true } diff --git a/crates/node/src/handlers/tee_attestation_admission.rs b/crates/node/src/handlers/tee_attestation_admission.rs index 9ebe356c1a..4398d24420 100644 --- a/crates/node/src/handlers/tee_attestation_admission.rs +++ b/crates/node/src/handlers/tee_attestation_admission.rs @@ -31,6 +31,20 @@ pub(crate) fn public_key_binding_hash(public_key: &PublicKey) -> [u8; 32] { /// /// Verifies the TDX quote, checks measurements against the group's TEE admission /// policy, and publishes a `MemberJoinedViaTeeAttestation` governance op if valid. +/// +/// # Error semantics (load-bearing) +/// +/// Returns `Ok(())` for both a successful admission **and** a definitively +/// rejected quote (verification ran and returned an invalid result). Returns +/// `Err` only for a *transient* failure where the verify could not be completed +/// — most importantly the outbound Intel-PCS collateral fetch erroring, but also +/// store/governance publish failures. The caller in the +/// `network_event::specialized` handler relies on this distinction: +/// it clears the throttle's in-memory dedup entry on `Err` (so a legitimate +/// re-announce can be re-verified after a transient outage) but keeps it on +/// `Ok` (so a replayed/invalid quote stays suppressed). Do not start returning +/// `Err` for an invalid-but-successfully-evaluated quote without updating that +/// caller. pub async fn handle_tee_attestation_announce( context_client: &calimero_context_client::client::ContextClient, source: libp2p::PeerId, diff --git a/crates/node/src/handlers/tee_attestation_throttle.rs b/crates/node/src/handlers/tee_attestation_throttle.rs new file mode 100644 index 0000000000..f7f4e72e62 --- /dev/null +++ b/crates/node/src/handlers/tee_attestation_throttle.rs @@ -0,0 +1,699 @@ +//! Admission-control throttle for the gossip `TeeAttestationAnnounce` path. +//! +//! Inbound `TeeAttestationAnnounce` broadcasts drive +//! [`tee_attestation_admission::handle_tee_attestation_announce`], which runs +//! the heavy [`calimero_tee_attestation::verify_attestation`] path — an +//! outbound Intel-PCS collateral fetch plus DCAP crypto-verify — *before* any +//! policy lookup. Without a guard, a malicious mesh peer on a TEE namespace +//! topic can replay a structurally-valid quote (varying the announce nonce to +//! beat gossipsub's message-id dedup) and amplify each cheap gossip frame into +//! a CPU verify + an outbound PCS request (TEE-01 / audit #48). +//! +//! This throttle is consulted *synchronously* on the `NodeManager` actor +//! thread before the verify task is spawned. It composes three independent +//! gates, each of which can reject an announce on its own: +//! +//! 1. **Per-group quote dedup** — a recently-seen `(group, quote_hash)` is +//! dropped, so replaying one captured quote (identical `quote_bytes` ⇒ +//! identical hash) under many nonces costs at most one verify per TTL +//! window. This complements the durable governance-store check +//! (`is_quote_hash_used`, which only knows *admitted* quotes) by also +//! covering not-yet-admitted replays. +//! 2. **Per-peer rate limit** — a lazily-refilled token bucket per source +//! peer bounds how fast any single peer can drive verifies. +//! 3. **Global inflight cap** — a bounded semaphore caps the number of +//! concurrent verifies across all peers/groups; the returned permit is +//! held for the lifetime of the spawned verify task. +//! +//! The struct is touched only on the actor thread, so the bookkeeping maps +//! need no locking; only the inflight [`Semaphore`] is shared (it is moved, +//! via an owned permit, into the spawned task). +//! +//! Per-peer state is created only when a peer's announce actually *proceeds* +//! (Gate 3 grant), never on a rejection. A consequence is that a peer whose +//! announces are *always* refused at the global inflight cap accumulates no +//! per-peer rate-limit debt — but that is harmless by construction: such a peer +//! triggers zero verifies (the global cap is doing the limiting), and the +//! moment one of its announces proceeds it is tracked and metered normally. +//! Tracking never-proceeding peers would instead let a unique-peer flood grow +//! the map for no protective gain. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use libp2p::PeerId; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +/// Maximum number of attestation verifies allowed to run concurrently across +/// all peers and groups. The verify path makes a blocking-ish outbound PCS +/// fetch, so this caps both CPU and outbound amplification fan-out. +pub const DEFAULT_MAX_INFLIGHT_VERIFIES: usize = 4; + +/// Per-peer token-bucket burst: a single peer may trigger this many verifies +/// back-to-back before being throttled to the refill rate. +pub const DEFAULT_PER_PEER_BURST: f64 = 5.0; + +/// Per-peer token-bucket refill interval — one token is restored per this +/// duration, up to [`DEFAULT_PER_PEER_BURST`]. At the default (1 token / 2s) +/// a saturating peer is held to ~30 verifies/min. +pub const DEFAULT_PER_PEER_REFILL: Duration = Duration::from_secs(2); + +/// How long a `(group, quote_hash)` is remembered for dedup. A replay of the +/// same quote within this window is dropped without a verify. +pub const DEFAULT_DEDUP_TTL: Duration = Duration::from_secs(300); + +/// Hard cap on the number of distinct peers and dedup entries retained, so a +/// flood of unique peers / quotes can't grow the maps without bound. When the +/// cap is hit, the oldest entries are pruned first. +const MAX_TRACKED_PEERS: usize = 4096; +const MAX_TRACKED_QUOTES: usize = 8192; + +/// Minimum spacing between time-based [`TeeAdmissionThrottle::prune`] sweeps. +/// `prune`'s `retain` pass is O(tracked entries); running it on every announce +/// would make the per-call cost grow with the map size under an adversarial +/// flood — the exact case this throttle defends against. Instead we sweep at +/// most once per this interval (a size-cap guard still forces an immediate +/// sweep if either map exceeds its hard cap, so memory stays bounded). +const PRUNE_INTERVAL: Duration = Duration::from_secs(1); + +/// Outcome of consulting the throttle for one announce. +#[derive(Debug)] +pub enum Decision { + /// Proceed with the verify; hold `permit` until the verify completes so + /// the global inflight cap stays accurate. + Proceed(OwnedSemaphorePermit), + /// The same `(group, quote_hash)` was seen recently — dropped. + Duplicate, + /// The source peer exceeded its per-peer rate limit — dropped. + RateLimited, + /// The global inflight-verify cap is saturated — dropped. + AtCapacity, +} + +struct PeerBucket { + tokens: f64, + /// Time of the last *granted* token. Drives the refill clock and is + /// committed only on `Decision::Proceed` (see `check`), so a rejection + /// never advances it. + last: Instant, + /// Time this peer was last *seen* (any `check` outcome, including + /// rejections). Used as the idle/LRU key for pruning so an active-but- + /// throttled peer is retained — evicting it would hand it a fresh full + /// bucket and reset its rate limit. + last_seen: Instant, +} + +/// Admission-control throttle. See the module docs for the gate ordering and +/// rationale. Construct one per node and consult it on the actor thread. +pub struct TeeAdmissionThrottle { + inflight: Arc, + peers: HashMap, + recent_quotes: HashMap<([u8; 32], [u8; 32]), Instant>, + per_peer_burst: f64, + per_peer_refill: Duration, + dedup_ttl: Duration, + /// `now` of the last time-based prune sweep, used to amortise `prune` to at + /// most once per [`PRUNE_INTERVAL`]. `None` until the first `check`. + last_prune: Option, +} + +impl std::fmt::Debug for TeeAdmissionThrottle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TeeAdmissionThrottle") + .field("available_permits", &self.inflight.available_permits()) + .field("tracked_peers", &self.peers.len()) + .field("tracked_quotes", &self.recent_quotes.len()) + .finish() + } +} + +impl Default for TeeAdmissionThrottle { + fn default() -> Self { + Self::new( + DEFAULT_MAX_INFLIGHT_VERIFIES, + DEFAULT_PER_PEER_BURST, + DEFAULT_PER_PEER_REFILL, + DEFAULT_DEDUP_TTL, + ) + } +} + +impl TeeAdmissionThrottle { + /// Construct a throttle with explicit gate parameters. + /// + /// # Panics + /// + /// Panics if `max_inflight == 0`, `per_peer_burst` is not in + /// `[1.0, u32::MAX]`, or `per_peer_refill == 0`: with no inflight permits no + /// announce could ever proceed, a sub-unit burst can never satisfy the + /// `tokens >= 1.0` gate, a burst beyond `u32::MAX` would overflow the + /// `burst * refill` idle-eviction window cast, and a zero refill interval + /// would make the refill rate non-finite (poisoning the lazy-refill + /// arithmetic with `NaN`) — so each renders the throttle useless or + /// ill-defined. These are construction-time programmer errors — the only + /// in-tree callers are [`Default`] and tests, both of which pass valid + /// constants — so they are asserted rather than surfaced as a runtime + /// `Result` the caller would have to thread through node startup. + pub fn new( + max_inflight: usize, + per_peer_burst: f64, + per_peer_refill: Duration, + dedup_ttl: Duration, + ) -> Self { + assert!(max_inflight > 0, "max_inflight must be positive"); + assert!( + per_peer_burst >= 1.0 && per_peer_burst <= u32::MAX as f64, + "per_peer_burst must be in [1, u32::MAX]" + ); + assert!( + per_peer_refill > Duration::ZERO, + "per_peer_refill must be positive" + ); + Self { + inflight: Arc::new(Semaphore::new(max_inflight)), + peers: HashMap::new(), + recent_quotes: HashMap::new(), + per_peer_burst, + per_peer_refill, + dedup_ttl, + last_prune: None, + } + } + + /// Tokens restored per second. `per_peer_refill` is asserted `> 0` in + /// [`Self::new`], so this is always finite (no `1.0 / 0.0` / `INFINITY`, + /// which would poison the `0 * rate` refill term with `NaN`). + fn refill_per_sec(&self) -> f64 { + 1.0 / self.per_peer_refill.as_secs_f64() + } + + /// Consult all three gates for an announce observed at `now`. + /// + /// On `Decision::Proceed` the `(group, quote_hash)` is recorded for dedup + /// and one per-peer token + one inflight permit are consumed. On any + /// rejection no token is consumed and nothing is recorded, so a legitimate + /// retry after the rate-limit/capacity pressure clears is not penalised. + pub fn check( + &mut self, + now: Instant, + source: PeerId, + group_id: [u8; 32], + quote_hash: [u8; 32], + ) -> Decision { + // Prune is O(tracked entries); amortise it to at most once per + // `PRUNE_INTERVAL` so an adversarial flood doesn't pay an O(N) sweep on + // every announce before the cheap gates below can reject it. A size-cap + // guard still forces an immediate sweep whenever either map is over its + // hard cap, so memory stays bounded regardless of call cadence. + let over_cap = + self.recent_quotes.len() > MAX_TRACKED_QUOTES || self.peers.len() > MAX_TRACKED_PEERS; + let due = self + .last_prune + .is_none_or(|last| now.saturating_duration_since(last) >= PRUNE_INTERVAL); + if over_cap || due { + self.prune(now); + self.last_prune = Some(now); + } + + // Gate 1: per-group quote dedup. Cheapest, and the most effective + // guard against single-quote replay floods. + let key = (group_id, quote_hash); + if let Some(seen) = self.recent_quotes.get(&key) { + if now.saturating_duration_since(*seen) < self.dedup_ttl { + return Decision::Duplicate; + } + } + + // Gate 2: per-peer rate limit. Compute the lazily-refilled token count + // but do NOT write it back yet — the bucket is mutated transactionally, + // only on `Proceed` (see the commit block below). A rejected announce + // (here or at Gate 3) therefore leaves `tokens` and `last` untouched, so + // it neither burns a token nor advances the peer's refill clock. + let burst = self.per_peer_burst; + let refill_per_sec = self.refill_per_sec(); + // Read the current bucket state *without* inserting. A brand-new peer + // is treated as a full bucket, so first contact is never rate-limited; + // crucially, a peer is only inserted on the commit (Proceed) path below, + // so a flood of unique source peers that never proceed can't grow + // `peers`. For an already-tracked peer, refresh `last_seen` on every + // call (any outcome) so an active-but-throttled peer isn't LRU-evicted. + let (cur_tokens, cur_last) = match self.peers.get_mut(&source) { + Some(bucket) => { + bucket.last_seen = now; + (bucket.tokens, bucket.last) + } + None => (burst, now), + }; + let elapsed = now.saturating_duration_since(cur_last).as_secs_f64(); + let refilled = (cur_tokens + elapsed * refill_per_sec).min(burst); + if refilled < 1.0 { + return Decision::RateLimited; + } + + // Gate 3: global inflight cap. Acquire last so a rejection here does + // not burn a per-peer token. No bucket has been inserted/mutated for a + // refused-here peer, so an `AtCapacity` return advances nothing. + // + // Intentional trade-off: a brand-new peer refused here is never + // inserted, so it retains a full bucket and can burst at the full rate + // the moment capacity opens, regardless of how many times it was refused. + // The global inflight cap — not per-peer fairness — is the binding + // constraint while saturated, and not tracking never-proceeding peers is + // what stops a unique-peer flood from growing `peers` (see module docs). + let permit = match Arc::clone(&self.inflight).try_acquire_owned() { + Ok(permit) => permit, + Err(_) => return Decision::AtCapacity, + }; + + // All gates passed: commit the side effects atomically. This is the + // only path that inserts/advances the bucket, so `tokens`/`last` reflect + // the previous *grant* plus this one, never an intervening rejection. + // The `or_insert_with` value is a placeholder used only for a genuinely + // new peer; the three assignments below always run and set the real + // post-grant state for both new and existing peers. + let bucket = self.peers.entry(source).or_insert_with(|| PeerBucket { + tokens: burst, + last: now, + last_seen: now, + }); + bucket.tokens = refilled - 1.0; + bucket.last = now; + bucket.last_seen = now; + let _ = self.recent_quotes.insert(key, now); + Decision::Proceed(permit) + } + + /// Forget a recorded `(group, quote_hash)` dedup entry so the next announce + /// for that quote is re-verified instead of being suppressed for the full + /// dedup TTL. + /// + /// Called after a verify that *errored* (a transient infrastructure failure + /// — e.g. the Intel-PCS collateral fetch failed), so a legitimate re-announce + /// of the identical quote (the fleet-join re-announce loop reuses the same + /// `quote_bytes`) can recover within its admission window rather than being + /// stuck until the TTL elapses. A quote that *failed verification* (a + /// definite invalid result, not an error) is deliberately left recorded, so + /// a replay of it stays suppressed. DoS remains bounded by the per-peer rate + /// limit and the global inflight cap either way. + pub fn forget_quote(&mut self, group_id: [u8; 32], quote_hash: [u8; 32]) { + let _ = self.recent_quotes.remove(&(group_id, quote_hash)); + } + + /// Drop expired dedup entries and full/idle peer buckets, and hard-cap map + /// sizes so adversarial churn can't grow memory without bound. + fn prune(&mut self, now: Instant) { + let dedup_ttl = self.dedup_ttl; + self.recent_quotes + .retain(|_, seen| now.saturating_duration_since(*seen) < dedup_ttl); + // Hard cap. Reaching it requires >MAX_TRACKED_QUOTES *distinct* quotes + // admitted within the TTL — and entries are only inserted on the commit + // (Proceed) path, gated by the per-peer rate limit and the global + // inflight cap, so a rejected flood can't grow this map at all. If the + // cap is somehow hit, evicting the oldest entries by *insertion time* + // (which, since the dedup TTL is uniform, are exactly those nearest to + // expiry) degrades dedup to best-effort for those quotes (a replay could + // trigger one more verify) — an intentional trade-off: the memory bound + // is the hard guarantee, and the rate-limit + inflight gates (plus the + // durable `is_quote_hash_used` check for admitted quotes) remain the + // real DoS backstop. Do not "fix" this by removing the cap. + if self.recent_quotes.len() > MAX_TRACKED_QUOTES { + Self::evict_oldest(&mut self.recent_quotes, MAX_TRACKED_QUOTES); + } + + // Drop a peer once it hasn't been *seen* for at least the time its + // bucket needs to fully refill from empty (`burst * per_peer_refill`). + // By then it is indistinguishable from a fresh peer, so its state is + // worth nothing — and this is exactly the window after which a fresh + // full bucket grants no more than the configured rate, so eviction can't + // become a rate-limit-reset bypass (a peer flooding faster is *seen* + // more recently and so retained). Keying on `last_seen` (updated every + // call, any outcome) rather than `last` (grant time) is what keeps an + // active-but-throttled peer; and because `last_seen >= last` always, a + // peer idle this long has provably refilled, so no separate fullness + // check is needed. + let full_refill = self + .per_peer_refill + .saturating_mul(self.per_peer_burst.ceil() as u32); + self.peers + .retain(|_, b| now.saturating_duration_since(b.last_seen) < full_refill); + if self.peers.len() > MAX_TRACKED_PEERS { + Self::evict_oldest_peers(&mut self.peers, MAX_TRACKED_PEERS); + } + } + + fn evict_oldest(map: &mut HashMap, keep: usize) { + let mut entries: Vec<(K, Instant)> = map.iter().map(|(k, v)| (k.clone(), *v)).collect(); + Self::evict_oldest_entries(map, &mut entries, keep); + } + + fn evict_oldest_peers(map: &mut HashMap, keep: usize) { + // Key on `last_seen` (last activity), not `last` (last grant), so the + // peers dropped under cap pressure are the genuinely-quiet ones rather + // than active-but-throttled peers (whose `last` is stale by design). + let mut entries: Vec<(PeerId, Instant)> = + map.iter().map(|(k, v)| (*k, v.last_seen)).collect(); + Self::evict_oldest_entries(map, &mut entries, keep); + } + + /// Remove the oldest entries from `map` until at most `keep` remain. `entries` + /// MUST be a full `(key, timestamp)` snapshot of `map` — both callers collect + /// every entry immediately before calling — so `remove` (derived from + /// `map.len()`) indexes the snapshot correctly. Uses `select_nth_unstable` + /// (O(n) average) rather than a full O(n log n) sort, since the cap is only + /// ever crossed under adversarial churn and the evicted entries are + /// discarded, so their relative order is irrelevant. + fn evict_oldest_entries( + map: &mut HashMap, + entries: &mut [(K, Instant)], + keep: usize, + ) { + debug_assert_eq!( + entries.len(), + map.len(), + "evict_oldest_entries requires a full snapshot of `map`" + ); + let remove = map.len().saturating_sub(keep); + if remove == 0 { + return; + } + let _ = entries.select_nth_unstable_by_key(remove - 1, |(_, t)| *t); + for (k, _) in &entries[..remove] { + let _ = map.remove(k); + } + // Post-condition: the cap is actually enforced. Catches a future caller + // that passes a stale snapshot (some `entries[..remove]` key already + // gone from `map`, so `remove` deletions don't bring it under `keep`). + debug_assert!( + map.len() <= keep, + "evict_oldest_entries left the map over its cap (stale snapshot?)" + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn peer(n: u8) -> PeerId { + // Deterministic distinct peers for tests. + let kp = libp2p::identity::Keypair::ed25519_from_bytes([n; 32]).expect("valid key"); + kp.public().to_peer_id() + } + + #[test] + fn proceeds_then_dedups_same_group_quote() { + let mut t = TeeAdmissionThrottle::default(); + let now = Instant::now(); + let g = [1u8; 32]; + let q = [2u8; 32]; + + // First announce verifies. + assert!(matches!(t.check(now, peer(1), g, q), Decision::Proceed(_))); + // Exact replay (same group+quote) is deduped, even from a different + // peer and a different time within the TTL. + assert!(matches!(t.check(now, peer(2), g, q), Decision::Duplicate)); + assert!(matches!( + t.check(now + Duration::from_secs(10), peer(1), g, q), + Decision::Duplicate + )); + } + + #[test] + fn dedup_is_per_group() { + let mut t = TeeAdmissionThrottle::default(); + let now = Instant::now(); + let q = [9u8; 32]; + assert!(matches!( + t.check(now, peer(1), [1u8; 32], q), + Decision::Proceed(_) + )); + // Same quote, different group: not a duplicate. + assert!(matches!( + t.check(now, peer(1), [2u8; 32], q), + Decision::Proceed(_) + )); + } + + #[test] + fn dedup_expires_after_ttl() { + let mut t = + TeeAdmissionThrottle::new(8, 100.0, Duration::from_secs(1), Duration::from_secs(60)); + let now = Instant::now(); + let g = [1u8; 32]; + let q = [2u8; 32]; + assert!(matches!(t.check(now, peer(1), g, q), Decision::Proceed(_))); + assert!(matches!( + t.check(now + Duration::from_secs(30), peer(1), g, q), + Decision::Duplicate + )); + // After the TTL the same quote is allowed through again. + assert!(matches!( + t.check(now + Duration::from_secs(61), peer(1), g, q), + Decision::Proceed(_) + )); + } + + #[test] + fn per_peer_rate_limit_bursts_then_throttles() { + // Burst 3, refill 1 token/sec. Big inflight cap so capacity isn't the + // gate. Each announce uses a distinct quote so dedup isn't the gate. + let mut t = + TeeAdmissionThrottle::new(1000, 3.0, Duration::from_secs(1), Duration::from_secs(300)); + let now = Instant::now(); + let g = [1u8; 32]; + let p = peer(1); + + for i in 0..3u8 { + assert!( + matches!(t.check(now, p, g, [i; 32]), Decision::Proceed(_)), + "burst token {i} should pass" + ); + } + // 4th within the same instant: bucket empty. + assert!(matches!( + t.check(now, p, g, [100u8; 32]), + Decision::RateLimited + )); + + // A different peer is unaffected (per-peer bucket). + assert!(matches!( + t.check(now, peer(2), g, [101u8; 32]), + Decision::Proceed(_) + )); + + // After 1s, one token refills. + assert!(matches!( + t.check(now + Duration::from_secs(1), p, g, [102u8; 32]), + Decision::Proceed(_) + )); + assert!(matches!( + t.check(now + Duration::from_secs(1), p, g, [103u8; 32]), + Decision::RateLimited + )); + } + + #[test] + fn rate_limited_announce_does_not_consume_token_or_dedup() { + let mut t = TeeAdmissionThrottle::new( + 1000, + 1.0, + Duration::from_secs(1000), + Duration::from_secs(300), + ); + let now = Instant::now(); + let g = [1u8; 32]; + let p = peer(1); + // Use the single burst token. + assert!(matches!( + t.check(now, p, g, [1u8; 32]), + Decision::Proceed(_) + )); + // Now rate-limited for quote [2;32]. + assert!(matches!( + t.check(now, p, g, [2u8; 32]), + Decision::RateLimited + )); + // That rejected quote was NOT recorded for dedup: once a token is + // available again it can proceed (proving no dedup side effect). + assert!(matches!( + t.check(now + Duration::from_secs(1000), p, g, [2u8; 32]), + Decision::Proceed(_) + )); + } + + #[test] + fn rejections_do_not_slow_token_recovery() { + // Regression guard: each `check` advances `bucket.last = now` *after* + // crediting the elapsed refill, so a rejection between two accepts must + // not steal recovery credit. Burst 1, refill 1 token/sec. + let mut t = + TeeAdmissionThrottle::new(1000, 1.0, Duration::from_secs(1), Duration::from_secs(300)); + let t0 = Instant::now(); + let g = [1u8; 32]; + let p = peer(1); + + // Spend the single token. + assert!(matches!(t.check(t0, p, g, [0u8; 32]), Decision::Proceed(_))); + // Half a second in: only 0.5 token refilled → rejected. The bucket is + // committed only on Proceed, so this rejection does NOT advance `last`; + // the full 1s of elapsed time is therefore credited at t0+1s. + assert!(matches!( + t.check(t0 + Duration::from_millis(500), p, g, [1u8; 32]), + Decision::RateLimited + )); + // At exactly t0+1s the bucket must be back to a full token despite the + // intervening rejection — recovery tracks wall-clock, not call count. + assert!(matches!( + t.check(t0 + Duration::from_secs(1), p, g, [2u8; 32]), + Decision::Proceed(_) + )); + } + + #[test] + fn refused_new_peer_is_not_tracked() { + // Inflight cap 1; hold the only permit so a second, brand-new peer can + // only ever hit AtCapacity. Generous burst/refill so Gate 2 never bites. + let mut t = TeeAdmissionThrottle::new( + 1, + 1000.0, + Duration::from_millis(1), + Duration::from_secs(300), + ); + let now = Instant::now(); + let g = [1u8; 32]; + let _held = match t.check(now, peer(1), g, [1u8; 32]) { + Decision::Proceed(p) => p, + other => panic!("expected Proceed, got {other:?}"), + }; + assert_eq!(t.peers.len(), 1, "the proceeding peer is tracked"); + // A new peer that is refused at Gate 3 must not be inserted, so a flood + // of unique never-proceeding peers can't grow the map. + assert!(matches!( + t.check(now, peer(2), g, [2u8; 32]), + Decision::AtCapacity + )); + assert_eq!(t.peers.len(), 1, "a refused new peer must not grow the map"); + } + + #[test] + fn idle_prune_drops_long_idle_peer_keeps_recently_seen() { + // burst 5, refill 1 token/sec → full-refill window = 5 * 1s = 5s. + let mut t = + TeeAdmissionThrottle::new(1000, 5.0, Duration::from_secs(1), Duration::from_secs(300)); + let now = Instant::now(); + let g = [0u8; 32]; + + // Peer 1: seen once at t0, then never again. + assert!(matches!( + t.check(now, peer(1), g, [1u8; 32]), + Decision::Proceed(_) + )); + // Peer 2: seen at t0, and again just before the prune so `last_seen` is + // recent (even though its bucket has also refilled by wall-clock). + assert!(matches!( + t.check(now, peer(2), g, [2u8; 32]), + Decision::Proceed(_) + )); + let _ = t.check(now + Duration::from_millis(4500), peer(2), g, [3u8; 32]); + + // Prune at t0+6s (> the 5s full-refill window). Peer 1 has been idle 6s + // (≥ window) → pruned; peer 2 was seen 1.5s ago (< window) → retained. + t.prune(now + Duration::from_secs(6)); + assert!( + !t.peers.contains_key(&peer(1)), + "peer idle past the full-refill window should be pruned" + ); + assert!( + t.peers.contains_key(&peer(2)), + "recently-seen peer must be retained regardless of bucket level" + ); + } + + #[test] + fn forget_quote_allows_reverify_after_transient_failure() { + // Generous burst/inflight so only the dedup gate is in play. + let mut t = TeeAdmissionThrottle::new( + 1000, + 100.0, + Duration::from_secs(1), + Duration::from_secs(300), + ); + let now = Instant::now(); + let g = [1u8; 32]; + let q = [2u8; 32]; + + // First announce proceeds and records the dedup entry. + assert!(matches!(t.check(now, peer(1), g, q), Decision::Proceed(_))); + // A re-announce of the same quote within the TTL is deduped. + assert!(matches!(t.check(now, peer(1), g, q), Decision::Duplicate)); + // The verify errored transiently → forget the entry. + t.forget_quote(g, q); + // The same quote may now be re-verified (retry resilience restored). + assert!(matches!(t.check(now, peer(1), g, q), Decision::Proceed(_))); + } + + #[test] + fn global_inflight_cap_blocks_when_saturated() { + // Cap of 2 concurrent verifies; generous per-peer + dedup so this is + // the only gate. Hold the returned permits to simulate in-flight work. + let mut t = TeeAdmissionThrottle::new( + 2, + 1000.0, + Duration::from_millis(1), + Duration::from_secs(300), + ); + let now = Instant::now(); + let g = [1u8; 32]; + + let p1 = match t.check(now, peer(1), g, [1u8; 32]) { + Decision::Proceed(p) => p, + other => panic!("expected Proceed, got {other:?}"), + }; + let _p2 = match t.check(now, peer(2), g, [2u8; 32]) { + Decision::Proceed(p) => p, + other => panic!("expected Proceed, got {other:?}"), + }; + // Both permits held → at capacity. + assert!(matches!( + t.check(now, peer(3), g, [3u8; 32]), + Decision::AtCapacity + )); + + // Completing one verify (drop its permit) frees a slot. + drop(p1); + assert!(matches!( + t.check(now, peer(4), g, [4u8; 32]), + Decision::Proceed(_) + )); + } + + #[test] + fn at_capacity_does_not_consume_token_or_dedup() { + let mut t = TeeAdmissionThrottle::new( + 1, + 1000.0, + Duration::from_millis(1), + Duration::from_secs(300), + ); + let now = Instant::now(); + let g = [1u8; 32]; + let p = peer(1); + let held = match t.check(now, p, g, [7u8; 32]) { + Decision::Proceed(perm) => perm, + other => panic!("expected Proceed, got {other:?}"), + }; + // Saturated. + assert!(matches!( + t.check(now, p, g, [8u8; 32]), + Decision::AtCapacity + )); + drop(held); + // The capacity-rejected quote [8;32] was not recorded for dedup, and + // the token was not burned: it proceeds once a slot is free. + assert!(matches!( + t.check(now, p, g, [8u8; 32]), + Decision::Proceed(_) + )); + } +} diff --git a/crates/node/src/manager.rs b/crates/node/src/manager.rs index 051159fedd..837ea98e0b 100644 --- a/crates/node/src/manager.rs +++ b/crates/node/src/manager.rs @@ -128,6 +128,17 @@ pub struct NodeManager { /// local residue, and the emitter publishes the node's own signed /// heartbeat (on-change + periodic) on the namespace topic. pub(crate) migration_emitter_addr: Option>, + /// Admission-control throttle for the inbound gossip + /// `TeeAttestationAnnounce` path (TEE-01 / audit #48). Consulted + /// synchronously on this actor thread before + /// `tee_attestation_admission::handle_tee_attestation_announce` is spawned, + /// so a malicious mesh peer cannot drive the heavy + /// `verify_attestation` (outbound Intel-PCS fetch + DCAP verify) path + /// unbounded by replaying structurally-valid quotes. Combines per-group + /// quote dedup, a per-peer rate limit, and a global inflight-verify cap. + /// See [`crate::handlers::tee_attestation_throttle`]. + pub(crate) tee_admission_throttle: + crate::handlers::tee_attestation_throttle::TeeAdmissionThrottle, } impl NodeManager { @@ -164,6 +175,8 @@ impl NodeManager { ns_beacon_sync_debounce: Arc::new(Mutex::new(HashMap::new())), migration_status_cache: Arc::new(MigrationStatusCache::default()), migration_emitter_addr: None, + tee_admission_throttle: + crate::handlers::tee_attestation_throttle::TeeAdmissionThrottle::default(), } } } diff --git a/crates/server/src/admin/handlers/tee.rs b/crates/server/src/admin/handlers/tee.rs index 3dd976d30f..b69126a004 100644 --- a/crates/server/src/admin/handlers/tee.rs +++ b/crates/server/src/admin/handlers/tee.rs @@ -5,6 +5,7 @@ mod attest; pub mod fleet_join; mod info; mod verify_quote; +mod verify_quote_throttle; pub fn service() -> Router { Router::new() diff --git a/crates/server/src/admin/handlers/tee/verify_quote.rs b/crates/server/src/admin/handlers/tee/verify_quote.rs index 3b2f463956..836c9de1e7 100644 --- a/crates/server/src/admin/handlers/tee/verify_quote.rs +++ b/crates/server/src/admin/handlers/tee/verify_quote.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use axum::response::IntoResponse; use axum::Extension; @@ -8,12 +8,27 @@ use calimero_server_primitives::admin::{ }; use calimero_tee_attestation::{verify_attestation, AttestationError}; use reqwest::StatusCode; -use tracing::{error, info}; +use tracing::{error, info, warn}; +use super::verify_quote_throttle::{Decision, VerifyQuoteThrottle}; use crate::admin::handlers::validation::ValidatedJson; use crate::admin::service::{ApiError, ApiResponse}; use crate::AdminState; +/// Process-global rate + concurrency limiter for the public, unauthenticated +/// `/verify-quote` endpoint (TEE-01 / audit #325). The endpoint must stay +/// public (the mdma manager proxies to it with no node admin token), so this +/// is the boundary that keeps an attacker from driving the heavy +/// `verify_attestation` (outbound Intel-PCS fetch + DCAP verify) path +/// unbounded. +static VERIFY_QUOTE_THROTTLE: LazyLock = + LazyLock::new(VerifyQuoteThrottle::default); + +/// Generic message returned for any internal (500) verification failure. The +/// real error is logged server-side; the response body never leaks internals +/// (PCS endpoint shape, cert-chain/TCB parser detail, etc.). +const GENERIC_INTERNAL_ERROR: &str = "Attestation verification failed"; + pub async fn handler( Extension(_state): Extension>, ValidatedJson(req): ValidatedJson, @@ -23,6 +38,22 @@ pub async fn handler( "Verifying TDX quote" ); + // Throttle BEFORE any decode/verify work so a flood is rejected cheaply. + // The permit is held for the lifetime of the verify so the global inflight + // cap stays accurate. + let _permit = match VERIFY_QUOTE_THROTTLE.check() { + Decision::Proceed(permit) => permit, + Decision::RateLimited | Decision::AtCapacity => { + warn!("Rejecting /verify-quote: attestation verification throttle exceeded"); + return ApiError { + status_code: StatusCode::TOO_MANY_REQUESTS, + message: "Too many attestation verification requests; please retry later" + .to_owned(), + } + .into_response(); + } + }; + match verify_quote(req).await { Ok(response) => ApiResponse { payload: response }.into_response(), Err(err) => err.into_response(), @@ -74,7 +105,7 @@ async fn verify_quote(req: TeeVerifyQuoteRequest) -> Result Result { - (StatusCode::BAD_REQUEST, err.to_string()) - } - AttestationError::CollateralFetchFailed(_) => { - (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - } - _ => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()), - }; + // Always log the real error server-side; never put it in the + // response body (TEE-01 / audit #325). See `scrub_verify_error`. error!(error=%err, "Attestation verification failed"); - ApiError { - status_code, - message, - } + scrub_verify_error(&err) })?; let is_valid = result.is_valid(); @@ -117,3 +138,72 @@ async fn verify_quote(req: TeeVerifyQuoteRequest) -> Result ApiError { + match err { + AttestationError::QuoteParsingFailed(_) => ApiError { + status_code: StatusCode::BAD_REQUEST, + message: "Invalid TDX quote".to_owned(), + }, + _ => ApiError { + status_code: StatusCode::INTERNAL_SERVER_ERROR, + message: GENERIC_INTERNAL_ERROR.to_owned(), + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A malformed-quote error is a 400 with a fixed message that does NOT + /// echo the underlying parser detail. + #[test] + fn quote_parsing_error_is_generic_400() { + let detail = "frobnicated header at offset 0x41414141"; + let api = scrub_verify_error(&AttestationError::QuoteParsingFailed(detail.to_owned())); + assert_eq!(api.status_code, StatusCode::BAD_REQUEST); + assert_eq!(api.message, "Invalid TDX quote"); + assert!( + !api.message.contains(detail), + "400 body must not leak parser detail" + ); + } + + /// The PCS collateral-fetch failure is the headline leak vector: it must + /// be a scrubbed 500 that never echoes the PCS URL / connection detail. + #[test] + fn collateral_fetch_error_is_scrubbed_500() { + let detail = "https://api.trustedservices.intel.com/... connection refused"; + let api = scrub_verify_error(&AttestationError::CollateralFetchFailed(detail.to_owned())); + assert_eq!(api.status_code, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(api.message, GENERIC_INTERNAL_ERROR); + assert!( + !api.message.contains("intel") && !api.message.contains(detail), + "500 body must not leak the PCS endpoint / fetch detail" + ); + } + + /// Any other internal error (e.g. signature verify) is also a scrubbed 500. + #[test] + fn other_internal_errors_are_scrubbed_500() { + for err in [ + AttestationError::QuoteVerificationFailed("cert chain X leaf Y".to_owned()), + AttestationError::QuoteConversionFailed("secret".to_owned()), + AttestationError::SystemTimeError("clock".to_owned()), + AttestationError::NotSupported, + ] { + let api = scrub_verify_error(&err); + assert_eq!(api.status_code, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(api.message, GENERIC_INTERNAL_ERROR); + } + } +} diff --git a/crates/server/src/admin/handlers/tee/verify_quote_throttle.rs b/crates/server/src/admin/handlers/tee/verify_quote_throttle.rs new file mode 100644 index 0000000000..877134245d --- /dev/null +++ b/crates/server/src/admin/handlers/tee/verify_quote_throttle.rs @@ -0,0 +1,249 @@ +//! Process-global throttle for the public `POST /tee/verify-quote` endpoint. +//! +//! `/verify-quote` is mounted unauthenticated (the mdma manager proxies to it +//! without a node admin token), and it drives the heavy +//! [`calimero_tee_attestation::verify_attestation`] path — an outbound +//! Intel-PCS collateral fetch + DCAP crypto-verify — on an attacker-supplied +//! quote (TEE-01 / audit #325). Since it must stay public, throttle it: a +//! global token bucket bounds the request rate and a bounded semaphore caps +//! concurrent verifies, so an unauthenticated caller cannot turn the endpoint +//! into a CPU-DoS / outbound-PCS amplifier. +//! +//! The limit is intentionally *global* (not per-client): the endpoint has no +//! authenticated identity to key on, and the amplification concern is the +//! aggregate outbound/CPU fan-out, not fairness between callers. + +use std::sync::Arc; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +/// Maximum concurrent `/verify-quote` verifies. Each does a blocking-ish +/// outbound PCS fetch, so this caps the aggregate fan-out. +pub const DEFAULT_MAX_INFLIGHT: usize = 4; + +/// Token-bucket burst: this many requests may be served back-to-back before +/// the caller is throttled to the refill rate. +pub const DEFAULT_BURST: f64 = 10.0; + +/// One token is restored per this interval, up to [`DEFAULT_BURST`]. At the +/// default (1 token / 2s) a saturating caller is held to ~30 verifies/min. +pub const DEFAULT_REFILL: Duration = Duration::from_secs(2); + +/// Outcome of consulting the throttle. +#[derive(Debug)] +pub enum Decision { + /// Proceed; hold `permit` until the verify completes. + Proceed(OwnedSemaphorePermit), + /// Token bucket empty — reject with 429. + RateLimited, + /// Inflight cap saturated — reject with 429. + AtCapacity, +} + +struct Bucket { + tokens: f64, + last: Instant, +} + +/// Global rate + concurrency limiter for `/verify-quote`. Interior mutability +/// (a small `Mutex` never held across `.await`) so a single shared instance +/// can be consulted from concurrent request handlers. +pub struct VerifyQuoteThrottle { + inflight: Arc, + bucket: Mutex, + burst: f64, + refill: Duration, +} + +impl VerifyQuoteThrottle { + /// Construct a throttle with explicit limits. + /// + /// # Panics + /// + /// Panics if `max_inflight == 0`, `burst < 1.0`, or `refill == 0`: with no + /// inflight permits no request could ever proceed, a sub-unit burst can + /// never satisfy the `tokens >= 1.0` gate, and a zero refill interval would + /// make the refill rate non-finite (poisoning the lazy-refill arithmetic + /// with `NaN`) — so each renders the endpoint useless. These are + /// construction-time programmer errors — the only in-tree callers are + /// [`Default`] and tests, both passing valid constants — so they are + /// asserted rather than surfaced as a `Result`. + pub fn new(max_inflight: usize, burst: f64, refill: Duration) -> Self { + assert!(max_inflight > 0, "max_inflight must be positive"); + assert!(burst >= 1.0, "burst must be >= 1"); + assert!(refill > Duration::ZERO, "refill must be positive"); + Self { + inflight: Arc::new(Semaphore::new(max_inflight)), + // Start full, anchored at construction. The first `check_at` credits + // elapsed-since-construction (clamped to `burst`), so the bucket is + // full on first use regardless of how long the `LazyLock` sat idle. + // `check_at` uses `saturating_duration_since`, so even a `now` passed + // slightly *before* this `Instant::now()` (e.g. a fixed `now` in + // tests captured just after construction) yields 0 elapsed, not a + // panic — the bucket simply stays full. + bucket: Mutex::new(Bucket { + tokens: burst, + last: Instant::now(), + }), + burst, + refill, + } + } + + /// Consult the throttle for a request observed at `now`. On + /// `Decision::Proceed` one token and one inflight permit are consumed; on + /// rejection nothing is consumed. + pub fn check_at(&self, now: Instant) -> Decision { + // `refill` is asserted `> 0` in `new`, so this is always finite. + let refill_per_sec = 1.0 / self.refill.as_secs_f64(); + + { + // Recover from a poisoned mutex rather than propagating the panic: + // the only data behind the lock is the token bucket, and a stale + // bucket from a thread that panicked mid-update is harmless (the + // refill below reconciles it against `now`). Letting the poison + // propagate via `.expect()` would instead wedge the endpoint + // permanently — every later request would re-panic (AGENTS.md: avoid + // `.expect()`). + let mut bucket = self + .bucket + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + // Compute the lazily-refilled token count but do NOT write it back + // yet — the bucket is mutated transactionally, only on `Proceed`. A + // rejected request (RateLimited or AtCapacity) leaves `tokens` and + // `last` untouched, so it neither burns a token nor advances the + // refill clock; the bucket keeps refilling from the last *grant*. + let elapsed = now.saturating_duration_since(bucket.last).as_secs_f64(); + let refilled = (bucket.tokens + elapsed * refill_per_sec).min(self.burst); + if refilled < 1.0 { + return Decision::RateLimited; + } + + // Acquire the inflight permit while holding the bucket lock so the + // token and permit are committed atomically. `try_acquire_owned` + // never blocks (and the lock is never held across an await), so the + // hold is bounded to a few instructions — releasing and re-taking + // the lock around the acquire would instead open a TOCTOU window + // where a concurrent grant could be clobbered, so the lock is held + // across the acquire intentionally. + // + // Lock ordering: this is the only place the bucket mutex and the + // semaphore's internal lock are held together, and the order is + // always bucket-mutex → semaphore (via `try_acquire_owned`). The + // permit is released only by dropping it in the request handler / + // spawned verify, which never holds the bucket mutex, so the + // ordering is never inverted and there is no deadlock risk. + match Arc::clone(&self.inflight).try_acquire_owned() { + Ok(permit) => { + // Commit only on success: the token and refill clock advance + // together with the permit grant. + bucket.tokens = refilled - 1.0; + bucket.last = now; + Decision::Proceed(permit) + } + // AtCapacity: leave the bucket untouched — no token is burned + // and `last` is not advanced, so the refill clock keeps running + // from the previous grant. + Err(_) => Decision::AtCapacity, + } + } + } + + /// Consult the throttle for a request observed now. + pub fn check(&self) -> Decision { + self.check_at(Instant::now()) + } +} + +impl Default for VerifyQuoteThrottle { + fn default() -> Self { + Self::new(DEFAULT_MAX_INFLIGHT, DEFAULT_BURST, DEFAULT_REFILL) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bursts_then_rate_limits() { + // Burst 3, refill 1 token/sec, generous inflight so capacity isn't the + // gate. Permits are dropped immediately so only the token bucket bites. + let t = VerifyQuoteThrottle::new(1000, 3.0, Duration::from_secs(1)); + let now = Instant::now(); + for i in 0..3 { + assert!( + matches!(t.check_at(now), Decision::Proceed(_)), + "burst {i} should pass" + ); + } + assert!(matches!(t.check_at(now), Decision::RateLimited)); + // One token refills after a second. + assert!(matches!( + t.check_at(now + Duration::from_secs(1)), + Decision::Proceed(_) + )); + assert!(matches!( + t.check_at(now + Duration::from_secs(1)), + Decision::RateLimited + )); + } + + #[test] + fn rejections_do_not_slow_token_recovery() { + // Regression guard for the transactional bucket: `tokens`/`last` are + // committed only on `Proceed`, so a rejected request must not steal + // refill credit. Burst 1, refill 1 token/sec, generous inflight. + let t = VerifyQuoteThrottle::new(1000, 1.0, Duration::from_secs(1)); + let t0 = Instant::now(); + assert!(matches!(t.check_at(t0), Decision::Proceed(_))); + // Half a second in: only 0.5 token → rejected, and this call must NOT + // advance the refill clock. + assert!(matches!( + t.check_at(t0 + Duration::from_millis(500)), + Decision::RateLimited + )); + // At exactly t0+1s a full token is back despite the intervening + // rejection — recovery tracks wall-clock from the last grant. + assert!(matches!( + t.check_at(t0 + Duration::from_secs(1)), + Decision::Proceed(_) + )); + } + + #[test] + fn inflight_cap_blocks_when_saturated() { + // Cap 2; huge burst so the rate limit isn't the gate. Hold the permits. + let t = VerifyQuoteThrottle::new(2, 1000.0, Duration::from_millis(1)); + let now = Instant::now(); + let p1 = match t.check_at(now) { + Decision::Proceed(p) => p, + other => panic!("expected Proceed, got {other:?}"), + }; + let _p2 = match t.check_at(now) { + Decision::Proceed(p) => p, + other => panic!("expected Proceed, got {other:?}"), + }; + assert!(matches!(t.check_at(now), Decision::AtCapacity)); + drop(p1); + assert!(matches!(t.check_at(now), Decision::Proceed(_))); + } + + #[test] + fn rate_limited_request_does_not_consume_permit() { + // Burst 1, refill never. After the first request the bucket is empty; + // the rejection must not have taken an inflight permit. + let t = VerifyQuoteThrottle::new(2, 1.0, Duration::from_secs(100_000)); + let now = Instant::now(); + let _held = match t.check_at(now) { + Decision::Proceed(p) => p, + other => panic!("expected Proceed, got {other:?}"), + }; + // Rate limited (bucket empty), even though an inflight slot is free. + assert!(matches!(t.check_at(now), Decision::RateLimited)); + assert_eq!(t.inflight.available_permits(), 1); + } +}