Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dwin-gharibi wants to merge 5 commits into
base: master
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7f6c39c
fix(CubeAPI/src/middleware/rate_limit.rs): key the rate limiter on th…
dwin-gharibi Aug 18, 2026
3ef5ef4
fix(CubeAPI/src/middleware/auth.rs): share one bucket per configured …
dwin-gharibi Aug 18, 2026
5511d29
test(CubeAPI/src/routes.rs): pin the rotating X-API-Key bypass at the…
dwin-gharibi Aug 18, 2026
2fdf4da
fix(CubeAPI/src/middleware/rate_limit.rs): assert unified_auth runs f…
dwin-gharibi Aug 18, 2026
f05fdd4
test(CubeAPI/src/routes.rs): cover callback-mode multi-tenant bucket …
dwin-gharibi Aug 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 28 additions & 2 deletions CubeAPI/src/middleware/auth.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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(),
Expand All @@ -113,6 +134,8 @@ pub async fn unified_auth(
"Invalid API key or token".to_string(),
));
}
let identity = identity_of(&credential);
request.extensions_mut().insert(RateLimitIdentity(identity));
}
}
// Mode 3: no auth (both unset) or simple-key match — pass through.
Expand Down Expand Up @@ -160,6 +183,9 @@ pub async fn unified_auth(
};

if callback_resp.status().as_u16() == 200 {
request

@cubesandboxbot cubesandboxbot Bot Aug 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The callback-mode identity stores the raw credential as the bucket key. For Bearer tokens this is often a multi-KB JWT, so every distinct token becomes a large DashMap entry that lingers until the GC sweep — and raw credentials stay resident in memory past the request. Since the limiter only needs to distinguish identities (not read the token back), key on a fixed-size hash instead, e.g. format!("bearer:{}", sha256(token)) / format!("apikey:{}", sha256(key)). That preserves per-tenant separation while bounding per-key memory and not retaining plaintext secrets. (This is also the shape of the unbounded-growth attack the GC in state.rs is meant to bound — hashing makes each entry small regardless.)

.extensions_mut()
.insert(RateLimitIdentity(identity_of(&credential)));
tracing::debug!(
path = %request_path,
method = %request_method,
Expand Down
10 changes: 4 additions & 6 deletions CubeAPI/src/middleware/rate_limit.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,11 @@ pub async fn rate_limit(
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();
.extensions()

@cubesandboxbot cubesandboxbot Bot Aug 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale doc comments now that the key comes from the validated identity, not the header. The module doc above still reads: "Per-API-key token bucket rate limiter middleware. Reads the X-API-Key header and checks the shared governor limiter." That's no longer accurate — this now keys on the RateLimitIdentity extension published by unified_auth (which is itself keyed on the validated credential). Same for AppState::rate_limiter's "Per-API-key rate limiter" field doc. Worth a one-line update so the docs don't mislead a future reader into re-adding header-based keying.

.get::<crate::middleware::auth::RateLimitIdentity>()
.map(|id| id.0.clone())
.unwrap_or_else(|| "unauthenticated".to_string());

@cubesandboxbot cubesandboxbot Bot Aug 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fallback masks a fragile layering invariant. The fix only works because, in with_auth_and_rate_limit (routes.rs:214), unified_auth is the outermost of the two layers (in axum the last .layer() added runs first), so the RateLimitIdentity extension is always populated by the time this runs. If those two .layer() calls are ever reordered, every request silently lands in this single shared "unauthenticated" bucket — re-introducing exactly the multi-tenant starvation this PR fixes — with no error, no log, and no test catching it. Current code is correct, but consider a comment on with_auth_and_rate_limit (or a debug assertion) documenting that unified_auth must run before rate_limit.

match state.rate_limiter.check_key(&key) {
Ok(_) => Ok(next.run(request).await),
Expand Down
11 changes: 11 additions & 0 deletions CubeAPI/src/state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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();

@cubesandboxbot cubesandboxbot Bot Aug 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This call won't compile against the pinned governor 0.6.3 (CubeAPI/Cargo.lock). In 0.6.x, KeyedRateLimiter::retain_recent is pub fn retain_recent(&self, up_to: Instant) -> usize — it takes an Instant argument, so limiter.retain_recent() is an arity error (E0061).

Also check the semantics once you add the argument: retain_recent(up_to) retains only keys whose last update is after up_to. Passing Instant::now() at each tick drops every key (every key's last update is before "now"), wiping the whole map — i.e. every client's bucket resets once a minute. That defeats the limiter for clients that are right at their burst boundary when the tick fires (they get a fresh full bucket). Pass a retention window instead, e.g.

let window = std::time::Instant::now() - std::time::Duration::from_secs(120);
limiter.retain_recent(window);

so recently-active keys (and the burst they've accrued) survive the sweep while abandoned keys are reclaimed.

}
});
}
Loading