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

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)
 

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,29 @@ fn extract_credential(request: &Request) -> Option<AuthCredential> {
     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<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 +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,

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),