cubeapi: key the rate limiter on the validated identity, not an unvalidated header

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<Self> {
         let _ = dotenvy::dotenv();
         let cfg = config::Config::builder()

CubeAPI/src/main.rs

@@ -98,7 +98,7 @@ struct Cli {
     #[arg(long, value_name = "PREFIX")]
     log_prefix: Option<String>,
 
-    /// 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")]

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,24 @@ fn extract_credential(request: &Request) -> Option<AuthCredential> {
     None
 }
 
+fn identity_of(credential: &AuthCredential) -> String {
+    match credential {
+        AuthCredential::Bearer(t) => format!("bearer:{}", t),
+        AuthCredential::ApiKey(k) => format!("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 +99,7 @@ fn extract_credential(request: &Request) -> Option<AuthCredential> {
 /// callback to enforce fine-grained (path + method) authorization.
 pub async fn unified_auth(
     State(state): State<AppState>,
-    request: Request,
+    mut request: Request,
     next: Next,
 ) -> Result<Response, AppError> {
     // Mode 1: callback auth — if a callback URL is configured, forward the
@@ -103,7 +124,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 +134,10 @@ pub async fn unified_auth(
                         "Invalid API key or token".to_string(),
                     ));
                 }
+                request.extensions_mut().insert(RateLimitIdentity(format!(
+                    "configured-key:{}",
+                    expected_key
+                )));
             }
         }
         // Mode 3: no auth (both unset) or simple-key match — pass through.
@@ -160,6 +185,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,

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<AppState>,
     request: Request,
     next: Next,
 ) -> Result<Response, AppError> {
-    // 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::<crate::middleware::auth::RateLimitIdentity>()
+        .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),

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

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<DefaultKeyedRateLimiter<String>>,
 
     /// 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,13 @@ impl AppState {
         }
     }
 }
+
+fn spawn_rate_limiter_gc(limiter: Arc<DefaultKeyedRateLimiter<String>>) {
+    tokio::spawn(async move {
+        let mut ticker = tokio::time::interval(std::time::Duration::from_secs(60));
+        loop {
+            ticker.tick().await;
+            limiter.retain_recent();
+        }
+    });
+}