use std::{sync::OnceLock, time::Duration}; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogProvider, }; use axum::http::Uri; use serde::{Deserialize, Serialize}; use serde_json::Value; use tracing::{debug, info, warn}; use crate::wallet_runtime::{ local_rejection_from_wallet_access, resolve_wallet_auth_gate_uncached, }; use crate::{AppState, GatewayError}; use super::super::GatewayControlDecision; use super::credentials::{ build_auth_context_cache_key, build_auth_context_cache_key_with_trusted_auth, current_unix_secs, extract_request_credentials, extract_request_credentials_with_trusted_auth, extract_trusted_admin_headers, hash_api_key, }; use super::gate::GatewayLocalAuthRejection; use super::principal::derive_principal_candidate; use super::types::{ GatewayCredentialCarrier, GatewayPrincipalCandidate, GatewayTrustedAuthHeaders, }; use crate::cache::{AuthContextCacheGeneration, AuthContextInflightRegistration}; use crate::headers::header_value_str; use crate::local_auth_token::{ decode_local_auth_token, local_auth_token_identity_matches_user, LocalAuthTokenType, }; const AUTH_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(60); const AUTH_CONTEXT_CACHE_REFRESH_INTERVAL: Duration = Duration::from_secs(10); const AUTH_CONTEXT_NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(10); const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 10_000; const AUTH_CONTEXT_CACHE_MAX_ENTRIES_ENV: &str = "AETHER_GATEWAY_AUTH_CONTEXT_CACHE_MAX_ENTRIES"; const AUTH_CONTEXT_CACHE_REFRESH_INTERVAL_SECS_ENV: &str = "AETHER_GATEWAY_AUTH_CONTEXT_CACHE_REFRESH_INTERVAL_SECS"; const AUTH_CONTEXT_NEGATIVE_CACHE_TTL_SECS_ENV: &str = "AETHER_GATEWAY_AUTH_CONTEXT_NEGATIVE_CACHE_TTL_SECS"; const AUTH_CONTEXT_NEGATIVE_CACHE_KEY_PREFIX: &str = "negative:"; #[derive(Debug, Clone, Deserialize)] struct AntigravityBearerBridgeConfig { #[serde(default)] enabled: bool, #[serde(default)] auth_user_id: String, #[serde(default)] auth_api_key_id: String, #[serde(default)] bearer_sha256_allowlist: Vec, #[serde(default)] allow_unverified_google_bearer: bool, } impl AntigravityBearerBridgeConfig { fn bearer_validation_mode(&self, raw_bearer: &str) -> Option<&'static str> { if !self.bearer_sha256_allowlist.is_empty() { let bearer_hash = hash_api_key(raw_bearer); return self .bearer_sha256_allowlist .iter() .any(|allowed| allowed.trim().eq_ignore_ascii_case(&bearer_hash)) .then_some("sha256_allowlist"); } self.allow_unverified_google_bearer .then_some("explicit_unverified") } } #[derive(Debug, Clone, Deserialize, Serialize)] pub(crate) struct GatewayControlAuthContext { pub(crate) user_id: String, pub(crate) api_key_id: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub(crate) username: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub(crate) api_key_name: Option, pub(crate) balance_remaining: Option, pub(crate) access_allowed: bool, #[serde(skip)] pub(crate) user_rate_limit: Option, #[serde(skip)] pub(crate) api_key_rate_limit: Option, #[serde(skip)] pub(crate) api_key_is_standalone: bool, #[serde(skip)] pub(crate) admin_bypass_limits: bool, #[serde(skip)] pub(crate) local_rejection: Option, #[serde(skip)] pub(crate) allowed_models: Option>, #[serde(skip)] pub(crate) ip_rules: Option>, /// Credential verifier that established this API-key identity. Long-lived /// executions use it to prove that a later row with the same IDs is still /// the record authenticated by the original request. #[serde(skip)] pub(crate) verified_api_key_hash: Option, } #[derive(Clone)] pub(crate) struct VerifiedApiKeyHash(String); impl VerifiedApiKeyHash { fn new(value: String) -> Self { Self(value) } fn as_str(&self) -> &str { self.0.as_str() } } impl std::fmt::Debug for VerifiedApiKeyHash { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter.write_str("VerifiedApiKeyHash([REDACTED])") } } #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct GatewayAdminPrincipalContext { pub(crate) user_id: String, pub(crate) user_role: String, pub(crate) session_id: Option, pub(crate) management_token_id: Option, pub(crate) management_token_permissions: Option>, } pub(in super::super) enum ControlDecisionAuthResolution { Resolved(GatewayControlDecision), } pub(in super::super) async fn resolve_control_decision_auth( state: &AppState, headers: &http::HeaderMap, uri: &Uri, trace_id: &str, decision: GatewayControlDecision, ) -> Result { resolve_control_decision_auth_with_trusted_auth( state, headers, uri, trace_id, decision, cfg!(test), ) .await } pub(in super::super) async fn resolve_control_decision_auth_with_trusted_auth( state: &AppState, headers: &http::HeaderMap, uri: &Uri, trace_id: &str, mut decision: GatewayControlDecision, trusted_auth_verified: bool, ) -> Result { if let Some(admin_principal) = resolve_trusted_admin_principal(headers, decision.auth_endpoint_signature.as_deref()) { log_admin_principal_resolution(trace_id, &decision, "trusted_headers", &admin_principal); decision.admin_principal = Some(admin_principal); } else if let Some(admin_principal) = resolve_local_admin_principal( state, headers, uri, decision.auth_endpoint_signature.as_deref(), ) .await? { log_admin_principal_resolution(trace_id, &decision, "local_session", &admin_principal); decision.admin_principal = Some(admin_principal); } let auth_context_cache_key = decision .auth_endpoint_signature .as_deref() .and_then(|signature| { build_auth_context_cache_key_with_trusted_auth( headers, uri, signature, trusted_auth_verified, ) }); let mut resolved_auth_context = None; if let Some(cache_key) = auth_context_cache_key.as_deref() { if let Some((auth_context, age)) = get_cached_auth_context_with_age(state, cache_key) { if auth_context_cache_refresh_due(state, age) { resolved_auth_context = Some( revalidate_cached_auth_context( state, cache_key, auth_context, decision.auth_endpoint_signature.as_deref(), headers, uri, trusted_auth_verified, ) .await?, ); } else { // The configured refresh interval is the bounded authorization // freshness window. This path remains a lock-free cache hit. resolved_auth_context = Some(auth_context); } } } if resolved_auth_context.is_none() { resolved_auth_context = resolve_data_backed_auth_context_cached( state, auth_context_cache_key.as_deref(), headers, uri, decision.auth_endpoint_signature.as_deref(), true, trusted_auth_verified, ) .await?; } if let Some(auth_context) = resolved_auth_context { apply_resolved_auth_context_to_decision(trace_id, &mut decision, auth_context); } if decision.local_auth_rejection.is_some() { log_local_auth_rejection(trace_id, &decision); return Ok(ControlDecisionAuthResolution::Resolved(decision)); } if decision.is_execution_runtime_candidate() { return Ok(ControlDecisionAuthResolution::Resolved(decision)); } if decision.auth_context.is_some() { return Ok(ControlDecisionAuthResolution::Resolved(decision)); } if allows_missing_data_backed_auth_context(&decision) { return Ok(ControlDecisionAuthResolution::Resolved(decision)); } Ok(ControlDecisionAuthResolution::Resolved(decision)) } fn log_admin_principal_resolution( trace_id: &str, decision: &GatewayControlDecision, resolution: &'static str, admin_principal: &GatewayAdminPrincipalContext, ) { debug!( event_name = "admin_principal_resolved", log_type = "debug", debug_context = "control_auth", trace_id = %trace_id, route_class = decision.route_class.as_deref().unwrap_or("unknown"), route_family = decision.route_family.as_deref().unwrap_or("unknown"), route_kind = decision.route_kind.as_deref().unwrap_or("unknown"), resolution, admin_user_id = admin_principal.user_id.as_str(), admin_user_role = admin_principal.user_role.as_str(), admin_session_id = admin_principal.session_id.as_deref().unwrap_or("-"), admin_management_token_id = admin_principal.management_token_id.as_deref().unwrap_or("-"), "resolved admin principal for control decision" ); } fn log_auth_context_resolution( trace_id: &str, decision: &GatewayControlDecision, auth_context: &GatewayControlAuthContext, ) { let balance_remaining = auth_context .balance_remaining .map(|value| format!("{value:.4}")) .unwrap_or_else(|| "-".to_string()); info!( event_name = "auth_context_resolved", log_type = "event", status = if auth_context.access_allowed { "allowed" } else { "blocked" }, trace_id = %trace_id, route_class = decision.route_class.as_deref().unwrap_or("unknown"), route_family = decision.route_family.as_deref().unwrap_or("unknown"), route_kind = decision.route_kind.as_deref().unwrap_or("unknown"), user_id = auth_context.user_id.as_str(), api_key_id = auth_context.api_key_id.as_str(), api_key_name = auth_context.api_key_name.as_deref().unwrap_or("-"), balance_remaining = balance_remaining.as_str(), access_allowed = auth_context.access_allowed, api_key_is_standalone = auth_context.api_key_is_standalone, has_local_rejection = auth_context.local_rejection.is_some(), "resolved data-backed auth context for control decision" ); } fn log_local_auth_rejection(trace_id: &str, decision: &GatewayControlDecision) { let Some(rejection) = decision.local_auth_rejection.as_ref() else { return; }; let (rejection_kind, rejection_detail) = match rejection { GatewayLocalAuthRejection::InvalidApiKey => ("invalid_api_key", "-".to_string()), GatewayLocalAuthRejection::LockedApiKey => ("locked_api_key", "-".to_string()), GatewayLocalAuthRejection::WalletUnavailable => ("wallet_unavailable", "-".to_string()), GatewayLocalAuthRejection::BalanceDenied { remaining } => ( "balance_denied", remaining .map(|value| format!("remaining_usd={value:.4}")) .unwrap_or_else(|| "remaining_usd=unknown".to_string()), ), GatewayLocalAuthRejection::ProviderNotAllowed { provider } => { ("provider_not_allowed", provider.clone()) } GatewayLocalAuthRejection::ApiFormatNotAllowed { api_format } => { ("api_format_not_allowed", api_format.clone()) } GatewayLocalAuthRejection::ModelNotAllowed { model } => { ("model_not_allowed", model.clone()) } GatewayLocalAuthRejection::IpNotAllowed { remote_ip } => { ("ip_not_allowed", remote_ip.clone()) } }; info!( event_name = "local_auth_rejected", log_type = "event", status = "rejected", trace_id = %trace_id, route_class = decision.route_class.as_deref().unwrap_or("unknown"), route_family = decision.route_family.as_deref().unwrap_or("unknown"), route_kind = decision.route_kind.as_deref().unwrap_or("unknown"), rejection_kind, rejection_detail = %rejection_detail, "rejected local control request during auth gate resolution" ); } fn allows_missing_data_backed_auth_context(decision: &GatewayControlDecision) -> bool { matches!( decision.route_kind.as_deref(), Some("chat" | "cli" | "compact") ) } fn resolve_trusted_admin_principal( headers: &http::HeaderMap, auth_endpoint_signature: Option<&str>, ) -> Option { if !auth_endpoint_signature .map(str::trim) .unwrap_or_default() .starts_with("admin:") { return None; } let trusted_headers = extract_trusted_admin_headers(headers)?; Some(GatewayAdminPrincipalContext { user_id: trusted_headers.user_id, user_role: trusted_headers.user_role, session_id: trusted_headers.session_id, management_token_id: trusted_headers.management_token_id, management_token_permissions: None, }) } async fn resolve_local_admin_principal( state: &AppState, headers: &http::HeaderMap, uri: &Uri, auth_endpoint_signature: Option<&str>, ) -> Result, GatewayError> { let Some(signature) = auth_endpoint_signature .map(str::trim) .filter(|value| value.starts_with("admin:")) else { return Ok(None); }; let extracted = extract_request_credentials(headers, uri, signature); let Some(access_token) = extracted.bundle.authorization_bearer.as_deref() else { return Ok(None); }; let claims = match decode_local_auth_token(access_token, LocalAuthTokenType::Access) { Ok(claims) => claims, Err(_) => return Ok(None), }; if claims .get("role") .and_then(Value::as_str) .is_some_and(|role| !crate::roles::can_access_admin_console(role)) { return Ok(None); } resolve_local_admin_principal_from_claims(state, headers, uri, &claims).await } pub(crate) async fn resolve_local_admin_session_principal( state: &AppState, headers: &http::HeaderMap, uri: &Uri, ) -> Result, GatewayError> { resolve_local_admin_principal(state, headers, uri, Some("admin:operational")).await } async fn resolve_local_admin_principal_from_claims( state: &AppState, headers: &http::HeaderMap, uri: &Uri, claims: &serde_json::Map, ) -> Result, GatewayError> { let Some(user_id) = claims.get("user_id").and_then(Value::as_str) else { return Ok(None); }; let Some(session_id) = claims.get("session_id").and_then(Value::as_str) else { return Ok(None); }; let Some(client_device_id) = extract_local_admin_client_device_id(headers, uri) else { return Ok(None); }; let Some(user) = state.find_user_auth_by_id(user_id).await? else { return Ok(None); }; if !user.is_active || user.is_deleted || !crate::roles::can_access_admin_console(&user.role) { return Ok(None); } if !local_auth_token_identity_matches_user(claims, &user) { return Ok(None); } let now = chrono::Utc::now(); let Some(session) = state.find_user_session(user_id, session_id).await? else { return Ok(None); }; if session.is_revoked() || session.is_expired(now) || session.security_version != user.security_version || session.client_device_id != client_device_id { return Ok(None); } if session.should_touch(now) { let _ = state .touch_user_session( user_id, session_id, now, None, local_admin_user_agent(headers).as_deref(), ) .await; } Ok(Some(GatewayAdminPrincipalContext { user_id: user.id, user_role: user.role, session_id: Some(session.id), management_token_id: None, management_token_permissions: None, })) } fn extract_local_admin_client_device_id(headers: &http::HeaderMap, uri: &Uri) -> Option { let header_value = header_value_str(headers, "x-client-device-id"); let query_value = uri.query().and_then(|query| { url::form_urlencoded::parse(query.as_bytes()) .find(|(key, _)| key == "client_device_id") .map(|(_, value)| value.into_owned()) }); let candidate = header_value.or(query_value)?; let candidate = candidate.trim(); if candidate.is_empty() || candidate.len() > 128 || !candidate .chars() .all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_') { return None; } Some(candidate.to_string()) } fn local_admin_user_agent(headers: &http::HeaderMap) -> Option { header_value_str(headers, http::header::USER_AGENT.as_str()) .map(|value| value.chars().take(1000).collect()) } pub(crate) async fn resolve_execution_runtime_auth_context( state: &AppState, decision: &GatewayControlDecision, headers: &http::HeaderMap, uri: &Uri, trace_id: &str, ) -> Result, GatewayError> { let _ = trace_id; if let Some(auth_context) = decision.auth_context.as_ref() { // Control-route auth resolution already refreshed and validated this context. // Revalidating here would perform a second snapshot/wallet lookup per request. return Ok(Some(auth_context.clone())); } let Some(auth_endpoint_signature) = decision.auth_endpoint_signature.as_deref() else { return Ok(None); }; let Some(cache_key) = build_auth_context_cache_key(headers, uri, auth_endpoint_signature) else { return Ok(None); }; if let Some((auth_context, age)) = get_cached_auth_context_with_age(state, &cache_key) { if auth_context_cache_refresh_due(state, age) { return revalidate_cached_auth_context( state, &cache_key, auth_context, Some(auth_endpoint_signature), headers, uri, cfg!(test), ) .await .map(Some); } return Ok(Some(auth_context)); } if let Some(auth_context) = resolve_data_backed_auth_context_cached( state, Some(cache_key.as_str()), headers, uri, Some(auth_endpoint_signature), true, cfg!(test), ) .await? { if auth_context.user_id.is_empty() || auth_context.api_key_id.is_empty() { return Ok(None); } return Ok(Some(auth_context)); } Ok(None) } async fn revalidate_cached_auth_context( state: &AppState, cache_key: &str, auth_context: GatewayControlAuthContext, auth_endpoint_signature: Option<&str>, headers: &http::HeaderMap, uri: &Uri, trusted_auth_verified: bool, ) -> Result { if is_negative_auth_context(&auth_context) || !auth_context.access_allowed || !state.has_auth_api_key_reader() || auth_context.user_id.trim().is_empty() || auth_context.api_key_id.trim().is_empty() || auth_endpoint_signature .map(str::trim) .filter(|value| !value.is_empty()) .is_none() { return Ok(auth_context); } loop { match state.auth_context_cache.register_inflight(cache_key) { AuthContextInflightRegistration::Leader(guard) => { let refreshed = match resolve_security_fresh_auth_context( state, headers, uri, auth_context.clone(), auth_endpoint_signature, trusted_auth_verified, ) .await { Ok(refreshed) => refreshed, Err(error) => { // A failed security refresh must not leave the old allow // available to this flight's followers. Publish the same // error so a failed backend read is not retried once per // follower. guard.fail(error.clone()); return Err(error); } }; if guard.generation_is_current() { put_cached_auth_context( state, cache_key.to_string(), refreshed.clone(), Some(guard.generation()), ); } return Ok(refreshed); } AuthContextInflightRegistration::Follower(waiter) => { waiter.wait().await?; if let Some((refreshed, age)) = get_cached_auth_context_with_age(state, cache_key) { // A cancelled/failed leader leaves the old due entry in // place. Only a newly published security-fresh value may // be reused by followers. if !auth_context_cache_refresh_due(state, age) { return Ok(refreshed); } } } AuthContextInflightRegistration::Bypass => { let refreshed = resolve_security_fresh_auth_context( state, headers, uri, auth_context, auth_endpoint_signature, trusted_auth_verified, ) .await; if refreshed.is_err() { state.auth_context_cache.invalidate(cache_key); } return refreshed; } } } } async fn resolve_security_fresh_auth_context( state: &AppState, headers: &http::HeaderMap, uri: &Uri, stale: GatewayControlAuthContext, auth_endpoint_signature: Option<&str>, trusted_auth_verified: bool, ) -> Result { if let Some(refreshed) = resolve_data_backed_auth_context_with_trusted_auth( state, headers, uri, auth_endpoint_signature, trusted_auth_verified, ) .await? { return Ok(refreshed); } let mut denied = stale; denied.access_allowed = false; denied.local_rejection = Some(GatewayLocalAuthRejection::InvalidApiKey); denied.balance_remaining = None; Ok(denied) } async fn resolve_data_backed_auth_context_cached( state: &AppState, cache_key: Option<&str>, headers: &http::HeaderMap, uri: &Uri, auth_endpoint_signature: Option<&str>, cache_negative: bool, trusted_auth_verified: bool, ) -> Result, GatewayError> { let Some(cache_key) = cache_key else { return resolve_data_backed_auth_context_with_trusted_auth( state, headers, uri, auth_endpoint_signature, trusted_auth_verified, ) .await; }; loop { match state.auth_context_cache.register_inflight(cache_key) { AuthContextInflightRegistration::Leader(guard) => { let resolved = match resolve_data_backed_auth_context_with_trusted_auth( state, headers, uri, auth_endpoint_signature, trusted_auth_verified, ) .await { Ok(resolved) => resolved, Err(error) => { guard.fail(error.clone()); return Err(error); } }; if let Some(auth_context) = resolved.as_ref() { if cache_negative || (!auth_context.user_id.is_empty() && !auth_context.api_key_id.is_empty()) { if guard.generation_is_current() { put_cached_auth_context( state, cache_key.to_string(), auth_context.clone(), Some(guard.generation()), ); } } } return Ok(resolved); } AuthContextInflightRegistration::Follower(notified) => { notified.wait().await?; if let Some(auth_context) = get_cached_auth_context(state, cache_key) { return Ok(Some(auth_context)); } if !cache_negative { return Ok(None); } } AuthContextInflightRegistration::Bypass => { return resolve_data_backed_auth_context_with_trusted_auth( state, headers, uri, auth_endpoint_signature, trusted_auth_verified, ) .await; } } } } pub(crate) async fn refresh_execution_runtime_auth_context( state: &AppState, auth_context: GatewayControlAuthContext, auth_endpoint_signature: Option<&str>, ) -> Result { refresh_execution_runtime_auth_context_with_snapshot( state, auth_context, auth_endpoint_signature, ) .await .map(|(auth_context, _)| auth_context) } /// Strongly refreshes the long-lived execution authorization context and /// returns the exact API-key snapshot that produced it. /// /// WebSocket turns need both values: using the refreshed context for RPM and /// balance checks while letting the planner independently read its normal /// cache can authorize a different provider/model snapshot for up to the cache /// TTL. Ordinary HTTP callers keep using [`refresh_execution_runtime_auth_context`]. pub(crate) async fn refresh_execution_runtime_auth_context_with_snapshot( state: &AppState, auth_context: GatewayControlAuthContext, auth_endpoint_signature: Option<&str>, ) -> Result< ( GatewayControlAuthContext, Option, ), GatewayError, > { if auth_context.local_rejection.is_some() || !auth_context.access_allowed { return Ok((auth_context, None)); } let Some(auth_endpoint_signature) = auth_endpoint_signature .map(str::trim) .filter(|value| !value.is_empty()) else { return Ok((auth_context, None)); }; if !state.has_auth_api_key_reader() || auth_context.user_id.trim().is_empty() || auth_context.api_key_id.trim().is_empty() { return Ok((auth_context, None)); } let verified_api_key_hash = auth_context.verified_api_key_hash.clone(); let snapshot = { let _permit = state.acquire_auth_snapshot_load_gate().await?; if let Some(key_hash) = verified_api_key_hash.as_ref() { state .data .read_auth_api_key_snapshot_by_key_hash_strong( key_hash.as_str(), current_unix_secs(), ) .await .map_err(|err| GatewayError::Internal(err.to_string()))? } else { state .data .read_auth_api_key_snapshot_strong( &auth_context.user_id, &auth_context.api_key_id, current_unix_secs(), ) .await .map_err(|err| GatewayError::Internal(err.to_string()))? } }; let Some(snapshot) = snapshot else { return Ok((deny_refreshed_auth_context(auth_context), None)); }; if snapshot.user_id != auth_context.user_id || snapshot.api_key_id != auth_context.api_key_id { return Ok((deny_refreshed_auth_context(auth_context), None)); }; let wallet_access = resolve_wallet_auth_gate_uncached(state, &snapshot).await?; let mut refreshed = build_data_backed_auth_context( state, snapshot.clone(), auth_endpoint_signature, Some(true), auth_context.balance_remaining, wallet_access, ) .await; refreshed.verified_api_key_hash = verified_api_key_hash; Ok((refreshed, Some(snapshot))) } fn deny_refreshed_auth_context( mut auth_context: GatewayControlAuthContext, ) -> GatewayControlAuthContext { auth_context.access_allowed = false; auth_context.local_rejection = Some(GatewayLocalAuthRejection::InvalidApiKey); auth_context.balance_remaining = None; auth_context } fn put_cached_auth_context( state: &AppState, cache_key: String, auth_context: GatewayControlAuthContext, generation: Option, ) { let (cache_key, ttl) = if is_negative_auth_context(&auth_context) { let ttl = auth_context_negative_cache_ttl(); if ttl.is_zero() { return; } ( negative_auth_context_cache_key(&cache_key), AUTH_CONTEXT_CACHE_TTL.max(ttl), ) } else { (cache_key, AUTH_CONTEXT_CACHE_TTL) }; if let Some(generation) = generation { state.auth_context_cache.insert_if_generation( cache_key, auth_context, ttl, auth_context_cache_max_entries(), &generation, ); } else { state.auth_context_cache.insert( cache_key, auth_context, ttl, auth_context_cache_max_entries(), ); } } fn auth_context_cache_max_entries() -> usize { static MAX_ENTRIES: OnceLock = OnceLock::new(); *MAX_ENTRIES.get_or_init(|| { std::env::var(AUTH_CONTEXT_CACHE_MAX_ENTRIES_ENV) .ok() .and_then(|value| value.trim().parse::().ok()) .filter(|value| *value > 0) .unwrap_or(AUTH_CONTEXT_CACHE_MAX_ENTRIES) }) } fn auth_context_cache_refresh_interval(state: &AppState) -> Duration { #[cfg(test)] if let Some(interval) = state.auth_context_cache.refresh_interval_for_tests() { return interval; } static REFRESH_INTERVAL: OnceLock = OnceLock::new(); *REFRESH_INTERVAL.get_or_init(|| { std::env::var(AUTH_CONTEXT_CACHE_REFRESH_INTERVAL_SECS_ENV) .ok() .and_then(|value| value.trim().parse::().ok()) .filter(|value| *value > 0) .map(Duration::from_secs) .unwrap_or(AUTH_CONTEXT_CACHE_REFRESH_INTERVAL) // Operators may tighten the window, but cannot expand the maximum // authorization staleness beyond the secure default. .min(AUTH_CONTEXT_CACHE_REFRESH_INTERVAL) }) } fn auth_context_cache_refresh_due(state: &AppState, age: Duration) -> bool { age >= auth_context_cache_refresh_interval(state) } fn auth_context_negative_cache_ttl() -> Duration { static NEGATIVE_TTL: OnceLock = OnceLock::new(); *NEGATIVE_TTL.get_or_init(|| { std::env::var(AUTH_CONTEXT_NEGATIVE_CACHE_TTL_SECS_ENV) .ok() .and_then(|value| value.trim().parse::().ok()) .map(Duration::from_secs) .unwrap_or(AUTH_CONTEXT_NEGATIVE_CACHE_TTL) }) } fn negative_auth_context_cache_key(cache_key: &str) -> String { format!("{AUTH_CONTEXT_NEGATIVE_CACHE_KEY_PREFIX}{cache_key}") } fn is_negative_auth_context(auth_context: &GatewayControlAuthContext) -> bool { auth_context.user_id.is_empty() || auth_context.api_key_id.is_empty() || matches!( auth_context.local_rejection, Some(GatewayLocalAuthRejection::InvalidApiKey) ) } fn apply_resolved_auth_context_to_decision( trace_id: &str, decision: &mut GatewayControlDecision, auth_context: GatewayControlAuthContext, ) { log_auth_context_resolution(trace_id, decision, &auth_context); decision.local_auth_rejection = auth_context.local_rejection.clone(); if !auth_context.user_id.is_empty() && !auth_context.api_key_id.is_empty() { decision.auth_context = Some(auth_context); } } pub(super) async fn resolve_data_backed_auth_context( state: &AppState, headers: &http::HeaderMap, uri: &Uri, auth_endpoint_signature: Option<&str>, ) -> Result, GatewayError> { resolve_data_backed_auth_context_with_trusted_auth( state, headers, uri, auth_endpoint_signature, cfg!(test), ) .await } async fn resolve_data_backed_auth_context_with_trusted_auth( state: &AppState, headers: &http::HeaderMap, uri: &Uri, auth_endpoint_signature: Option<&str>, trusted_auth_verified: bool, ) -> Result, GatewayError> { let Some(signature) = auth_endpoint_signature .map(str::trim) .filter(|value| !value.is_empty()) else { return Ok(None); }; if !state.has_auth_api_key_reader() { return Ok(None); } let extracted = extract_request_credentials_with_trusted_auth( headers, uri, signature, trusted_auth_verified, ); let principal = derive_principal_candidate(&extracted); let now_unix_secs = current_unix_secs(); match principal { Some(GatewayPrincipalCandidate::TrustedHeaders(trusted_headers)) => { resolve_trusted_auth_context(state, signature, trusted_headers, now_unix_secs).await } Some(GatewayPrincipalCandidate::ApiKeyHash { key_hash, .. }) => { let snapshot = { let _permit = state.acquire_auth_snapshot_load_gate().await?; state .data .read_auth_api_key_snapshot_by_key_hash_strong(&key_hash, now_unix_secs) .await .map_err(|err| GatewayError::Internal(err.to_string()))? }; let Some(snapshot) = snapshot else { return Ok(Some(GatewayControlAuthContext { user_id: String::new(), api_key_id: String::new(), username: None, api_key_name: None, balance_remaining: None, access_allowed: false, user_rate_limit: None, api_key_rate_limit: None, api_key_is_standalone: false, admin_bypass_limits: false, local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey), allowed_models: None, ip_rules: None, verified_api_key_hash: None, })); }; state .touch_auth_api_key_last_used_best_effort(&snapshot.api_key_id) .await; let wallet_access = resolve_wallet_auth_gate_uncached(state, &snapshot).await?; let mut auth_context = build_data_backed_auth_context( state, snapshot, signature, None, None, wallet_access, ) .await; auth_context.verified_api_key_hash = Some(VerifiedApiKeyHash::new(key_hash)); Ok(Some(auth_context)) } Some(GatewayPrincipalCandidate::DeferredBearerToken { raw, carrier }) => { if let Some(auth_context) = resolve_antigravity_bearer_bridge_auth_context( state, signature, raw.as_str(), carrier, now_unix_secs, ) .await? { return Ok(Some(auth_context)); } Ok(None) } Some(GatewayPrincipalCandidate::DeferredCookieHeader { .. }) => Ok(None), None => Ok(None), } } async fn resolve_antigravity_bearer_bridge_auth_context( state: &AppState, auth_endpoint_signature: &str, raw_bearer: &str, carrier: GatewayCredentialCarrier, now_unix_secs: u64, ) -> Result, GatewayError> { if carrier != GatewayCredentialCarrier::AuthorizationBearer || !auth_endpoint_signature .trim() .eq_ignore_ascii_case("antigravity:v1internal") { return Ok(None); } let config_value = { let _permit = state.acquire_auth_snapshot_load_gate().await?; state .data .find_system_config_value_strong(crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY) .await .map_err(|err| GatewayError::Internal(err.to_string()))? }; let Some(config_value) = config_value else { return Ok(None); }; if config_value.is_null() { return Ok(None); } let config: AntigravityBearerBridgeConfig = serde_json::from_value(config_value).map_err(|err| { GatewayError::Internal(format!( "{} invalid: {err}", crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY )) })?; if !config.enabled { return Ok(None); } let Some(validation_mode) = config.bearer_validation_mode(raw_bearer) else { return Ok(None); }; let user_id = config.auth_user_id.trim(); let api_key_id = config.auth_api_key_id.trim(); if user_id.is_empty() || api_key_id.is_empty() { return Err(GatewayError::Internal(format!( "{} requires auth_user_id and auth_api_key_id", crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY ))); } let snapshot = { let _permit = state.acquire_auth_snapshot_load_gate().await?; state .data .read_auth_api_key_snapshot_strong(user_id, api_key_id, now_unix_secs) .await .map_err(|err| GatewayError::Internal(err.to_string()))? }; let Some(snapshot) = snapshot else { return Ok(Some(GatewayControlAuthContext { user_id: user_id.to_string(), api_key_id: api_key_id.to_string(), username: None, api_key_name: None, balance_remaining: None, access_allowed: false, user_rate_limit: None, api_key_rate_limit: None, api_key_is_standalone: false, admin_bypass_limits: false, local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey), allowed_models: None, ip_rules: None, verified_api_key_hash: None, })); }; let wallet_access = resolve_wallet_auth_gate_uncached(state, &snapshot).await?; let auth_context = build_data_backed_auth_context( state, snapshot, auth_endpoint_signature, None, None, wallet_access, ) .await; info!( event_name = "antigravity_bearer_bridge_auth_context_resolved", log_type = "event", validation_mode, user_id = auth_context.user_id.as_str(), api_key_id = auth_context.api_key_id.as_str(), access_allowed = auth_context.access_allowed, has_local_rejection = auth_context.local_rejection.is_some(), "resolved Antigravity bearer bridge auth context" ); Ok(Some(auth_context)) } async fn resolve_trusted_auth_context( state: &AppState, auth_endpoint_signature: &str, trusted_headers: GatewayTrustedAuthHeaders, now_unix_secs: u64, ) -> Result, GatewayError> { let snapshot = { let _permit = state.acquire_auth_snapshot_load_gate().await?; state .data .read_auth_api_key_snapshot_strong( &trusted_headers.user_id, &trusted_headers.api_key_id, now_unix_secs, ) .await .map_err(|err| GatewayError::Internal(err.to_string()))? }; let Some(snapshot) = snapshot else { return Ok(Some(GatewayControlAuthContext { user_id: trusted_headers.user_id, api_key_id: trusted_headers.api_key_id, username: None, api_key_name: None, balance_remaining: trusted_headers.balance_remaining, access_allowed: false, user_rate_limit: None, api_key_rate_limit: None, api_key_is_standalone: false, admin_bypass_limits: false, local_rejection: Some(GatewayLocalAuthRejection::InvalidApiKey), allowed_models: None, ip_rules: None, verified_api_key_hash: None, })); }; let wallet_access = resolve_wallet_auth_gate_uncached(state, &snapshot).await?; Ok(Some( build_data_backed_auth_context( state, snapshot, auth_endpoint_signature, trusted_headers.access_allowed, trusted_headers.balance_remaining, wallet_access, ) .await, )) } async fn build_data_backed_auth_context( state: &AppState, snapshot: crate::data::auth::GatewayAuthApiKeySnapshot, auth_endpoint_signature: &str, header_access_allowed: Option, balance_remaining: Option, wallet_access: Option, ) -> GatewayControlAuthContext { let allowed_models = snapshot .effective_allowed_models() .map(|items| items.to_vec()); let invalid_api_key = !snapshot.user_is_active || snapshot.user_is_deleted || !snapshot.api_key_is_active || api_key_is_expired(snapshot.api_key_expires_at_unix_secs, current_unix_secs()); let locked_api_key = snapshot.api_key_is_locked && !snapshot.api_key_is_standalone; let key_access_allowed = header_access_allowed .map(|value| value && snapshot.currently_usable) .unwrap_or(snapshot.currently_usable); let wallet_remaining = wallet_access .as_ref() .and_then(|decision| decision.remaining); let requested_provider = auth_endpoint_signature .split_once(':') .map(|(provider, _)| provider) .unwrap_or(auth_endpoint_signature) .trim(); let identity_only = auth_gate_identity_only(auth_endpoint_signature); let requested_provider_allowed = identity_only || auth_snapshot_allows_requested_provider(state, &snapshot, auth_endpoint_signature).await; let local_rejection = if invalid_api_key { Some(GatewayLocalAuthRejection::InvalidApiKey) } else if locked_api_key { Some(GatewayLocalAuthRejection::LockedApiKey) } else if let Some(rejection) = wallet_access .as_ref() .and_then(local_rejection_from_wallet_access) { Some(rejection) } else if header_access_allowed.is_some_and(|value| !value) && snapshot.currently_usable { Some(GatewayLocalAuthRejection::BalanceDenied { remaining: balance_remaining.or(wallet_remaining), }) } else if !requested_provider.is_empty() && !requested_provider_allowed { Some(GatewayLocalAuthRejection::ProviderNotAllowed { provider: requested_provider.to_string(), }) } else if !identity_only && snapshot .effective_allowed_api_formats() .is_some_and(|allowed| { !contains_api_format_or_alias( allowed, auth_gate_api_format(auth_endpoint_signature).as_str(), ) }) { Some(GatewayLocalAuthRejection::ApiFormatNotAllowed { api_format: auth_endpoint_signature.to_string(), }) } else { None }; GatewayControlAuthContext { username: Some(snapshot.username.clone()), api_key_name: snapshot.api_key_name.clone(), user_id: snapshot.user_id, api_key_id: snapshot.api_key_id, balance_remaining: wallet_remaining.or(balance_remaining), access_allowed: key_access_allowed && local_rejection.is_none(), user_rate_limit: snapshot.user_rate_limit, api_key_rate_limit: snapshot.api_key_rate_limit, api_key_is_standalone: snapshot.api_key_is_standalone, admin_bypass_limits: snapshot.user_role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone, local_rejection, allowed_models, ip_rules: snapshot.api_key_ip_rules, verified_api_key_hash: None, } } fn api_key_is_expired(expires_at_unix_secs: Option, now_unix_secs: u64) -> bool { expires_at_unix_secs.is_some_and(|expires_at| expires_at <= now_unix_secs) } fn contains_api_format_or_alias(items: &[String], target: &str) -> bool { items.iter().any(|item| api_format_matches(item, target)) } fn normalize_api_format_alias(value: &str) -> String { crate::ai_serving::normalize_api_format_alias(value) } fn auth_gate_api_format(auth_endpoint_signature: &str) -> String { let normalized = normalize_api_format_alias(auth_endpoint_signature); match normalized.as_str() { "antigravity:v1internal" => "gemini:generate_content".to_string(), _ => normalized, } } fn auth_gate_identity_only(auth_endpoint_signature: &str) -> bool { matches!( auth_endpoint_signature.trim().to_ascii_lowercase().as_str(), "aether:ccswitch_usage" ) } fn api_format_matches(left: &str, right: &str) -> bool { aether_scheduler_core::api_format_matches_allowed_value(left, right) } async fn auth_snapshot_allows_requested_provider( state: &AppState, snapshot: &crate::data::auth::GatewayAuthApiKeySnapshot, auth_endpoint_signature: &str, ) -> bool { let Some(allowed_providers) = snapshot.effective_allowed_providers() else { return true; }; let requested_api_format = normalize_api_format_alias(auth_endpoint_signature); let requested_provider = requested_api_format .split_once(':') .map(|(provider, _)| provider) .unwrap_or(requested_api_format.as_str()) .trim(); if requested_provider.is_empty() { return true; } if allowed_providers.is_empty() { return false; } if allowed_providers .iter() .any(|value| allowed_provider_value_matches_requested_provider(value, requested_provider)) { return true; } if !state.has_provider_catalog_data_reader() { debug!( "deny requested provider {}: provider catalog is unavailable for allowlist resolution", requested_provider ); return false; } let providers = match state.list_provider_catalog_providers(true).await { Ok(value) => value, Err(err) => { warn!( "deny requested provider {}: provider catalog lookup failed: {:?}", requested_provider, err ); return false; } }; let allowed_catalog_providers = providers .into_iter() .filter(|provider| { allowed_providers.iter().any(|value| { aether_scheduler_core::provider_matches_allowed_value( value, &provider.id, &provider.name, &provider.provider_type, ) }) }) .collect::>(); if allowed_catalog_providers .iter() .any(|provider| provider_matches_requested_provider(provider, requested_provider)) { return true; } let allowed_provider_ids = allowed_catalog_providers .iter() .map(|provider| provider.id.clone()) .collect::>(); if allowed_provider_ids.is_empty() { return false; } let endpoints = match state .list_provider_catalog_endpoints_by_provider_ids(&allowed_provider_ids) .await { Ok(value) => value, Err(err) => { warn!( "deny requested provider {}: provider endpoint lookup failed: {:?}", requested_provider, err ); return false; } }; endpoints.iter().any(|endpoint| { endpoint_matches_requested_provider(endpoint, &requested_api_format, requested_provider) }) } fn allowed_provider_value_matches_requested_provider( allowed_value: &str, requested_provider: &str, ) -> bool { aether_scheduler_core::provider_matches_allowed_value( allowed_value, requested_provider, requested_provider, requested_provider, ) } fn provider_matches_requested_provider( provider: &StoredProviderCatalogProvider, requested_provider: &str, ) -> bool { aether_scheduler_core::provider_matches_allowed_value( requested_provider, &provider.id, &provider.name, &provider.provider_type, ) } fn endpoint_matches_requested_provider( endpoint: &StoredProviderCatalogEndpoint, requested_api_format: &str, requested_provider: &str, ) -> bool { if !endpoint.is_active { return false; } if api_format_matches(&endpoint.api_format, requested_api_format) { return true; } let endpoint_api_format = normalize_api_format_alias(&endpoint.api_format); if crate::ai_serving::request_conversion_kind(requested_api_format, &endpoint_api_format) .is_some() { return true; } if endpoint.api_family.as_deref().is_some_and(|family| { allowed_provider_value_matches_requested_provider(family, requested_provider) }) { return true; } let endpoint_provider = endpoint_api_format .split_once(':') .map(|(provider, _)| provider) .unwrap_or(endpoint_api_format.as_str()); allowed_provider_value_matches_requested_provider(endpoint_provider, requested_provider) } fn get_cached_auth_context(state: &AppState, cache_key: &str) -> Option { get_cached_auth_context_with_age(state, cache_key).map(|(auth_context, _)| auth_context) } fn get_cached_auth_context_with_age( state: &AppState, cache_key: &str, ) -> Option<(GatewayControlAuthContext, Duration)> { let negative_ttl = auth_context_negative_cache_ttl(); if !negative_ttl.is_zero() { if let Some(auth_context) = state .auth_context_cache .get_fresh_with_age(&negative_auth_context_cache_key(cache_key), negative_ttl) { return Some(auth_context); } } state .auth_context_cache .get_fresh_with_age(cache_key, AUTH_CONTEXT_CACHE_TTL) } #[cfg(test)] mod tests { use std::sync::Arc; use std::time::Duration; use aether_data::repository::auth::{ AuthApiKeyWriteRepository, CreateUserApiKeyRecord, InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, }; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::repository::wallet::{ InMemoryWalletRepository, StoredWalletSnapshot, WalletReadRepository, }; use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogProvider, }; use aether_runtime::ConcurrencyGate; use axum::http::{HeaderMap, Uri}; use futures_util::future::join_all; use super::{ api_key_is_expired, get_cached_auth_context, refresh_execution_runtime_auth_context_with_snapshot, resolve_control_decision_auth, resolve_data_backed_auth_context, resolve_execution_runtime_auth_context, ControlDecisionAuthResolution, GatewayLocalAuthRejection, }; use crate::control::auth::credentials::{build_auth_context_cache_key, hash_api_key}; use crate::control::GatewayControlDecision; use crate::data::{GatewayDataConfig, GatewayDataState}; use crate::AppState; fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot { StoredAuthApiKeySnapshot::new( user_id.to_string(), "alice".to_string(), Some("alice@example.com".to_string()), "user".to_string(), "local".to_string(), true, false, Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:chat"])), Some(serde_json::json!(["gpt-4.1"])), api_key_id.to_string(), Some("default".to_string()), true, false, false, Some(60), Some(5), Some(4_102_444_800), Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:chat"])), Some(serde_json::json!(["gpt-4.1"])), ) .expect("snapshot should build") } fn uri(path: &str) -> Uri { path.parse().expect("uri should parse") } #[test] fn api_key_expiry_is_inclusive_at_the_declared_second() { assert!(!api_key_is_expired(None, 100)); assert!(!api_key_is_expired(Some(101), 100)); assert!(api_key_is_expired(Some(100), 100)); assert!(api_key_is_expired(Some(99), 100)); } fn sample_provider(id: &str, name: &str, provider_type: &str) -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( id.to_string(), name.to_string(), None, provider_type.to_string(), ) .expect("provider should build") } fn sample_endpoint( id: &str, provider_id: &str, api_format: &str, ) -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( id.to_string(), provider_id.to_string(), api_format.to_string(), None, None, true, ) .expect("endpoint should build") } struct PostgresAuthConfigNodes { first: AppState, second: AppState, _database: aether_testkit::ManagedPostgresServer, } async fn postgres_auth_config_nodes( auth_repository: Arc, ) -> PostgresAuthConfigNodes { let server = aether_testkit::ManagedPostgresServer::start() .await .expect("temporary PostgreSQL should start"); let mut pool = SqlPoolConfig::default(); pool.min_connections = 0; pool.max_connections = 4; let database = SqlDatabaseConfig::new(DatabaseDriver::Postgres, server.database_url(), pool) .expect("postgres config should build"); let config = GatewayDataConfig::from_database_config(database); let first_data = GatewayDataState::from_config(config.clone()) .expect("first data state should build") .with_auth_api_key_reader(auth_repository.clone()) .without_wallet_reader_for_tests(); assert!(first_data .run_database_migrations() .await .expect("postgres migrations should run")); let second_data = GatewayDataState::from_config(config) .expect("second data state should build") .with_auth_api_key_reader(auth_repository) .without_wallet_reader_for_tests(); PostgresAuthConfigNodes { first: AppState::new() .expect("first app state should build") .with_data_state_for_tests(first_data), second: AppState::new() .expect("second app state should build") .with_data_state_for_tests(second_data), _database: server, } } #[tokio::test] async fn strong_system_config_read_bypasses_app_and_data_caches() { let nodes = postgres_auth_config_nodes(Arc::new(InMemoryAuthApiKeySnapshotRepository::seed([]))) .await; let key = format!("test.auth.strong-read.{}", uuid::Uuid::new_v4()); let old_value = serde_json::json!({"version": "old"}); let new_value = serde_json::json!({"version": "new"}); nodes .first .data .upsert_system_config_value(&key, &old_value, None) .await .expect("initial config should write"); assert_eq!( nodes .first .read_system_config_json_value(&key) .await .expect("initial app config read should succeed"), Some(old_value.clone()) ); nodes .second .data .upsert_system_config_value(&key, &new_value, None) .await .expect("cross-node config update should write"); assert_eq!( nodes .first .read_system_config_json_value(&key) .await .expect("stale app config read should succeed"), Some(old_value.clone()), "node-local AppState cache should still contain the old value" ); assert_eq!( nodes .first .data .find_system_config_value(&key) .await .expect("stale data config read should succeed"), Some(old_value), "node-local data cache should still contain the old value" ); assert_eq!( nodes .first .data .find_system_config_value_strong(&key) .await .expect("strong config read should succeed"), Some(new_value), "strong reads must reach the shared repository" ); } #[tokio::test] async fn control_auth_caches_invalid_api_key_rejections() { let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed([])); let data = GatewayDataState::with_auth_api_key_repository_for_tests(repository); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, "Bearer sk-missing-for-negative-cache".parse().unwrap(), ); let request_uri = uri("/v1/chat/completions"); let decision = GatewayControlDecision::synthetic( "/v1/chat/completions", Some("ai_public".to_string()), Some("openai".to_string()), Some("chat".to_string()), Some("openai:chat".to_string()), ); let ControlDecisionAuthResolution::Resolved(first) = resolve_control_decision_auth( &state, &headers, &request_uri, "trace-invalid-auth-cache", decision, ) .await .expect("auth resolution should succeed"); assert_eq!( first.local_auth_rejection, Some(GatewayLocalAuthRejection::InvalidApiKey) ); let cache_key = build_auth_context_cache_key(&headers, &request_uri, "openai:chat") .expect("cache key should exist"); let cached = get_cached_auth_context(&state, &cache_key) .expect("invalid API key rejection should be cached"); assert_eq!( cached.local_rejection, Some(GatewayLocalAuthRejection::InvalidApiKey) ); assert!(cached.user_id.is_empty()); assert!(cached.api_key_id.is_empty()); } #[tokio::test] async fn uncached_invalid_api_key_lookup_waits_for_auth_snapshot_gate() { let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed([])); let data = GatewayDataState::with_auth_api_key_repository_for_tests(repository.clone()); let mut state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); state.auth_snapshot_load_gate = Some(Arc::new(ConcurrencyGate::new( "test_invalid_auth_snapshot_lookup", 1, ))); let held = state .acquire_auth_snapshot_load_gate() .await .expect("auth gate acquisition should succeed") .expect("auth gate should be configured"); let api_key = "sk-missing-auth-gate"; let mut headers = HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, format!("Bearer {api_key}").parse().unwrap(), ); let request_uri = uri("/v1/chat/completions"); let blocked = tokio::time::timeout( Duration::from_millis(25), resolve_data_backed_auth_context(&state, &headers, &request_uri, Some("openai:chat")), ) .await; assert!( blocked.is_err(), "an uncached invalid-key strong read must wait for auth DB admission" ); assert_eq!(repository.key_hash_lookup_count(&hash_api_key(api_key)), 0); drop(held); let resolved = tokio::time::timeout( Duration::from_secs(1), resolve_data_backed_auth_context(&state, &headers, &request_uri, Some("openai:chat")), ) .await .expect("auth lookup should resume after releasing the gate") .expect("auth resolution should succeed") .expect("invalid API key should resolve to a rejection context"); assert_eq!( resolved.local_rejection, Some(GatewayLocalAuthRejection::InvalidApiKey) ); assert_eq!(repository.key_hash_lookup_count(&hash_api_key(api_key)), 1); } #[tokio::test] async fn data_backed_api_key_auth_touches_last_used_once_per_throttle_window() { let api_key = "sk-test-touch"; let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), sample_snapshot("key-1", "user-1"), )])); let data = GatewayDataState::with_auth_api_key_repository_for_tests(repository.clone()); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, format!("Bearer {api_key}").parse().unwrap(), ); let first = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/chat/completions"), Some("openai:chat"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(first.user_id, "user-1"); assert_eq!(first.api_key_id, "key-1"); assert_eq!(repository.touch_count("key-1"), 1); let second = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/chat/completions"), Some("openai:chat"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(second.api_key_id, "key-1"); assert_eq!(repository.touch_count("key-1"), 1); } #[tokio::test] async fn long_lived_refresh_rejects_same_ids_recreated_with_a_different_credential() { let old_api_key = "sk-old-websocket-credential"; let new_api_key = "sk-new-websocket-credential"; let old_key_hash = hash_api_key(old_api_key); let new_key_hash = hash_api_key(new_api_key); let mut old_snapshot = sample_snapshot("key-stable-id", "user-stable-id"); old_snapshot.user_allowed_api_formats = Some(vec!["openai:responses".to_string()]); old_snapshot.api_key_allowed_api_formats = Some(vec!["openai:responses".to_string()]); let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(old_key_hash.clone()), old_snapshot, )])); let data = GatewayDataState::with_auth_api_key_repository_for_tests(repository.clone()); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, format!("Bearer {old_api_key}").parse().unwrap(), ); let original = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/responses"), Some("openai:responses"), ) .await .expect("initial auth resolution should succeed") .expect("the old API key should authenticate"); assert!(original.access_allowed); assert!(original.verified_api_key_hash.is_some()); assert!( !format!("{original:?}").contains(&old_key_hash), "the credential verifier must stay redacted from Debug output" ); assert!(repository .delete_user_api_key("user-stable-id", "key-stable-id") .await .expect("old API key deletion should succeed")); repository .create_user_api_key(CreateUserApiKeyRecord { user_id: "user-stable-id".to_string(), api_key_id: "key-stable-id".to_string(), key_hash: new_key_hash, key_encrypted: None, name: Some("restored-with-new-secret".to_string()), allowed_providers: Some(vec!["openai".to_string()]), allowed_api_formats: Some(vec!["openai:responses".to_string()]), allowed_models: Some(vec!["gpt-4.1".to_string()]), ip_rules: None, rate_limit: 60, concurrent_limit: Some(5), force_capabilities: None, feature_settings: None, is_active: true, expires_at_unix_secs: Some(4_102_444_800), auto_delete_on_expiry: false, total_requests: 0, total_tokens: 0, total_cost_usd: 0.0, }) .await .expect("same-ID API key recreation should resolve") .expect("same-ID API key recreation should persist"); let (refreshed, snapshot) = refresh_execution_runtime_auth_context_with_snapshot( &state, original, Some("openai:responses"), ) .await .expect("long-lived auth refresh should resolve"); assert!(!refreshed.access_allowed); assert_eq!( refreshed.local_rejection, Some(GatewayLocalAuthRejection::InvalidApiKey) ); assert!(snapshot.is_none()); assert_eq!(repository.key_hash_lookup_count(&old_key_hash), 1); assert_eq!( repository.snapshot_lookup_count("key-stable-id"), 0, "a bound long-lived credential must not fall back to identity-only lookup" ); } #[tokio::test] async fn control_auth_context_singleflights_concurrent_cache_misses() { let api_key = "sk-test-concurrent-auth-miss"; let repository = Arc::new( InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), sample_snapshot("key-concurrent-auth-miss", "user-concurrent-auth-miss"), )]) .with_lookup_delay_for_tests(Duration::from_millis(20)), ); let data = GatewayDataState::with_auth_api_key_repository_for_tests(repository.clone()); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, format!("Bearer {api_key}").parse().unwrap(), ); let request_uri = uri("/v1/chat/completions"); let tasks = (0..32).map(|index| { let decision = GatewayControlDecision::synthetic( "/v1/chat/completions", Some("ai_public".to_string()), Some("openai".to_string()), Some("chat".to_string()), Some("openai:chat".to_string()), ); let trace_id = format!("trace-concurrent-auth-miss-{index}"); let state = &state; let headers = &headers; let request_uri = &request_uri; async move { resolve_control_decision_auth(state, headers, request_uri, &trace_id, decision) .await } }); for result in join_all(tasks).await { let ControlDecisionAuthResolution::Resolved(decision) = result.expect("auth resolution should succeed"); let auth_context = decision .auth_context .expect("auth context should be resolved"); assert_eq!(auth_context.user_id, "user-concurrent-auth-miss"); assert_eq!(auth_context.api_key_id, "key-concurrent-auth-miss"); } assert_eq!( repository.key_hash_lookup_count(&hash_api_key(api_key)), 1, "concurrent cache misses for one auth context should only load one snapshot" ); } #[tokio::test] async fn data_backed_auth_context_marks_wallet_denial_as_not_allowed() { let api_key = "sk-test-empty-wallet"; let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), sample_snapshot("key-empty-wallet", "user-empty-wallet"), )])); let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![ StoredWalletSnapshot::new( "wallet-empty".to_string(), Some("user-empty-wallet".to_string()), None, 0.0, 0.0, "finite".to_string(), "USD".to_string(), "active".to_string(), 0.0, 0.0, 0.0, 0.0, 100, ) .expect("wallet should build"), ])); let data = GatewayDataState::with_auth_and_wallet_for_tests(auth_repository, wallet_repository); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, format!("Bearer {api_key}").parse().unwrap(), ); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/chat/completions"), Some("openai:chat"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!( auth_context.local_rejection, Some(GatewayLocalAuthRejection::BalanceDenied { remaining: Some(0.0), }) ); assert!(!auth_context.access_allowed); } #[tokio::test] async fn execution_runtime_auth_context_revalidates_cached_wallet_state() { let api_key = "sk-test-runtime-wallet-cache"; let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), sample_snapshot("key-runtime-wallet-cache", "user-runtime-wallet-cache"), )])); let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![ StoredWalletSnapshot::new( "wallet-runtime-cache".to_string(), Some("user-runtime-wallet-cache".to_string()), None, 10.0, 0.0, "finite".to_string(), "USD".to_string(), "active".to_string(), 10.0, 0.0, 0.0, 0.0, 100, ) .expect("wallet should build"), ])); let data = GatewayDataState::with_auth_and_wallet_for_tests( auth_repository.clone(), Arc::clone(&wallet_repository), ); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let decision = GatewayControlDecision::synthetic( "/v1/chat/completions", Some("ai_public".to_string()), Some("openai".to_string()), Some("chat".to_string()), Some("openai:chat".to_string()), ); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let first = resolve_execution_runtime_auth_context( &state, &decision, &headers, &uri("/v1/chat/completions"), "trace-runtime-wallet-cache", ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert!(first.access_allowed); state .auth_context_cache .set_refresh_interval_for_tests(Duration::from_millis(10)); wallet_repository .update_auth_user_wallet_snapshot( "user-runtime-wallet-cache", 0.0, 0.0, "finite", "USD", "active", 10.0, 10.0, 0.0, 0.0, Some(101), ) .await .expect("wallet update should succeed") .expect("wallet should exist"); tokio::time::sleep(Duration::from_millis(15)).await; let second = tokio::time::timeout( Duration::from_millis(100), resolve_execution_runtime_auth_context( &state, &decision, &headers, &uri("/v1/chat/completions"), "trace-runtime-wallet-cache", ), ) .await .expect("due security refresh should complete") .expect("resolution should succeed") .expect("auth context should exist"); assert!( !second.access_allowed, "the due request must observe denial" ); let cache_key = build_auth_context_cache_key(&headers, &uri("/v1/chat/completions"), "openai:chat") .expect("cache key should exist"); let refreshed = get_cached_auth_context(&state, &cache_key) .expect("synchronous refresh should publish the wallet denial"); assert_eq!( refreshed.local_rejection, Some(GatewayLocalAuthRejection::BalanceDenied { remaining: Some(0.0), }), "auth refresh should publish current wallet state" ); assert!(!refreshed.access_allowed); assert_eq!( auth_repository.key_hash_lookup_count(&hash_api_key(api_key)), 2, "initial auth and due refresh must each validate the presented credential" ); assert_eq!( auth_repository.snapshot_lookup_count("key-runtime-wallet-cache"), 0, "due refresh must not trust the cached user/key mapping" ); } #[tokio::test] async fn due_auth_context_refresh_observes_cross_node_key_lock() { let api_key = "sk-test-cross-node-lock"; let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), sample_snapshot("key-cross-node-lock", "user-cross-node-lock"), )])); let data = GatewayDataState::with_auth_api_key_repository_for_tests(auth_repository.clone()); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); state .auth_context_cache .set_refresh_interval_for_tests(Duration::from_millis(10)); let decision = GatewayControlDecision::synthetic( "/v1/chat/completions", Some("ai_public".to_string()), Some("openai".to_string()), Some("chat".to_string()), Some("openai:chat".to_string()), ); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let request_uri = uri("/v1/chat/completions"); let first = resolve_execution_runtime_auth_context( &state, &decision, &headers, &request_uri, "trace-cross-node-lock-prime", ) .await .expect("initial resolution should succeed") .expect("auth context should exist"); assert!(first.access_allowed); assert!(auth_repository .set_user_api_key_locked("user-cross-node-lock", "key-cross-node-lock", true) .await .expect("simulated cross-node lock should succeed")); tokio::time::sleep(Duration::from_millis(15)).await; let refreshed = resolve_execution_runtime_auth_context( &state, &decision, &headers, &request_uri, "trace-cross-node-lock-refresh", ) .await .expect("due resolution should succeed") .expect("auth context should exist"); assert!(!refreshed.access_allowed); assert_eq!( refreshed.local_rejection, Some(GatewayLocalAuthRejection::LockedApiKey) ); } #[tokio::test] async fn auth_context_hard_miss_bypasses_fresh_snapshot_allow_caches() { let api_key = "sk-test-hard-miss-cross-node-lock"; let key_hash = hash_api_key(api_key); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(key_hash.clone()), sample_snapshot("key-hard-miss-lock", "user-hard-miss-lock"), )])); let data = GatewayDataState::with_cached_auth_api_key_repository_for_tests( auth_repository.clone(), ); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let primed = state .read_cached_auth_api_key_snapshot_by_key_hash(&key_hash, super::current_unix_secs()) .await .expect("snapshot cache prime should succeed") .expect("snapshot should exist"); assert!(primed.currently_usable); assert!(auth_repository .set_user_api_key_locked("user-hard-miss-lock", "key-hard-miss-lock", true) .await .expect("simulated cross-node lock should succeed")); let decision = GatewayControlDecision::synthetic( "/v1/chat/completions", Some("ai_public".to_string()), Some("openai".to_string()), Some("chat".to_string()), Some("openai:chat".to_string()), ); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let resolved = resolve_execution_runtime_auth_context( &state, &decision, &headers, &uri("/v1/chat/completions"), "trace-hard-miss-cross-node-lock", ) .await .expect("hard-miss resolution should succeed") .expect("auth context should exist"); assert!(!resolved.access_allowed); assert_eq!( resolved.local_rejection, Some(GatewayLocalAuthRejection::LockedApiKey) ); assert_eq!( auth_repository.key_hash_lookup_count(&key_hash), 2, "hard miss must bypass both fresh snapshot cache layers" ); } #[tokio::test] async fn execution_runtime_auth_context_singleflights_concurrent_cache_refreshes() { let api_key = "sk-test-runtime-auth-refresh"; let auth_repository = Arc::new( InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), sample_snapshot("key-runtime-auth-refresh", "user-runtime-auth-refresh"), )]) .with_lookup_delay_for_tests(Duration::from_millis(200)), ); let data = GatewayDataState::with_auth_api_key_repository_for_tests(auth_repository.clone()); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let decision = GatewayControlDecision::synthetic( "/v1/chat/completions", Some("ai_public".to_string()), Some("openai".to_string()), Some("chat".to_string()), Some("openai:chat".to_string()), ); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let request_uri = uri("/v1/chat/completions"); let first = resolve_execution_runtime_auth_context( &state, &decision, &headers, &request_uri, "trace-runtime-auth-refresh-prime", ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(first.api_key_id, "key-runtime-auth-refresh"); assert_eq!( auth_repository.key_hash_lookup_count(&hash_api_key(api_key)), 1 ); state .auth_context_cache .set_refresh_interval_for_tests(Duration::from_millis(100)); tokio::time::sleep(Duration::from_millis(110)).await; let tasks = (0..32).map(|index| { let trace_id = format!("trace-runtime-auth-refresh-{index}"); let state = &state; let decision = &decision; let headers = &headers; let request_uri = &request_uri; async move { resolve_execution_runtime_auth_context( state, decision, headers, request_uri, &trace_id, ) .await } }); let results = tokio::time::timeout(Duration::from_secs(1), join_all(tasks)) .await .expect("same-key security refreshes should complete through one flight"); for result in results { let auth_context = result .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(auth_context.user_id, "user-runtime-auth-refresh"); assert_eq!(auth_context.api_key_id, "key-runtime-auth-refresh"); } assert_eq!( auth_repository.key_hash_lookup_count(&hash_api_key(api_key)), 2, "one due security refresh should revalidate the original API key hash" ); assert_eq!( auth_repository.snapshot_lookup_count("key-runtime-auth-refresh"), 0, "full credential revalidation must not trust the cached user/key mapping" ); } #[tokio::test] async fn execution_runtime_auth_context_reuses_control_resolved_context() { let api_key = "sk-test-control-execution-auth-reuse"; let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), sample_snapshot( "key-control-execution-auth-reuse", "user-control-execution-auth-reuse", ), )])); let data = GatewayDataState::with_auth_api_key_repository_for_tests(auth_repository.clone()); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let request_uri = uri("/v1/chat/completions"); let decision = GatewayControlDecision::synthetic( "/v1/chat/completions", Some("ai_public".to_string()), Some("openai".to_string()), Some("chat".to_string()), Some("openai:chat".to_string()), ) .with_execution_runtime_candidate(true); let ControlDecisionAuthResolution::Resolved(first) = resolve_control_decision_auth( &state, &headers, &request_uri, "trace-control-execution-auth-reuse-first", decision, ) .await .expect("control auth resolution should succeed"); assert!(first.auth_context.is_some()); assert_eq!( auth_repository.snapshot_lookup_count("key-control-execution-auth-reuse"), 0, "initial auth resolution should use the key-hash lookup" ); let mut second = None; for index in 0..32 { let decision = GatewayControlDecision::synthetic( "/v1/chat/completions", Some("ai_public".to_string()), Some("openai".to_string()), Some("chat".to_string()), Some("openai:chat".to_string()), ) .with_execution_runtime_candidate(true); let ControlDecisionAuthResolution::Resolved(resolved) = resolve_control_decision_auth( &state, &headers, &request_uri, &format!("trace-control-execution-auth-reuse-{index}"), decision, ) .await .expect("control auth resolution should succeed"); second = Some(resolved); } let second = second.expect("fresh cache hit should resolve a decision"); assert!(second.auth_context.is_some()); let snapshot_lookups_after_control = auth_repository.snapshot_lookup_count("key-control-execution-auth-reuse"); assert_eq!( snapshot_lookups_after_control, 0, "continuous fresh hits before soft TTL must not refresh auth" ); let execution = resolve_execution_runtime_auth_context( &state, &second, &headers, &request_uri, "trace-control-execution-auth-reuse-execution", ) .await .expect("execution auth resolution should succeed") .expect("execution auth context should exist"); assert_eq!(execution.api_key_id, "key-control-execution-auth-reuse"); let control_context = second .auth_context .as_ref() .expect("control auth context should exist"); assert_eq!( execution.access_allowed, control_context.access_allowed, "execution must preserve the control-stage access decision" ); assert_eq!( execution.local_rejection, control_context.local_rejection, "execution must preserve the control-stage local rejection" ); assert_eq!( auth_repository.snapshot_lookup_count("key-control-execution-auth-reuse"), snapshot_lookups_after_control, "execution should reuse the control-resolved auth context without another snapshot lookup" ); } #[tokio::test] async fn data_backed_auth_context_allows_provider_id_for_matching_provider_type() { let api_key = "sk-test-provider-id"; let mut snapshot = sample_snapshot("key-2", "user-2"); snapshot.user_allowed_providers = Some(vec!["provider-openai-1".to_string()]); snapshot.api_key_allowed_providers = Some(vec!["provider-openai-1".to_string()]); let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), snapshot, )])); let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider( "provider-openai-1", "OpenAI Pool 1", "openai", )], Vec::new(), Vec::new(), )); let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository) .with_provider_catalog_reader(provider_catalog); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/chat/completions"), Some("openai:chat"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(auth_context.local_rejection, None); } #[tokio::test] async fn data_backed_auth_context_allows_provider_id_for_matching_endpoint_format() { let api_key = "sk-test-provider-endpoint"; let mut snapshot = sample_snapshot("key-4", "user-4"); snapshot.user_allowed_providers = Some(vec!["provider-custom-claude".to_string()]); snapshot.api_key_allowed_providers = Some(vec!["provider-custom-claude".to_string()]); snapshot.user_allowed_api_formats = None; snapshot.api_key_allowed_api_formats = None; let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), snapshot, )])); let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider( "provider-custom-claude", "Custom Claude Gateway", "custom", )], vec![sample_endpoint( "endpoint-custom-claude", "provider-custom-claude", "claude:messages", )], Vec::new(), )); let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository) .with_provider_catalog_reader(provider_catalog); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/messages"), Some("claude:messages"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(auth_context.local_rejection, None); } #[tokio::test] async fn data_backed_auth_context_denies_unresolved_provider_id_without_catalog_reader() { let api_key = "sk-test-provider-no-catalog"; let mut snapshot = sample_snapshot("key-no-catalog", "user-no-catalog"); snapshot.user_allowed_providers = Some(vec!["provider-custom-claude".to_string()]); snapshot.api_key_allowed_providers = Some(vec!["provider-custom-claude".to_string()]); snapshot.user_allowed_api_formats = None; snapshot.api_key_allowed_api_formats = None; let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), snapshot, )])); let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/messages"), Some("claude:messages"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert!(!auth_context.access_allowed); assert_eq!( auth_context.local_rejection, Some(GatewayLocalAuthRejection::ProviderNotAllowed { provider: "claude".to_string(), }) ); } #[tokio::test] async fn due_antigravity_bearer_refresh_observes_cross_node_allowlist_revocation() { let raw_bearer = "google-oauth-access-token-revoked-cross-node"; let mut snapshot = sample_snapshot( "key-antigravity-bearer-revocation", "user-antigravity-bearer-revocation", ); snapshot.user_allowed_providers = Some(vec!["antigravity".to_string()]); snapshot.api_key_allowed_providers = Some(vec!["antigravity".to_string()]); snapshot.user_allowed_api_formats = Some(vec!["gemini:generate_content".to_string()]); snapshot.api_key_allowed_api_formats = Some(vec!["gemini:generate_content".to_string()]); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( None, snapshot, )])); let nodes = postgres_auth_config_nodes(auth_repository.clone()).await; nodes .first .auth_context_cache .set_refresh_interval_for_tests(Duration::from_millis(10)); let allowed_config = serde_json::json!({ "enabled": true, "auth_user_id": "user-antigravity-bearer-revocation", "auth_api_key_id": "key-antigravity-bearer-revocation", "bearer_sha256_allowlist": [hash_api_key(raw_bearer)] }); nodes .first .data .upsert_system_config_value( crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY, &allowed_config, None, ) .await .expect("initial bearer bridge config should write"); assert_eq!( nodes .first .read_system_config_json_value( crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY, ) .await .expect("initial bridge config read should succeed"), Some(allowed_config.clone()) ); let decision = GatewayControlDecision::synthetic( "/v1internal:streamGenerateContent", Some("ai_public".to_string()), Some("antigravity".to_string()), Some("v1internal".to_string()), Some("antigravity:v1internal".to_string()), ); let mut headers = HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, format!("Bearer {raw_bearer}").parse().unwrap(), ); let request_uri = uri("/v1internal:streamGenerateContent?alt=sse"); let initial = resolve_execution_runtime_auth_context( &nodes.first, &decision, &headers, &request_uri, "trace-antigravity-bearer-revocation-prime", ) .await .expect("initial bearer auth should resolve") .expect("initial bearer auth context should exist"); assert!(initial.access_allowed); let revoked_config = serde_json::json!({ "enabled": true, "auth_user_id": "user-antigravity-bearer-revocation", "auth_api_key_id": "key-antigravity-bearer-revocation", "bearer_sha256_allowlist": [hash_api_key("different-bearer")] }); nodes .second .data .upsert_system_config_value( crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY, &revoked_config, None, ) .await .expect("cross-node bearer revocation should write"); assert_eq!( nodes .first .read_system_config_json_value( crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY, ) .await .expect("stale app config read should succeed"), Some(allowed_config), "the node-local AppState cache should remain stale for the regression setup" ); assert_eq!( nodes .first .data .find_system_config_value(crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY,) .await .expect("stale data config read should succeed"), Some(serde_json::json!({ "enabled": true, "auth_user_id": "user-antigravity-bearer-revocation", "auth_api_key_id": "key-antigravity-bearer-revocation", "bearer_sha256_allowlist": [hash_api_key(raw_bearer)] })), "the node-local data cache should remain stale for the regression setup" ); tokio::time::sleep(Duration::from_millis(15)).await; let refreshed = resolve_execution_runtime_auth_context( &nodes.first, &decision, &headers, &request_uri, "trace-antigravity-bearer-revocation-refresh", ) .await .expect("due bearer auth refresh should resolve") .expect("due bearer auth context should exist"); assert!(!refreshed.access_allowed); assert_eq!( refreshed.local_rejection, Some(GatewayLocalAuthRejection::InvalidApiKey) ); assert_eq!( auth_repository.snapshot_lookup_count("key-antigravity-bearer-revocation"), 1, "a revoked bearer must be rejected before its mapped API key is reloaded" ); } #[tokio::test] async fn data_backed_auth_context_allows_antigravity_v1internal_for_gemini_generate_content_keys( ) { let api_key = "sk-test-antigravity-v1internal"; let mut snapshot = sample_snapshot("key-ant-v1internal", "user-ant-v1internal"); snapshot.user_allowed_providers = Some(vec!["antigravity".to_string()]); snapshot.api_key_allowed_providers = Some(vec!["antigravity".to_string()]); snapshot.user_allowed_api_formats = Some(vec!["gemini:generate_content".to_string()]); snapshot.api_key_allowed_api_formats = Some(vec!["gemini:generate_content".to_string()]); let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), snapshot, )])); let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider( "provider-antigravity-1", "Antigravity", "antigravity", )], vec![sample_endpoint( "endpoint-antigravity-1", "provider-antigravity-1", "gemini:generate_content", )], Vec::new(), )); let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository) .with_provider_catalog_reader(provider_catalog); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); headers.insert( http::header::AUTHORIZATION, "Bearer google-oauth-access-token".parse().unwrap(), ); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1internal:streamGenerateContent?alt=sse"), Some("antigravity:v1internal"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(auth_context.local_rejection, None); } #[tokio::test] async fn data_backed_auth_context_allows_provider_id_for_convertible_endpoint_format() { let api_key = "sk-test-provider-convertible-endpoint"; let mut snapshot = sample_snapshot("key-9", "user-9"); snapshot.api_key_is_standalone = true; snapshot.user_allowed_providers = None; snapshot.api_key_allowed_providers = Some(vec!["provider-custom-openai".to_string()]); snapshot.user_allowed_api_formats = None; snapshot.api_key_allowed_api_formats = None; let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), snapshot, )])); let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider( "provider-custom-openai", "Custom OpenAI Responses Gateway", "custom", )], vec![sample_endpoint( "endpoint-custom-openai-responses", "provider-custom-openai", "openai:responses", )], Vec::new(), )); let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository) .with_provider_catalog_reader(provider_catalog); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/messages?beta=true"), Some("claude:messages"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(auth_context.local_rejection, None); } #[tokio::test] async fn data_backed_auth_context_denies_retired_anthropic_provider_alias_for_claude_route() { let api_key = "sk-test-provider-retired-anthropic-alias"; let mut snapshot = sample_snapshot("key-5", "user-5"); snapshot.user_allowed_providers = Some(vec!["anthropic".to_string()]); snapshot.api_key_allowed_providers = Some(vec!["anthropic".to_string()]); snapshot.user_allowed_api_formats = None; snapshot.api_key_allowed_api_formats = None; let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), snapshot, )])); let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider("provider-claude", "Claude", "custom")], Vec::new(), Vec::new(), )); let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository) .with_provider_catalog_reader(provider_catalog); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/messages"), Some("claude:messages"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!( auth_context.local_rejection, Some(GatewayLocalAuthRejection::ProviderNotAllowed { provider: "claude".to_string(), }) ); } #[tokio::test] async fn data_backed_auth_context_treats_empty_allowed_lists_as_unrestricted() { let api_key = "sk-test-empty-restrictions"; let mut snapshot = sample_snapshot("key-6", "user-6"); snapshot.api_key_is_standalone = true; snapshot.user_allowed_providers = Some(vec!["openai".to_string()]); snapshot.user_allowed_api_formats = Some(vec!["openai:chat".to_string()]); snapshot.user_allowed_models = Some(vec!["gpt-4.1".to_string()]); snapshot.api_key_allowed_providers = Some(Vec::new()); snapshot.api_key_allowed_api_formats = Some(Vec::new()); snapshot.api_key_allowed_models = Some(Vec::new()); let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), snapshot, )])); let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider("provider-claude", "Claude", "custom")], vec![sample_endpoint( "endpoint-claude", "provider-claude", "claude:messages", )], Vec::new(), )); let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository) .with_provider_catalog_reader(provider_catalog); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert("x-api-key", api_key.parse().unwrap()); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/messages"), Some("claude:messages"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!(auth_context.local_rejection, None); } #[tokio::test] async fn data_backed_auth_context_denies_provider_type_without_matching_allowed_provider() { let api_key = "sk-test-provider-miss"; let mut snapshot = sample_snapshot("key-3", "user-3"); snapshot.user_allowed_providers = Some(vec!["provider-claude-1".to_string()]); snapshot.api_key_allowed_providers = Some(vec!["provider-claude-1".to_string()]); let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(api_key)), snapshot, )])); let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![ sample_provider("provider-openai-1", "OpenAI Pool 1", "openai"), sample_provider("provider-claude-1", "Claude Pool 1", "claude"), ], Vec::new(), Vec::new(), )); let data = GatewayDataState::with_auth_api_key_reader_for_tests(repository) .with_provider_catalog_reader(provider_catalog); let state = AppState::new() .expect("state should build") .with_data_state_for_tests(data); let mut headers = HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, format!("Bearer {api_key}").parse().unwrap(), ); let auth_context = resolve_data_backed_auth_context( &state, &headers, &uri("/v1/chat/completions"), Some("openai:chat"), ) .await .expect("resolution should succeed") .expect("auth context should exist"); assert_eq!( auth_context.local_rejection, Some(GatewayLocalAuthRejection::ProviderNotAllowed { provider: "openai".to_string(), }) ); } }