diff --git a/CubeAPI/Cargo.lock b/CubeAPI/Cargo.lock index cc2d84fc2..af0773bda 100644 --- a/CubeAPI/Cargo.lock +++ b/CubeAPI/Cargo.lock @@ -536,6 +536,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "sha2", "thiserror 1.0.69", "tokio", "tower 0.4.13", diff --git a/CubeAPI/Cargo.toml b/CubeAPI/Cargo.toml index 96a084e29..2a98432ad 100644 --- a/CubeAPI/Cargo.toml +++ b/CubeAPI/Cargo.toml @@ -59,6 +59,7 @@ uuid = { version = "1", features = ["v4", "serde"] } # ── High-concurrency in-memory state ────────────────────────────────────── # Lock-free concurrent HashMap: O(1) reads without global lock dashmap = "5" +sha2 = "0.10" # Atomic counters / flags # (std::sync::atomic is sufficient for simple cases) diff --git a/CubeAPI/src/config/mod.rs b/CubeAPI/src/config/mod.rs index 0401c4781..aad4c07c4 100644 --- a/CubeAPI/src/config/mod.rs +++ b/CubeAPI/src/config/mod.rs @@ -18,7 +18,7 @@ pub struct ServerConfig { #[serde(default = "default_worker_threads")] pub worker_threads: usize, - /// Rate limit: max requests per second per API key + /// Rate limit: max requests per second per validated identity #[serde(default = "default_rate_limit")] pub rate_limit_per_sec: u32, @@ -111,6 +111,13 @@ fn default_log_prefix() -> String { } impl ServerConfig { + pub fn auth_is_configured(&self) -> bool { + self.auth_callback_url + .as_deref() + .is_some_and(|u| !u.is_empty()) + || self.cube_api_key.as_deref().is_some_and(|k| !k.is_empty()) + } + pub fn from_env() -> anyhow::Result { let _ = dotenvy::dotenv(); let cfg = config::Config::builder() diff --git a/CubeAPI/src/main.rs b/CubeAPI/src/main.rs index 6dce99544..7c87706cd 100644 --- a/CubeAPI/src/main.rs +++ b/CubeAPI/src/main.rs @@ -98,7 +98,7 @@ struct Cli { #[arg(long, value_name = "PREFIX")] log_prefix: Option, - /// Rate limit: max requests per second per API key (default: 100). + /// Rate limit: max requests per second per validated identity (default: 100). /// /// Overrides the RATE_LIMIT_PER_SEC environment variable. #[arg(long, value_name = "N")] diff --git a/CubeAPI/src/middleware/auth.rs b/CubeAPI/src/middleware/auth.rs index e402c9ef7..4d4e7939e 100644 --- a/CubeAPI/src/middleware/auth.rs +++ b/CubeAPI/src/middleware/auth.rs @@ -9,6 +9,9 @@ use axum::{ response::Response, }; +#[derive(Debug, Clone)] +pub struct RateLimitIdentity(pub String); + /// Auth credential extracted from the request headers. #[derive(Debug)] enum AuthCredential { @@ -47,6 +50,29 @@ fn extract_credential(request: &Request) -> Option { None } +fn identity_hash(kind: &str, credential: &str) -> String { + use sha2::{Digest, Sha256}; + format!("{}:{:x}", kind, Sha256::digest(credential.as_bytes())) +} + +fn identity_of(credential: &AuthCredential) -> String { + match credential { + AuthCredential::Bearer(t) => identity_hash("bearer", t), + AuthCredential::ApiKey(k) => identity_hash("apikey", k), + } +} + +fn constant_time_eq(a: &[u8], b: &[u8]) -> bool { + if a.len() != b.len() { + return false; + } + let mut diff = 0u8; + for (x, y) in a.iter().zip(b.iter()) { + diff |= x ^ y; + } + diff == 0 +} + /// Unified auth middleware. /// /// Behavior (priority order): @@ -78,7 +104,7 @@ fn extract_credential(request: &Request) -> Option { /// callback to enforce fine-grained (path + method) authorization. pub async fn unified_auth( State(state): State, - request: Request, + mut request: Request, next: Next, ) -> Result { // Mode 1: callback auth — if a callback URL is configured, forward the @@ -103,7 +129,7 @@ pub async fn unified_auth( AuthCredential::Bearer(t) => t.as_str(), AuthCredential::ApiKey(k) => k.as_str(), }; - if provided != expected_key { + if !constant_time_eq(provided.as_bytes(), expected_key.as_bytes()) { tracing::warn!( path = %request.uri().path(), method = %request.method(), @@ -113,6 +139,9 @@ pub async fn unified_auth( "Invalid API key or token".to_string(), )); } + request + .extensions_mut() + .insert(RateLimitIdentity("configured-key".to_string())); } } // Mode 3: no auth (both unset) or simple-key match — pass through. @@ -160,6 +189,9 @@ pub async fn unified_auth( }; if callback_resp.status().as_u16() == 200 { + request + .extensions_mut() + .insert(RateLimitIdentity(identity_of(&credential))); tracing::debug!( path = %request_path, method = %request_method, @@ -183,6 +215,51 @@ pub async fn unified_auth( #[cfg(test)] mod tests { + use super::{identity_of, AuthCredential}; + + #[test] + fn rate_limit_identities_never_contain_the_raw_credential() { + let token = "eyJhbGciOiJIUzI1NiJ9.super-secret-tenant-token.signature"; + let key = "sk-live-super-secret-api-key"; + + let bearer = identity_of(&AuthCredential::Bearer(token.to_string())); + let apikey = identity_of(&AuthCredential::ApiKey(key.to_string())); + + assert!( + !bearer.contains(token), + "bearer identity leaks the token: {bearer}" + ); + assert!( + !apikey.contains(key), + "apikey identity leaks the key: {apikey}" + ); + assert!(bearer.starts_with("bearer:")); + assert!(apikey.starts_with("apikey:")); + assert_eq!( + bearer.len(), + "bearer:".len() + 64, + "expected a hex sha256 digest" + ); + assert_eq!( + apikey.len(), + "apikey:".len() + 64, + "expected a hex sha256 digest" + ); + } + + #[test] + fn the_same_credential_always_maps_to_the_same_identity() { + let a = identity_of(&AuthCredential::Bearer("tok".to_string())); + let b = identity_of(&AuthCredential::Bearer("tok".to_string())); + let c = identity_of(&AuthCredential::ApiKey("tok".to_string())); + + assert_eq!(a, b, "identity must be stable or buckets would churn"); + assert_ne!( + a, c, + "a bearer token and an api key with the same value must not collide" + ); + } + use super::*; use crate::{ config::ServerConfig, diff --git a/CubeAPI/src/middleware/rate_limit.rs b/CubeAPI/src/middleware/rate_limit.rs index f399d9d86..168864866 100644 --- a/CubeAPI/src/middleware/rate_limit.rs +++ b/CubeAPI/src/middleware/rate_limit.rs @@ -10,21 +10,27 @@ use axum::{ response::Response, }; -/// Per-API-key token bucket rate limiter middleware. -/// Reads the X-API-Key header and checks the shared governor limiter. -/// Returns 429 if the key has exceeded its quota. +/// Per-identity token bucket rate limiter middleware. +/// Reads the `RateLimitIdentity` published by `unified_auth` after it validated +/// the credential, and checks the shared governor limiter. +/// Returns 429 if that identity has exceeded its quota. pub async fn rate_limit( State(state): State, request: Request, next: Next, ) -> Result { - // Extract key; fall back to IP or "anonymous" - let key = request - .headers() - .get("X-API-Key") - .and_then(|v| v.to_str().ok()) - .unwrap_or("anonymous") - .to_string(); + let identity = request + .extensions() + .get::() + .map(|id| id.0.clone()); + + debug_assert!( + identity.is_some() || !state.config.auth_is_configured(), + "unified_auth must run before rate_limit: no RateLimitIdentity was published, \ + so every request would share one bucket" + ); + + let key = identity.unwrap_or_else(|| "unauthenticated".to_string()); match state.rate_limiter.check_key(&key) { Ok(_) => Ok(next.run(request).await), diff --git a/CubeAPI/src/routes.rs b/CubeAPI/src/routes.rs index 4774bf632..2e29e99c8 100644 --- a/CubeAPI/src/routes.rs +++ b/CubeAPI/src/routes.rs @@ -36,16 +36,7 @@ const PAUSE_RESUME_ROUTE_TIMEOUT: Duration = Duration::from_secs(120); const SNAPSHOT_LONG_ROUTE_TIMEOUT: Duration = Duration::from_secs(240); pub fn build_router(state: AppState) -> Router { - let auth_configured = state - .config - .auth_callback_url - .as_deref() - .is_some_and(|u| !u.is_empty()) - || state - .config - .cube_api_key - .as_deref() - .is_some_and(|k| !k.is_empty()); + let auth_configured = state.config.auth_is_configured(); let standard_router = apply_http_layers( Router::new().merge(build_e2b_router(&state, auth_configured)), @@ -246,13 +237,170 @@ mod tests { }; use axum::{ extract::Json, - http::{header::RETRY_AFTER, StatusCode}, + http::{ + header::{AUTHORIZATION, RETRY_AFTER}, + HeaderName, HeaderValue, StatusCode, + }, routing::delete, Router, }; use axum_test::TestServer; use serde_json::Value; + async fn spawn_approving_callback() -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("callback listener should bind"); + let address = listener.local_addr().expect("callback address"); + tokio::spawn(async move { + axum::serve( + listener, + Router::new().route("/auth", axum::routing::any(|| async { StatusCode::OK })), + ) + .await + .expect("callback server should run"); + }); + format!("http://{address}/auth") + } + + async fn callback_mode_server(callback_url: &str, rate_limit_per_sec: u32) -> TestServer { + let mut config = ServerConfig::default(); + config.cubemaster_url = "http://127.0.0.1:9".to_string(); + config.auth_callback_url = Some(callback_url.to_string()); + config.cube_api_key = None; + config.rate_limit_per_sec = rate_limit_per_sec; + + let state = AppState::new(config, arc(NoopLogger)).await; + TestServer::new(build_router(state)).expect("router should build") + } + + async fn throttled_count(server: &TestServer, token: &str, requests: usize) -> usize { + let mut throttled = 0; + for _ in 0..requests { + let response = server + .get("/sandboxes") + .add_header( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {token}")).expect("valid header"), + ) + .await; + if response.status_code() == StatusCode::TOO_MANY_REQUESTS { + throttled += 1; + } + } + throttled + } + + #[tokio::test] + async fn callback_mode_gives_distinct_tokens_independent_buckets() { + let callback_url = spawn_approving_callback().await; + let server = callback_mode_server(&callback_url, 3).await; + + let noisy = throttled_count(&server, "tenant-a-token", 30).await; + assert!( + noisy > 20, + "an abusive tenant was not throttled: {noisy}/30" + ); + + let quiet = throttled_count(&server, "tenant-b-token", 3).await; + assert_eq!( + quiet, 0, + "a quiet tenant was starved by another tenant's traffic: {quiet}/3 throttled" + ); + } + + async fn rate_limited_server(rate_limit_per_sec: u32) -> TestServer { + let mut config = ServerConfig::default(); + config.cubemaster_url = "http://127.0.0.1:9".to_string(); + config.auth_callback_url = None; + config.cube_api_key = Some("supersecret".to_string()); + config.rate_limit_per_sec = rate_limit_per_sec; + + let state = AppState::new(config, arc(NoopLogger)).await; + TestServer::new(build_router(state)).expect("router should build") + } + + #[tokio::test] + async fn rotating_an_unvalidated_api_key_header_cannot_refresh_the_bucket() { + let server = rate_limited_server(3).await; + + let mut statuses = Vec::new(); + for i in 0..30 { + let response = server + .get("/sandboxes") + .add_header( + AUTHORIZATION, + HeaderValue::from_static("Bearer supersecret"), + ) + .add_header( + HeaderName::from_static("x-api-key"), + HeaderValue::from_str(&format!("rotating-{i}")).expect("valid header"), + ) + .await; + statuses.push(response.status_code()); + } + + let throttled = statuses + .iter() + .filter(|s| **s == StatusCode::TOO_MANY_REQUESTS) + .count(); + assert!( + throttled > 20, + "rotating X-API-Key bypassed the limiter: only {throttled}/30 throttled, statuses {statuses:?}" + ); + } + + #[tokio::test] + async fn a_bearer_client_is_throttled_without_any_api_key_header() { + let server = rate_limited_server(3).await; + + let mut throttled = 0; + for _ in 0..30 { + let response = server + .get("/sandboxes") + .add_header( + AUTHORIZATION, + HeaderValue::from_static("Bearer supersecret"), + ) + .await; + if response.status_code() == StatusCode::TOO_MANY_REQUESTS { + throttled += 1; + } + } + assert!( + throttled > 20, + "bearer client was not throttled: {throttled}/30" + ); + } + + #[tokio::test] + async fn alternating_header_styles_share_one_bucket_in_simple_key_mode() { + let server = rate_limited_server(3).await; + + let mut throttled = 0; + for i in 0..30 { + let request = server.get("/sandboxes"); + let request = if i % 2 == 0 { + request.add_header( + AUTHORIZATION, + HeaderValue::from_static("Bearer supersecret"), + ) + } else { + request.add_header( + HeaderName::from_static("x-api-key"), + HeaderValue::from_static("supersecret"), + ) + }; + if request.await.status_code() == StatusCode::TOO_MANY_REQUESTS { + throttled += 1; + } + } + assert!( + throttled > 20, + "alternating Bearer and X-API-Key doubled the quota: only {throttled}/30 throttled" + ); + } + async fn test_server() -> TestServer { let mut config = ServerConfig::default(); config.cubemaster_url = "http://127.0.0.1:9".to_string(); diff --git a/CubeAPI/src/state.rs b/CubeAPI/src/state.rs index 154b57140..57cb86281 100644 --- a/CubeAPI/src/state.rs +++ b/CubeAPI/src/state.rs @@ -14,7 +14,7 @@ use std::sync::Arc; /// on every request, so real data must live behind Arc. #[derive(Clone)] pub struct AppState { - /// Per-API-key rate limiter (token bucket). + /// Per-identity rate limiter (token bucket), keyed on the validated credential. pub rate_limiter: Arc>, /// Shared reqwest connection pool. @@ -38,6 +38,7 @@ impl AppState { pub async fn new(config: crate::config::ServerConfig, logger: ArcLogger) -> Self { let quota = Quota::per_second(NonZeroU32::new(config.rate_limit_per_sec.max(1)).unwrap()); let rate_limiter = Arc::new(RateLimiter::keyed(quota)); + spawn_rate_limiter_gc(rate_limiter.clone()); let http_client = reqwest::Client::builder() .pool_max_idle_per_host(100) @@ -57,3 +58,73 @@ impl AppState { } } } + +fn spawn_rate_limiter_gc(limiter: Arc>) { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(std::time::Duration::from_secs(60)); + loop { + ticker.tick().await; + limiter.retain_recent(); + } + }); +} + +#[cfg(test)] +mod gc_tests { + use governor::{DefaultKeyedRateLimiter, Quota, RateLimiter}; + use std::num::NonZeroU32; + + fn limiter(per_sec: u32) -> DefaultKeyedRateLimiter { + RateLimiter::keyed(Quota::per_second(NonZeroU32::new(per_sec).unwrap())) + } + + #[test] + fn retain_recent_does_not_hand_an_active_key_a_fresh_bucket() { + let l = limiter(1); + let key = "active".to_string(); + + assert!(l.check_key(&key).is_ok(), "first request should pass"); + assert!( + l.check_key(&key).is_err(), + "second request should be throttled" + ); + + l.retain_recent(); + + assert!( + l.check_key(&key).is_err(), + "retain_recent reset an active key's bucket, so a client at its burst \ + boundary would get a full quota back on every sweep" + ); + } + + #[test] + fn retain_recent_reclaims_idle_keys() { + let l = limiter(1000); + for i in 0..64 { + assert!(l.check_key(&format!("idle-{i}")).is_ok()); + } + assert_eq!(l.len(), 64, "keys should be tracked before the sweep"); + + std::thread::sleep(std::time::Duration::from_millis(50)); + l.retain_recent(); + + assert_eq!( + l.len(), + 0, + "idle keys were not reclaimed, so the map grows without bound" + ); + } + + #[test] + fn distinct_keys_do_not_share_a_bucket() { + let l = limiter(1); + + assert!(l.check_key(&"a".to_string()).is_ok()); + assert!(l.check_key(&"a".to_string()).is_err()); + assert!( + l.check_key(&"b".to_string()).is_ok(), + "a second identity was throttled by the first one's traffic" + ); + } +}