use std::time::{SystemTime, UNIX_EPOCH}; use axum::body::Bytes; use axum::http::Uri; use sha2::{Digest, Sha256}; use url::form_urlencoded; use crate::{ ai_serving::extract_gemini_model_from_path, headers::{decoded_request_body_bytes, header_value_str, is_json_request}, }; use super::super::GatewayControlDecision; use super::types::{ GatewayCredentialBundle, GatewayCredentialCarrier, GatewayExtractedCredentials, GatewayPrimaryCredential, GatewayTrustedAdminHeaders, GatewayTrustedAuthHeaders, }; pub(crate) fn extract_requested_model( decision: &GatewayControlDecision, uri: &Uri, headers: &http::HeaderMap, body: &Bytes, ) -> Option { if decision.route_family.as_deref() == Some("gemini") { if let Some(model) = extract_gemini_model_from_path(uri.path()) { return Some(model); } } if !is_json_request(headers) || body.is_empty() { return None; } let body = decoded_request_body_bytes(headers, body.as_ref()).ok()?; serde_json::from_slice::(body.as_ref()) .ok() .and_then(|payload| { payload .get("model") .and_then(|value| value.as_str()) .map(|value| value.trim().to_string()) }) .filter(|value| !value.is_empty()) } pub(super) fn extract_request_credentials( headers: &http::HeaderMap, uri: &Uri, auth_endpoint_signature: &str, ) -> GatewayExtractedCredentials { let bundle = GatewayCredentialBundle { authorization_bearer: header_value_str(headers, http::header::AUTHORIZATION.as_str()) .as_deref() .and_then(extract_bearer_token) .map(ToOwned::to_owned), x_api_key: header_value_str(headers, "x-api-key"), api_key: header_value_str(headers, "api-key"), x_goog_api_key: header_value_str(headers, "x-goog-api-key"), query_key: extract_query_api_key(uri), cookie_header: header_value_str(headers, http::header::COOKIE.as_str()), }; let trusted_headers = extract_trusted_auth_headers(headers); let trusted_admin_headers = extract_trusted_admin_headers(headers); let primary = select_primary_credential(auth_endpoint_signature, &bundle); GatewayExtractedCredentials { trusted_headers, trusted_admin_headers, bundle, primary, } } pub(in crate::control) fn resolve_gateway_credential_carrier( headers: &http::HeaderMap, uri: &Uri, auth_endpoint_signature: &str, ) -> Option { extract_request_credentials(headers, uri, auth_endpoint_signature) .primary .map(|credential| match credential { GatewayPrimaryCredential::ProviderApiKey { carrier, .. } | GatewayPrimaryCredential::BearerToken { carrier, .. } | GatewayPrimaryCredential::CookieHeader { carrier, .. } => carrier, }) } fn has_trusted_gateway_marker(headers: &http::HeaderMap) -> bool { header_value_str(headers, crate::constants::GATEWAY_HEADER) .unwrap_or_default() .trim() .to_ascii_lowercase() .starts_with("rust-phase3") } pub(super) fn build_auth_context_cache_key( headers: &http::HeaderMap, uri: &Uri, auth_endpoint_signature: &str, ) -> Option { let signature = auth_endpoint_signature.trim(); if signature.is_empty() { return None; } let extracted = extract_request_credentials(headers, uri, signature); let trusted_headers = extracted.trusted_headers; let bundle = extracted.bundle; if bundle.authorization_bearer.is_none() && bundle.x_api_key.is_none() && bundle.api_key.is_none() && bundle.x_goog_api_key.is_none() && bundle.query_key.is_none() && bundle.cookie_header.is_none() && trusted_headers.is_none() { return None; } let (trusted_user_id, trusted_api_key_id, trusted_balance_remaining, trusted_access_allowed) = trusted_headers .map(|trusted| { ( trusted.user_id, trusted.api_key_id, trusted .balance_remaining .map(|value| value.to_string()) .unwrap_or_default(), trusted .access_allowed .map(|value| value.to_string()) .unwrap_or_default(), ) }) .unwrap_or_default(); Some(format!( "{signature}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}", bundle.authorization_bearer.unwrap_or_default(), bundle.x_api_key.unwrap_or_default(), bundle.api_key.unwrap_or_default(), bundle.x_goog_api_key.unwrap_or_default(), bundle.query_key.unwrap_or_default(), bundle.cookie_header.unwrap_or_default(), trusted_user_id, trusted_api_key_id, trusted_balance_remaining, trusted_access_allowed, )) } fn extract_trusted_auth_headers(headers: &http::HeaderMap) -> Option { if !has_trusted_gateway_marker(headers) { return None; } let user_id = header_value_str(headers, crate::constants::TRUSTED_AUTH_USER_ID_HEADER) .filter(|value| !value.is_empty())?; let api_key_id = header_value_str(headers, crate::constants::TRUSTED_AUTH_API_KEY_ID_HEADER) .filter(|value| !value.is_empty())?; let balance_remaining = header_value_str(headers, crate::constants::TRUSTED_AUTH_BALANCE_HEADER) .as_deref() .and_then(parse_f64_header); let access_allowed = header_value_str( headers, crate::constants::TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, ) .as_deref() .and_then(parse_bool_header); Some(GatewayTrustedAuthHeaders { user_id, api_key_id, balance_remaining, access_allowed, }) } #[cfg(not(test))] pub(super) fn extract_trusted_admin_headers( _headers: &http::HeaderMap, ) -> Option { // The public gateway has no authenticated upstream that is allowed to // assert an administrator principal. `x-aether-gateway` is also emitted // on public responses, so it cannot serve as proof that these headers were // produced by a trusted hop. Production requests must authenticate with a // real admin session or management bearer token instead. None } #[cfg(test)] pub(super) fn extract_trusted_admin_headers( headers: &http::HeaderMap, ) -> Option { if !has_trusted_gateway_marker(headers) { return None; } let user_id = header_value_str(headers, crate::constants::TRUSTED_ADMIN_USER_ID_HEADER)? .trim() .to_string(); if user_id.is_empty() { return None; } let user_role = header_value_str(headers, crate::constants::TRUSTED_ADMIN_USER_ROLE_HEADER)? .trim() .to_string(); if !crate::roles::can_access_admin_console(&user_role) { return None; } let session_id = header_value_str(headers, crate::constants::TRUSTED_ADMIN_SESSION_ID_HEADER) .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()); let management_token_id = header_value_str( headers, crate::constants::TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, ) .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()); if session_id.is_none() && management_token_id.is_none() { return None; } Some(GatewayTrustedAdminHeaders { user_id, user_role, session_id, management_token_id, }) } fn select_primary_credential( auth_endpoint_signature: &str, bundle: &GatewayCredentialBundle, ) -> Option { let signature = auth_endpoint_signature.trim().to_ascii_lowercase(); if signature.starts_with("gemini:") { return select_gemini_credential(bundle); } if signature.starts_with("antigravity:") { return select_antigravity_credential(bundle); } if signature.starts_with("claude:") { return select_claude_messages_credential(bundle); } if signature.starts_with("openai:") || signature.starts_with("codex:") { return select_openai_credential(bundle); } if signature.starts_with("aether:") { return select_openai_credential(bundle); } select_generic_credential(bundle) } fn select_antigravity_credential( bundle: &GatewayCredentialBundle, ) -> Option { first_provider_api_key( bundle, &[ GatewayCredentialCarrier::XApiKey, GatewayCredentialCarrier::ApiKey, ], ) .or_else(|| first_bearer_token(bundle)) .or_else(|| select_cookie_credential(bundle)) } fn select_openai_credential(bundle: &GatewayCredentialBundle) -> Option { first_provider_api_key( bundle, &[ GatewayCredentialCarrier::AuthorizationBearer, GatewayCredentialCarrier::XApiKey, GatewayCredentialCarrier::ApiKey, GatewayCredentialCarrier::XGoogApiKey, GatewayCredentialCarrier::QueryKey, ], ) .or_else(|| select_cookie_credential(bundle)) } fn select_claude_messages_credential( bundle: &GatewayCredentialBundle, ) -> Option { first_provider_api_key( bundle, &[ GatewayCredentialCarrier::XApiKey, GatewayCredentialCarrier::ApiKey, GatewayCredentialCarrier::AuthorizationBearer, ], ) .or_else(|| first_bearer_token(bundle)) .or_else(|| select_cookie_credential(bundle)) } fn select_gemini_credential(bundle: &GatewayCredentialBundle) -> Option { first_provider_api_key( bundle, &[ GatewayCredentialCarrier::QueryKey, GatewayCredentialCarrier::XGoogApiKey, GatewayCredentialCarrier::XApiKey, GatewayCredentialCarrier::ApiKey, ], ) .or_else(|| first_bearer_token(bundle)) .or_else(|| select_cookie_credential(bundle)) } fn select_generic_credential(bundle: &GatewayCredentialBundle) -> Option { first_bearer_token(bundle) .or_else(|| { first_provider_api_key( bundle, &[ GatewayCredentialCarrier::XApiKey, GatewayCredentialCarrier::ApiKey, GatewayCredentialCarrier::XGoogApiKey, GatewayCredentialCarrier::QueryKey, ], ) }) .or_else(|| select_cookie_credential(bundle)) } fn first_provider_api_key( bundle: &GatewayCredentialBundle, carriers: &[GatewayCredentialCarrier], ) -> Option { for carrier in carriers { if let Some(raw) = credential_value(bundle, *carrier) { return Some(GatewayPrimaryCredential::ProviderApiKey { raw, carrier: *carrier, }); } } None } fn first_bearer_token(bundle: &GatewayCredentialBundle) -> Option { credential_value(bundle, GatewayCredentialCarrier::AuthorizationBearer).map(|raw| { GatewayPrimaryCredential::BearerToken { raw, carrier: GatewayCredentialCarrier::AuthorizationBearer, } }) } fn select_cookie_credential(bundle: &GatewayCredentialBundle) -> Option { credential_value(bundle, GatewayCredentialCarrier::CookieHeader).map(|raw| { GatewayPrimaryCredential::CookieHeader { raw, carrier: GatewayCredentialCarrier::CookieHeader, } }) } fn credential_value( bundle: &GatewayCredentialBundle, carrier: GatewayCredentialCarrier, ) -> Option { match carrier { GatewayCredentialCarrier::AuthorizationBearer => bundle.authorization_bearer.clone(), GatewayCredentialCarrier::XApiKey => bundle.x_api_key.clone(), GatewayCredentialCarrier::ApiKey => bundle.api_key.clone(), GatewayCredentialCarrier::XGoogApiKey => bundle.x_goog_api_key.clone(), GatewayCredentialCarrier::QueryKey => bundle.query_key.clone(), GatewayCredentialCarrier::CookieHeader => bundle.cookie_header.clone(), } } fn extract_query_api_key(uri: &Uri) -> Option { let query = uri.query()?; form_urlencoded::parse(query.as_bytes()) .find(|(key, value)| key == "key" && !value.trim().is_empty()) .map(|(_, value)| value.into_owned()) } fn extract_bearer_token(value: &str) -> Option<&str> { let trimmed = value.trim(); let (scheme, token) = trimmed.split_once(' ')?; if !scheme.eq_ignore_ascii_case("bearer") { return None; } let token = token.trim(); if token.is_empty() { None } else { Some(token) } } pub(super) fn hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); hasher.update(value.as_bytes()); format!("{:x}", hasher.finalize()) } pub(super) fn contains_string(items: &[String], target: &str) -> bool { items .iter() .any(|item| item.trim().eq_ignore_ascii_case(target.trim())) } pub(super) fn parse_bool_header(value: &str) -> Option { match value.trim().to_ascii_lowercase().as_str() { "true" | "1" | "yes" => Some(true), "false" | "0" | "no" => Some(false), _ => None, } } pub(super) fn parse_f64_header(value: &str) -> Option { value.trim().parse::().ok() } pub(super) fn current_unix_secs() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs() } #[cfg(test)] mod tests { use super::{ build_auth_context_cache_key, extract_request_credentials, extract_requested_model, GatewayCredentialCarrier, GatewayPrimaryCredential, GatewayTrustedAdminHeaders, GatewayTrustedAuthHeaders, }; use crate::control::GatewayControlDecision; use axum::body::Bytes; use axum::http::{self, Uri}; fn uri(path: &str) -> Uri { path.parse().expect("uri should parse") } #[test] fn extract_requested_model_reads_zstd_encoded_json_body() { let decision = GatewayControlDecision::synthetic( "/v1/responses", Some("ai_public".to_string()), Some("openai".to_string()), Some("responses".to_string()), Some("openai:responses".to_string()), ); let mut headers = http::HeaderMap::new(); headers.insert( http::header::CONTENT_TYPE, http::HeaderValue::from_static("application/json"), ); headers.insert( http::header::CONTENT_ENCODING, http::HeaderValue::from_static("zstd"), ); let encoded = zstd::stream::encode_all(br#"{"model":"gpt-5.4","input":"hello"}"#.as_slice(), 0) .expect("zstd body should encode"); let requested_model = extract_requested_model( &decision, &uri("/v1/responses"), &headers, &Bytes::from(encoded), ); assert_eq!(requested_model.as_deref(), Some("gpt-5.4")); } #[test] fn selects_openai_bearer_as_provider_api_key() { let mut headers = http::HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, "Bearer sk-openai".parse().unwrap(), ); let extracted = extract_request_credentials(&headers, &uri("/v1/chat/completions"), "openai:chat"); assert_eq!( extracted.primary, Some(GatewayPrimaryCredential::ProviderApiKey { raw: "sk-openai".to_string(), carrier: GatewayCredentialCarrier::AuthorizationBearer, }) ); } #[test] fn selects_codex_live_bearer_as_provider_api_key() { let mut headers = http::HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, "Bearer sk-codex-live".parse().unwrap(), ); let extracted = extract_request_credentials(&headers, &uri("/v1/live?model=gpt-live"), "codex:live"); assert_eq!( extracted.primary, Some(GatewayPrimaryCredential::ProviderApiKey { raw: "sk-codex-live".to_string(), carrier: GatewayCredentialCarrier::AuthorizationBearer, }) ); } #[test] fn prefers_claude_chat_x_api_key_over_bearer() { let mut headers = http::HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, "Bearer cli-token".parse().unwrap(), ); headers.insert("x-api-key", "claude-key".parse().unwrap()); let extracted = extract_request_credentials(&headers, &uri("/v1/messages"), "claude:messages"); assert_eq!( extracted.primary, Some(GatewayPrimaryCredential::ProviderApiKey { raw: "claude-key".to_string(), carrier: GatewayCredentialCarrier::XApiKey, }) ); } #[test] fn selects_claude_cli_bearer_as_provider_api_key() { let mut headers = http::HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, "Bearer cli-token".parse().unwrap(), ); let extracted = extract_request_credentials(&headers, &uri("/v1/messages"), "claude:messages"); assert_eq!( extracted.primary, Some(GatewayPrimaryCredential::ProviderApiKey { raw: "cli-token".to_string(), carrier: GatewayCredentialCarrier::AuthorizationBearer, }) ); } #[test] fn prefers_antigravity_aether_api_key_over_google_bearer() { let mut headers = http::HeaderMap::new(); headers.insert( http::header::AUTHORIZATION, "Bearer google-oauth-access-token".parse().unwrap(), ); headers.insert("x-api-key", "sk-aether-antigravity".parse().unwrap()); let extracted = extract_request_credentials( &headers, &uri("/v1internal:streamGenerateContent?alt=sse"), "antigravity:v1internal", ); assert_eq!( extracted.primary, Some(GatewayPrimaryCredential::ProviderApiKey { raw: "sk-aether-antigravity".to_string(), carrier: GatewayCredentialCarrier::XApiKey, }) ); } #[test] fn prefers_gemini_query_key_over_header_key() { let mut headers = http::HeaderMap::new(); headers.insert("x-goog-api-key", "gemini-header".parse().unwrap()); let extracted = extract_request_credentials( &headers, &uri("/v1beta/models?key=gemini-query"), "gemini:generate_content", ); assert_eq!( extracted.primary, Some(GatewayPrimaryCredential::ProviderApiKey { raw: "gemini-query".to_string(), carrier: GatewayCredentialCarrier::QueryKey, }) ); } #[test] fn extracts_cookie_as_fallback_credential() { let mut headers = http::HeaderMap::new(); headers.insert(http::header::COOKIE, "session=abc123".parse().unwrap()); let extracted = extract_request_credentials(&headers, &uri("/v1/chat/completions"), "internal:session"); assert_eq!( extracted.primary, Some(GatewayPrimaryCredential::CookieHeader { raw: "session=abc123".to_string(), carrier: GatewayCredentialCarrier::CookieHeader, }) ); } #[test] fn cache_key_includes_cookie_header() { let mut headers = http::HeaderMap::new(); headers.insert(http::header::COOKIE, "session=abc123".parse().unwrap()); let cache_key = build_auth_context_cache_key( &headers, &uri("/v1/chat/completions"), "internal:session", ) .expect("cache key should exist"); assert!(cache_key.contains("session=abc123")); } #[test] fn cache_key_includes_trusted_auth_headers() { let mut first_headers = http::HeaderMap::new(); first_headers.insert( crate::constants::GATEWAY_HEADER, "rust-phase3b".parse().unwrap(), ); first_headers.insert( crate::constants::TRUSTED_AUTH_USER_ID_HEADER, "user-1".parse().unwrap(), ); first_headers.insert( crate::constants::TRUSTED_AUTH_API_KEY_ID_HEADER, "key-1".parse().unwrap(), ); first_headers.insert( crate::constants::TRUSTED_AUTH_BALANCE_HEADER, "1.5".parse().unwrap(), ); first_headers.insert( crate::constants::TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, "true".parse().unwrap(), ); let mut second_headers = first_headers.clone(); second_headers.insert( crate::constants::TRUSTED_AUTH_USER_ID_HEADER, "user-2".parse().unwrap(), ); second_headers.insert( crate::constants::TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, "false".parse().unwrap(), ); let first = build_auth_context_cache_key( &first_headers, &uri("/v1/chat/completions"), "openai:chat", ) .expect("trusted cache key should exist"); let second = build_auth_context_cache_key( &second_headers, &uri("/v1/chat/completions"), "openai:chat", ) .expect("trusted cache key should exist"); assert_ne!(first, second); assert!(first.contains("user-1")); assert!(first.contains("key-1")); assert!(first.contains("1.5")); assert!(first.contains("true")); assert!(second.contains("user-2")); assert!(second.contains("false")); } #[test] fn cache_key_ignores_untrusted_auth_identity_headers() { let mut trusted_headers = http::HeaderMap::new(); trusted_headers.insert( crate::constants::GATEWAY_HEADER, "rust-phase3b".parse().unwrap(), ); trusted_headers.insert( crate::constants::TRUSTED_AUTH_USER_ID_HEADER, "user-1".parse().unwrap(), ); trusted_headers.insert( crate::constants::TRUSTED_AUTH_API_KEY_ID_HEADER, "key-1".parse().unwrap(), ); let mut untrusted_headers = http::HeaderMap::new(); untrusted_headers.insert( crate::constants::TRUSTED_AUTH_USER_ID_HEADER, "user-1".parse().unwrap(), ); untrusted_headers.insert( crate::constants::TRUSTED_AUTH_API_KEY_ID_HEADER, "key-1".parse().unwrap(), ); let trusted = build_auth_context_cache_key( &trusted_headers, &uri("/v1/chat/completions"), "openai:chat", ); let untrusted = build_auth_context_cache_key( &untrusted_headers, &uri("/v1/chat/completions"), "openai:chat", ); assert!(trusted.is_some()); assert_eq!(untrusted, None); } #[test] fn extracts_trusted_auth_headers() { let mut headers = http::HeaderMap::new(); headers.insert( crate::constants::GATEWAY_HEADER, "rust-phase3b".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_AUTH_USER_ID_HEADER, "user-1".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_AUTH_API_KEY_ID_HEADER, "key-1".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_AUTH_BALANCE_HEADER, "1.5".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, "true".parse().unwrap(), ); let extracted = extract_request_credentials(&headers, &uri("/v1/chat/completions"), "openai:chat"); assert_eq!( extracted.trusted_headers, Some(GatewayTrustedAuthHeaders { user_id: "user-1".to_string(), api_key_id: "key-1".to_string(), balance_remaining: Some(1.5), access_allowed: Some(true), }) ); assert_eq!(extracted.trusted_admin_headers, None); } #[test] fn ignores_trusted_auth_headers_without_gateway_marker() { let mut headers = http::HeaderMap::new(); headers.insert( crate::constants::TRUSTED_AUTH_USER_ID_HEADER, "user-1".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_AUTH_API_KEY_ID_HEADER, "key-1".parse().unwrap(), ); let extracted = extract_request_credentials(&headers, &uri("/v1/chat/completions"), "openai:chat"); assert_eq!(extracted.trusted_headers, None); } #[test] fn extracts_trusted_admin_headers() { let mut headers = http::HeaderMap::new(); headers.insert( crate::constants::GATEWAY_HEADER, "rust-phase3b".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-1".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_ADMIN_USER_ROLE_HEADER, "admin".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_ADMIN_SESSION_ID_HEADER, "sess-1".parse().unwrap(), ); let extracted = extract_request_credentials( &headers, &uri("/api/admin/endpoints/health/api-formats"), "admin:endpoints_health", ); assert_eq!( extracted.trusted_admin_headers, Some(GatewayTrustedAdminHeaders { user_id: "admin-user-1".to_string(), user_role: "admin".to_string(), session_id: Some("sess-1".to_string()), management_token_id: None, }) ); } #[test] fn extracts_trusted_audit_admin_headers() { let mut headers = http::HeaderMap::new(); headers.insert( crate::constants::GATEWAY_HEADER, "rust-phase3b".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_ADMIN_USER_ID_HEADER, "audit-admin-1".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_ADMIN_USER_ROLE_HEADER, "audit_admin".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_ADMIN_SESSION_ID_HEADER, "sess-audit-1".parse().unwrap(), ); let extracted = extract_request_credentials( &headers, &uri("/api/admin/endpoints/health/api-formats"), "admin:endpoints_health", ); assert_eq!( extracted.trusted_admin_headers, Some(GatewayTrustedAdminHeaders { user_id: "audit-admin-1".to_string(), user_role: "audit_admin".to_string(), session_id: Some("sess-audit-1".to_string()), management_token_id: None, }) ); } #[test] fn ignores_trusted_admin_headers_without_gateway_marker() { let mut headers = http::HeaderMap::new(); headers.insert( crate::constants::TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-1".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_ADMIN_USER_ROLE_HEADER, "admin".parse().unwrap(), ); headers.insert( crate::constants::TRUSTED_ADMIN_SESSION_ID_HEADER, "sess-1".parse().unwrap(), ); let extracted = extract_request_credentials( &headers, &uri("/api/admin/endpoints/health/api-formats"), "admin:endpoints_health", ); assert_eq!(extracted.trusted_admin_headers, None); } }