From 8170e0ed187c8dd99903bfa9f889d09cc0c0d63e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 24 Jun 2026 12:33:39 -0700 Subject: [PATCH] feat(ai-gateway): bounded (200) TTL key cache --- .../crates/ai-gateway/src/auth/cache.rs | 153 ++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100644 litellm-rust/crates/ai-gateway/src/auth/cache.rs diff --git a/litellm-rust/crates/ai-gateway/src/auth/cache.rs b/litellm-rust/crates/ai-gateway/src/auth/cache.rs new file mode 100644 index 00000000000..fe3fb459328 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/auth/cache.rs @@ -0,0 +1,153 @@ +//! A bounded, in-memory cache of verified keys. +//! +//! This bounds two things at once: **memory** (at most [`MAX_ENTRIES`] entries, +//! FIFO-evicted) and **staleness** (each entry expires after [`KeyCache::ttl`], so +//! a revoked or re-budgeted key is re-verified within the TTL window). Both are +//! deliberate: the gateway trades a small staleness window for not hitting the +//! auth backend on every request. +//! +//! Std-only (no external cache crate): a `Mutex` for lookup plus a +//! `VecDeque` recording insertion order for FIFO eviction. + +use std::collections::{HashMap, VecDeque}; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use crate::auth::UserApiKeyAuth; + +/// Hard cap on cached entries; the 201st insert evicts the oldest. Keeps memory +/// bounded regardless of how many distinct keys hit the gateway. +const MAX_ENTRIES: usize = 200; +/// Default time-to-live when `LITELLM_AUTH_CACHE_TTL_SECS` is unset/invalid. +const DEFAULT_TTL_SECS: u64 = 60; +/// Env var overriding the entry TTL (seconds). +const TTL_ENV: &str = "LITELLM_AUTH_CACHE_TTL_SECS"; + +struct Inner { + /// hash → (identity, insertion time). Insertion time drives TTL expiry. + entries: HashMap<[u8; 32], (UserApiKeyAuth, Instant)>, + /// Insertion order, oldest first, for FIFO eviction. + order: VecDeque<[u8; 32]>, +} + +/// Bounded TTL cache keyed by the SHA-256 of the API key. +pub struct KeyCache { + inner: Mutex, + ttl: Duration, +} + +impl KeyCache { + /// Build with the TTL from `LITELLM_AUTH_CACHE_TTL_SECS` (falling back to + /// [`DEFAULT_TTL_SECS`] when unset or unparseable). + pub fn new() -> Self { + let ttl_secs = std::env::var(TTL_ENV) + .ok() + .and_then(|raw| raw.trim().parse::().ok()) + .unwrap_or(DEFAULT_TTL_SECS); + Self::with_ttl(Duration::from_secs(ttl_secs)) + } + + /// Build with an explicit TTL (used by tests for tiny windows). + pub fn with_ttl(ttl: Duration) -> Self { + Self { + inner: Mutex::new(Inner { + entries: HashMap::new(), + order: VecDeque::new(), + }), + ttl, + } + } + + /// Return the cached identity for `hash`, or `None`. A present-but-expired + /// entry is evicted and treated as a miss, so the caller re-verifies. + pub fn get(&self, hash: &[u8; 32]) -> Option { + let mut inner = self.inner.lock().expect("auth cache mutex poisoned"); + let expired = match inner.entries.get(hash) { + Some((_, inserted)) => inserted.elapsed() >= self.ttl, + None => return None, + }; + if expired { + inner.entries.remove(hash); + inner.order.retain(|h| h != hash); + return None; + } + inner.entries.get(hash).map(|(auth, _)| auth.clone()) + } + + /// Insert (or refresh) an entry, evicting the oldest if at capacity. + pub fn insert(&self, hash: [u8; 32], auth: UserApiKeyAuth) { + let mut inner = self.inner.lock().expect("auth cache mutex poisoned"); + + // Refreshing an existing key: drop its old position so we don't double-count + // it in `order` (which would corrupt FIFO and the capacity check). + if inner.entries.contains_key(&hash) { + inner.order.retain(|h| h != &hash); + } + + while inner.entries.len() >= MAX_ENTRIES { + match inner.order.pop_front() { + Some(oldest) => { + inner.entries.remove(&oldest); + } + None => break, + } + } + + inner.entries.insert(hash, (auth, Instant::now())); + inner.order.push_back(hash); + } +} + +impl Default for KeyCache { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn hash(byte: u8) -> [u8; 32] { + [byte; 32] + } + + #[test] + fn get_returns_inserted_entry() { + let cache = KeyCache::with_ttl(Duration::from_secs(60)); + let auth = UserApiKeyAuth { + user_id: Some("u".to_string()), + ..UserApiKeyAuth::default() + }; + cache.insert(hash(1), auth); + + let got = cache.get(&hash(1)).expect("entry present"); + assert_eq!(got.user_id.as_deref(), Some("u")); + assert!(cache.get(&hash(2)).is_none()); + } + + #[test] + fn evicts_oldest_at_capacity() { + let cache = KeyCache::with_ttl(Duration::from_secs(60)); + // Fill exactly to capacity. + for i in 0..MAX_ENTRIES { + cache.insert(hash(i as u8), UserApiKeyAuth::default()); + } + assert!(cache.get(&hash(0)).is_some()); + + // The 201st distinct insert evicts the oldest (hash(0)). + cache.insert(hash(255), UserApiKeyAuth::default()); + assert!(cache.get(&hash(0)).is_none()); + assert!(cache.get(&hash(255)).is_some()); + // A middle entry survives. + assert!(cache.get(&hash(1)).is_some()); + } + + #[test] + fn get_returns_none_after_ttl() { + let cache = KeyCache::with_ttl(Duration::from_millis(10)); + cache.insert(hash(1), UserApiKeyAuth::default()); + std::thread::sleep(Duration::from_millis(25)); + assert!(cache.get(&hash(1)).is_none()); + } +}