From bf139af37ae8e176c7c0923a7a8015495e7aec75 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sun, 25 Jan 2026 17:52:42 -0500 Subject: [PATCH] [smg] use already published auth crate for better compilation speed (#17723) --- sgl-model-gateway/Cargo.toml | 10 +- sgl-model-gateway/src/auth/audit.rs | 357 --------------- sgl-model-gateway/src/auth/config.rs | 345 --------------- sgl-model-gateway/src/auth/jwks.rs | 517 ---------------------- sgl-model-gateway/src/auth/jwt.rs | 536 ----------------------- sgl-model-gateway/src/auth/middleware.rs | 399 ----------------- sgl-model-gateway/src/auth/mod.rs | 20 - sgl-model-gateway/src/lib.rs | 2 +- 8 files changed, 2 insertions(+), 2184 deletions(-) delete mode 100644 sgl-model-gateway/src/auth/audit.rs delete mode 100644 sgl-model-gateway/src/auth/config.rs delete mode 100644 sgl-model-gateway/src/auth/jwks.rs delete mode 100644 sgl-model-gateway/src/auth/jwt.rs delete mode 100644 sgl-model-gateway/src/auth/middleware.rs delete mode 100644 sgl-model-gateway/src/auth/mod.rs diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index bae768800..1bbabda10 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -71,7 +71,6 @@ metrics-exporter-prometheus = "0.17.0" uuid = { version = "1.10", features = ["v4", "serde"] } ulid = "1.2.1" parking_lot = "0.12.4" -rayon = "1.10" thiserror = "2.0.12" @@ -81,19 +80,14 @@ url = "2.5.4" validator = { version = "0.20.0", features = ["derive"] } tokio-stream = { version = "0.1", features = ["sync"] } anyhow = "1.0" -tokenizers = { version = "0.22.0" } -tiktoken-rs = { version = "0.7.0" } reasoning-parser = "1.0.0" openai-protocol = { version = "1.0.0", features = ["axum"] } tool-parser = "1.0.0" llm-tokenizer = "1.0.0" -minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins"] } -minijinja-contrib = { version = "2.0", features = ["pycompat"] } +smg-auth = "1.0.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } -tokio-rustls = { version = "0.26", default-features = false, features = ["ring", "logging", "tls12"] } rustls-pemfile = "2.2" openssl = "0.10.73" -hf-hub = { version = "0.4.3", features = ["tokio"] } rmcp = { version = "0.8.3", features = ["client", "server", "transport-child-process", "transport-sse-client-reqwest", @@ -106,7 +100,6 @@ serde_yaml = "0.9" oracle = { version = "0.6.3", features = ["chrono"] } subtle = "2.6" jsonwebtoken = { version = "9.3", default-features = false, features = ["use_pem"] } -rustpython-parser = "0.4.0" num-traits = "0.2" num-bigint = "0.4" image = { version = "0.25.4", default-features = false, features = ["png", "jpeg", "gif", "bmp", "ico", "tiff", "webp"] } @@ -143,7 +136,6 @@ aho-corasick = "1.1.4" [build-dependencies] tonic-prost-build = "0.14.2" -prost-build = "0.14.1" chrono = { version = "0.4", features = ["clock"] } toml = "0.9" diff --git a/sgl-model-gateway/src/auth/audit.rs b/sgl-model-gateway/src/auth/audit.rs deleted file mode 100644 index 0decc250b..000000000 --- a/sgl-model-gateway/src/auth/audit.rs +++ /dev/null @@ -1,357 +0,0 @@ -//! Audit logging for control plane operations. -//! -//! Provides structured audit events for security monitoring and compliance. -//! -//! Security features: -//! - Input sanitization to prevent log injection attacks -//! - Structured logging for safe parsing - -use chrono::{DateTime, Utc}; -use serde::Serialize; -use tracing::{info, span, Level}; - -use super::config::Role; - -/// Maximum length for sanitized strings to prevent log flooding -const MAX_SANITIZED_LENGTH: usize = 1024; - -/// Sanitize a string for safe logging. -/// Removes/escapes control characters and newlines that could be used for log injection. -fn sanitize_for_log(input: &str) -> String { - let mut result = String::with_capacity(input.len().min(MAX_SANITIZED_LENGTH)); - - for ch in input.chars().take(MAX_SANITIZED_LENGTH) { - match ch { - // Replace newlines and carriage returns with escaped versions - '\n' => result.push_str("\\n"), - '\r' => result.push_str("\\r"), - // Replace tabs with escaped version - '\t' => result.push_str("\\t"), - // Escape backslashes - '\\' => result.push_str("\\\\"), - // Remove other control characters (ASCII 0-31, 127) - c if c.is_control() => { - result.push_str(&format!("\\x{:02x}", c as u32)); - } - // Keep printable characters as-is - c => result.push(c), - } - } - - // Indicate truncation if input was too long - if input.len() > MAX_SANITIZED_LENGTH { - result.push_str("...[TRUNCATED]"); - } - - result -} - -/// Outcome of an audited operation. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] -#[serde(rename_all = "lowercase")] -pub enum AuditOutcome { - /// Operation completed successfully - Success, - /// Operation denied (authorization failure) - Denied, -} - -impl std::fmt::Display for AuditOutcome { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - AuditOutcome::Success => write!(f, "success"), - AuditOutcome::Denied => write!(f, "denied"), - } - } -} - -/// Audit event for control plane operations. -#[derive(Debug, Clone, Serialize)] -pub struct AuditEvent { - /// Timestamp of the event - pub timestamp: DateTime, - - /// Principal who performed the action (subject ID or API key ID) - pub principal: String, - - /// Authentication method used (jwt, api_key) - pub auth_method: String, - - /// Role of the principal - pub role: Role, - - /// HTTP method - pub method: String, - - /// Request path - pub path: String, - - /// Resource being accessed (e.g., worker ID, wasm module ID) - #[serde(skip_serializing_if = "Option::is_none")] - pub resource: Option, - - /// Operation outcome - pub outcome: AuditOutcome, - - /// Request ID for correlation - #[serde(skip_serializing_if = "Option::is_none")] - pub request_id: Option, - - /// Additional details or error message - #[serde(skip_serializing_if = "Option::is_none")] - pub details: Option, -} - -/// Context for audit logging containing request and principal information. -#[derive(Debug, Clone)] -pub struct AuditContext<'a> { - /// Principal who performed the action - pub principal: &'a str, - /// Authentication method used (jwt, api_key) - pub auth_method: &'a str, - /// Role of the principal - pub role: Role, - /// HTTP method - pub method: &'a str, - /// Request path - pub path: &'a str, - /// Request ID for correlation - pub request_id: Option<&'a str>, -} - -impl<'a> AuditContext<'a> { - /// Create a new audit context. - pub fn new( - principal: &'a str, - auth_method: &'a str, - role: Role, - method: &'a str, - path: &'a str, - request_id: Option<&'a str>, - ) -> Self { - Self { - principal, - auth_method, - role, - method, - path, - request_id, - } - } -} - -impl AuditEvent { - /// Create a new audit event from context. - fn from_context( - ctx: &AuditContext<'_>, - outcome: AuditOutcome, - resource: Option<&str>, - details: Option<&str>, - ) -> Self { - Self { - timestamp: Utc::now(), - principal: sanitize_for_log(ctx.principal), - auth_method: sanitize_for_log(ctx.auth_method), - role: ctx.role, - method: sanitize_for_log(ctx.method), - path: sanitize_for_log(ctx.path), - resource: resource.map(sanitize_for_log), - outcome, - request_id: ctx.request_id.map(sanitize_for_log), - details: details.map(sanitize_for_log), - } - } - - /// Create an audit event for unauthenticated requests. - fn unauthenticated(method: &str, path: &str, reason: &str, request_id: Option<&str>) -> Self { - Self { - timestamp: Utc::now(), - principal: "unauthenticated".to_string(), - auth_method: "none".to_string(), - role: Role::User, - method: sanitize_for_log(method), - path: sanitize_for_log(path), - resource: None, - outcome: AuditOutcome::Denied, - request_id: request_id.map(sanitize_for_log), - details: Some(sanitize_for_log(reason)), - } - } -} - -/// Audit logger that emits structured audit events. -#[derive(Clone, Default)] -pub struct AuditLogger { - enabled: bool, -} - -impl AuditLogger { - /// Create a new audit logger. - pub fn new(enabled: bool) -> Self { - Self { enabled } - } - - /// Check if audit logging is enabled. - pub fn is_enabled(&self) -> bool { - self.enabled - } - - /// Log an audit event. - pub fn log(&self, event: &AuditEvent) { - if !self.enabled { - return; - } - - // Create a span for structured logging - let _span = span!( - Level::INFO, - "audit", - principal = %event.principal, - auth_method = %event.auth_method, - role = %event.role, - method = %event.method, - path = %event.path, - outcome = %event.outcome, - ) - .entered(); - - // Log the event - info!( - target: "smg::audit", - timestamp = %event.timestamp.to_rfc3339(), - principal = %event.principal, - auth_method = %event.auth_method, - role = %event.role, - method = %event.method, - path = %event.path, - resource = ?event.resource, - outcome = %event.outcome, - request_id = ?event.request_id, - details = ?event.details, - "control_plane_audit" - ); - } - - /// Log a successful operation. - pub fn log_success(&self, ctx: &AuditContext<'_>, resource: Option<&str>) { - self.log(&AuditEvent::from_context( - ctx, - AuditOutcome::Success, - resource, - None, - )); - } - - /// Log a denied operation (authorization failure). - pub fn log_denied(&self, ctx: &AuditContext<'_>, reason: &str) { - self.log(&AuditEvent::from_context( - ctx, - AuditOutcome::Denied, - None, - Some(reason), - )); - } - - /// Log an authentication failure (before principal is known). - pub fn log_auth_failure( - &self, - method: &str, - path: &str, - reason: &str, - request_id: Option<&str>, - ) { - self.log(&AuditEvent::unauthenticated( - method, path, reason, request_id, - )); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_audit_event_from_context() { - let ctx = AuditContext::new( - "user@example.com", - "jwt", - Role::Admin, - "POST", - "/workers", - Some("req-abc"), - ); - let event = AuditEvent::from_context(&ctx, AuditOutcome::Success, Some("worker-123"), None); - - assert_eq!(event.principal, "user@example.com"); - assert_eq!(event.auth_method, "jwt"); - assert_eq!(event.role, Role::Admin); - assert_eq!(event.method, "POST"); - assert_eq!(event.path, "/workers"); - assert_eq!(event.resource, Some("worker-123".to_string())); - assert_eq!(event.outcome, AuditOutcome::Success); - assert_eq!(event.request_id, Some("req-abc".to_string())); - } - - #[test] - fn test_audit_outcome_display() { - assert_eq!(AuditOutcome::Success.to_string(), "success"); - assert_eq!(AuditOutcome::Denied.to_string(), "denied"); - } - - #[test] - fn test_sanitize_for_log_normal_input() { - assert_eq!(sanitize_for_log("normal string"), "normal string"); - assert_eq!(sanitize_for_log("user@example.com"), "user@example.com"); - assert_eq!(sanitize_for_log("/api/v1/workers"), "/api/v1/workers"); - } - - #[test] - fn test_sanitize_for_log_newlines() { - assert_eq!(sanitize_for_log("line1\nline2"), "line1\\nline2"); - assert_eq!(sanitize_for_log("line1\r\nline2"), "line1\\r\\nline2"); - } - - #[test] - fn test_sanitize_for_log_control_chars() { - assert_eq!(sanitize_for_log("test\x00null"), "test\\x00null"); - assert_eq!(sanitize_for_log("test\x1bescape"), "test\\x1bescape"); - } - - #[test] - fn test_sanitize_for_log_backslashes() { - assert_eq!(sanitize_for_log("path\\to\\file"), "path\\\\to\\\\file"); - } - - #[test] - fn test_sanitize_for_log_truncation() { - let long_string = "a".repeat(2000); - let sanitized = sanitize_for_log(&long_string); - assert!(sanitized.len() < 2000); - assert!(sanitized.ends_with("...[TRUNCATED]")); - } - - #[test] - fn test_audit_event_sanitizes_inputs() { - let ctx = AuditContext::new( - "user\ninjected", - "jwt", - Role::User, - "GET", - "/workers\r\nfake_log_entry", - None, - ); - let event = AuditEvent::from_context( - &ctx, - AuditOutcome::Denied, - None, - Some("error\x00with\x1bnull"), - ); - - // Verify sanitization was applied - assert!(!event.principal.contains('\n')); - assert!(event.principal.contains("\\n")); - assert!(!event.path.contains('\r')); - assert!(event.path.contains("\\r")); - assert!(event.details.as_ref().unwrap().contains("\\x00")); - } -} diff --git a/sgl-model-gateway/src/auth/config.rs b/sgl-model-gateway/src/auth/config.rs deleted file mode 100644 index 774896d19..000000000 --- a/sgl-model-gateway/src/auth/config.rs +++ /dev/null @@ -1,345 +0,0 @@ -//! Configuration types for control plane authentication. -//! -//! Security features: -//! - API keys are hashed at load time (never stored in plaintext in memory) -//! - Constant-time comparison for API key validation - -use std::collections::HashMap; - -use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; -use subtle::ConstantTimeEq; - -/// Role-based access control for control plane APIs. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] -#[serde(rename_all = "lowercase")] -pub enum Role { - /// Full access to all control plane APIs (workers, wasm, tokenizers, etc.) - Admin, - /// Access to inference/data plane APIs only (default for backward compatibility) - #[default] - User, -} - -impl Role { - /// Check if this role has admin privileges for control plane APIs. - pub fn is_admin(&self) -> bool { - matches!(self, Role::Admin) - } -} - -impl std::fmt::Display for Role { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Role::Admin => write!(f, "admin"), - Role::User => write!(f, "user"), - } - } -} - -impl std::str::FromStr for Role { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "admin" => Ok(Role::Admin), - "user" => Ok(Role::User), - _ => Err(format!("Invalid role: {}. Valid roles: admin, user", s)), - } - } -} - -/// JWT/OIDC configuration for external identity provider integration. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct JwtConfig { - /// OIDC issuer URL (e.g., "https://login.microsoftonline.com/{tenant}/v2.0") - /// Used to discover JWKS endpoint via .well-known/openid-configuration - pub issuer: String, - - /// Expected audience claim (usually the client ID or API identifier) - pub audience: String, - - /// Optional explicit JWKS URI. If not provided, discovered from issuer. - #[serde(skip_serializing_if = "Option::is_none")] - pub jwks_uri: Option, - - /// Claim name containing the role (default: "role" or "roles") - #[serde(default = "default_role_claim")] - pub role_claim: String, - - /// Mapping from IDP role/group names to gateway roles. - /// Example: {"Gateway.Admin": "admin", "Gateway.User": "user"} - #[serde(default)] - pub role_mapping: HashMap, - - /// Clock skew tolerance in seconds (default: 30) - #[serde(default = "default_leeway_secs")] - pub leeway_secs: u64, - - /// JWKS cache TTL in seconds (default: 3600 = 1 hour) - #[serde(default = "default_jwks_cache_ttl_secs")] - pub jwks_cache_ttl_secs: u64, -} - -fn default_role_claim() -> String { - "roles".to_string() -} - -fn default_leeway_secs() -> u64 { - 30 -} - -fn default_jwks_cache_ttl_secs() -> u64 { - 3600 -} - -impl JwtConfig { - /// Create a new JWT config with required fields. - pub fn new(issuer: impl Into, audience: impl Into) -> Self { - Self { - issuer: issuer.into(), - audience: audience.into(), - jwks_uri: None, - role_claim: default_role_claim(), - role_mapping: HashMap::new(), - leeway_secs: default_leeway_secs(), - jwks_cache_ttl_secs: default_jwks_cache_ttl_secs(), - } - } - - /// Set explicit JWKS URI instead of using OIDC discovery. - pub fn with_jwks_uri(mut self, jwks_uri: impl Into) -> Self { - self.jwks_uri = Some(jwks_uri.into()); - self - } - - /// Add a role mapping from IDP role to gateway role. - pub fn with_role_mapping(mut self, idp_role: impl Into, gateway_role: Role) -> Self { - self.role_mapping.insert(idp_role.into(), gateway_role); - self - } -} - -/// Hash an API key using SHA-256. -/// Returns the hash as a fixed-size byte array. -fn hash_api_key(key: &str) -> [u8; 32] { - let mut hasher = Sha256::new(); - hasher.update(key.as_bytes()); - hasher.finalize().into() -} - -/// API key entry for service account authentication. -/// The key is hashed at construction time and never stored in plaintext. -#[derive(Clone, Serialize, Deserialize)] -pub struct ApiKeyEntry { - /// Unique identifier for this key - pub id: String, - - /// Human-readable name/description - pub name: String, - - /// SHA-256 hash of the API key (never store plaintext) - #[serde(skip)] - key_hash: [u8; 32], - - /// Role assigned to this API key - #[serde(default)] - pub role: Role, -} - -impl std::fmt::Debug for ApiKeyEntry { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ApiKeyEntry") - .field("id", &self.id) - .field("name", &self.name) - .field("key_hash", &"[REDACTED]") - .field("role", &self.role) - .finish() - } -} - -impl ApiKeyEntry { - /// Create a new API key entry. - /// The key is immediately hashed and the plaintext is not stored. - pub fn new( - id: impl Into, - name: impl Into, - key: impl Into, - role: Role, - ) -> Self { - let key_str = key.into(); - Self { - id: id.into(), - name: name.into(), - key_hash: hash_api_key(&key_str), - role, - } - } - - /// Check if the provided key matches this entry. - /// Uses constant-time comparison to prevent timing attacks. - pub fn verify(&self, key: &str) -> bool { - let provided_hash = hash_api_key(key); - self.key_hash.ct_eq(&provided_hash).into() - } -} - -/// Complete control plane authentication configuration. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct ControlPlaneAuthConfig { - /// JWT/OIDC configuration for external IDP. - #[serde(skip_serializing_if = "Option::is_none")] - pub jwt: Option, - - /// API keys for service accounts (parsed from CLI). - /// Format: "role:key" where role is "admin" or "user" - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub api_keys: Vec, - - /// Enable audit logging for control plane operations - #[serde(default = "default_audit_enabled")] - pub audit_enabled: bool, -} - -fn default_audit_enabled() -> bool { - true -} - -impl ControlPlaneAuthConfig { - /// Check if any authentication method is configured. - pub fn is_enabled(&self) -> bool { - self.jwt.is_some() || !self.api_keys.is_empty() - } - - /// Check if JWT authentication is configured. - pub fn has_jwt(&self) -> bool { - self.jwt.is_some() - } - - /// Check if any API keys are configured. - pub fn has_api_keys(&self) -> bool { - !self.api_keys.is_empty() - } - - /// Find an API key by its value and return the entry if found. - /// Uses constant-time hash comparison to prevent timing attacks. - pub fn find_api_key(&self, key: &str) -> Option<&ApiKeyEntry> { - // Iterate through all keys to prevent timing leaks about key existence - // We use a variable to track the match to ensure constant-time behavior - let mut found: Option<&ApiKeyEntry> = None; - for entry in &self.api_keys { - if entry.verify(key) { - found = Some(entry); - // Don't break early to maintain constant-time behavior - } - } - found - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_role_parsing() { - assert_eq!("admin".parse::().unwrap(), Role::Admin); - assert_eq!("ADMIN".parse::().unwrap(), Role::Admin); - assert_eq!("user".parse::().unwrap(), Role::User); - assert_eq!("USER".parse::().unwrap(), Role::User); - assert!("invalid".parse::().is_err()); - } - - #[test] - fn test_role_display() { - assert_eq!(Role::Admin.to_string(), "admin"); - assert_eq!(Role::User.to_string(), "user"); - } - - #[test] - fn test_role_is_admin() { - assert!(Role::Admin.is_admin()); - assert!(!Role::User.is_admin()); - } - - #[test] - fn test_jwt_config_builder() { - let config = JwtConfig::new("https://issuer.example.com", "api://my-app") - .with_jwks_uri("https://issuer.example.com/.well-known/jwks.json") - .with_role_mapping("Gateway.Admin", Role::Admin) - .with_role_mapping("Gateway.User", Role::User); - - assert_eq!(config.issuer, "https://issuer.example.com"); - assert_eq!(config.audience, "api://my-app"); - assert!(config.jwks_uri.is_some()); - assert_eq!(config.role_mapping.len(), 2); - } - - #[test] - fn test_control_plane_auth_config() { - let mut config = ControlPlaneAuthConfig::default(); - assert!(!config.is_enabled()); - assert!(!config.has_jwt()); - assert!(!config.has_api_keys()); - - config.jwt = Some(JwtConfig::new("https://issuer.example.com", "api://test")); - assert!(config.is_enabled()); - assert!(config.has_jwt()); - - config.api_keys.push(ApiKeyEntry::new( - "test-key", - "Test Key", - "secret123", - Role::Admin, - )); - assert!(config.has_api_keys()); - } - - #[test] - fn test_find_api_key() { - let config = ControlPlaneAuthConfig { - jwt: None, - api_keys: vec![ - ApiKeyEntry::new("key1", "Key 1", "secret1", Role::Admin), - ApiKeyEntry::new("key2", "Key 2", "secret2", Role::User), - ], - audit_enabled: true, - }; - - let found = config.find_api_key("secret1"); - assert!(found.is_some()); - assert_eq!(found.unwrap().id, "key1"); - assert_eq!(found.unwrap().role, Role::Admin); - - let found = config.find_api_key("secret2"); - assert!(found.is_some()); - assert_eq!(found.unwrap().role, Role::User); - - assert!(config.find_api_key("invalid").is_none()); - } - - #[test] - fn test_api_key_hashing() { - let entry = ApiKeyEntry::new("test", "Test Key", "my-secret-key", Role::Admin); - - // Verify the correct key works - assert!(entry.verify("my-secret-key")); - - // Verify wrong keys don't work - assert!(!entry.verify("wrong-key")); - assert!(!entry.verify("my-secret-ke")); // prefix - assert!(!entry.verify("my-secret-keyy")); // suffix - assert!(!entry.verify("")); // empty - } - - #[test] - fn test_api_key_debug_redacts_hash() { - let entry = ApiKeyEntry::new("test", "Test Key", "secret", Role::Admin); - let debug_str = format!("{:?}", entry); - - // Should contain REDACTED, not the actual hash - assert!(debug_str.contains("[REDACTED]")); - // Should not contain the key or hash bytes - assert!(!debug_str.contains("secret")); - } -} diff --git a/sgl-model-gateway/src/auth/jwks.rs b/sgl-model-gateway/src/auth/jwks.rs deleted file mode 100644 index 26611e53f..000000000 --- a/sgl-model-gateway/src/auth/jwks.rs +++ /dev/null @@ -1,517 +0,0 @@ -//! JWKS (JSON Web Key Set) fetching and caching. -//! -//! Handles: -//! - OIDC discovery to find JWKS endpoint -//! - Fetching and parsing JWKS -//! - Caching with TTL and automatic refresh -//! -//! Security features: -//! - SSRF protection: only HTTPS URLs allowed, private IPs blocked -//! - Response size limits to prevent DoS -//! - Redirect prevention - -use std::{ - net::{IpAddr, Ipv4Addr, Ipv6Addr}, - time::{Duration, Instant}, -}; - -use jsonwebtoken::jwk::{Jwk, JwkSet}; -use parking_lot::RwLock; -use tracing::{debug, info, warn}; -use url::Url; - -/// Maximum allowed JWKS response size (1 MB) -const MAX_JWKS_RESPONSE_SIZE: u64 = 1024 * 1024; - -/// Error types for JWKS operations. -#[derive(Debug, thiserror::Error)] -pub enum JwksError { - #[error("Failed to fetch OIDC discovery document: {0}")] - DiscoveryFetch(String), - - #[error("Failed to parse OIDC discovery document: {0}")] - DiscoveryParse(String), - - #[error("JWKS URI not found in discovery document")] - JwksUriNotFound, - - #[error("Failed to fetch JWKS: {0}")] - JwksFetch(String), - - #[error("Failed to parse JWKS: {0}")] - JwksParse(String), - - #[error("Key not found for kid: {0}")] - KeyNotFound(String), - - #[error("Invalid URL: {0}")] - InvalidUrl(String), - - #[error("SSRF protection: {0}")] - SsrfBlocked(String), - - #[error("Response too large: {0} bytes (max: {1})")] - ResponseTooLarge(u64, u64), - - #[error("Failed to create HTTP client: {0}")] - HttpClientError(String), -} - -/// Check if an IP address is private/internal (SSRF protection). -fn is_private_ip(ip: &IpAddr) -> bool { - match ip { - IpAddr::V4(ipv4) => { - ipv4.is_private() // 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16 - || ipv4.is_loopback() // 127.0.0.0/8 - || ipv4.is_link_local() // 169.254.0.0/16 - || ipv4.is_broadcast() // 255.255.255.255 - || ipv4.is_unspecified() // 0.0.0.0 - || is_shared_address(ipv4) // 100.64.0.0/10 (CGNAT) - || is_documentation_v4(ipv4) // 192.0.2.0/24, 198.51.100.0/24, 203.0.113.0/24 - || is_cloud_metadata(ipv4) // 169.254.169.254 - } - IpAddr::V6(ipv6) => { - ipv6.is_loopback() // ::1 - || ipv6.is_unspecified() // :: - || is_unique_local(ipv6) // fc00::/7 - || is_link_local_v6(ipv6) // fe80::/10 - } - } -} - -/// Check for CGNAT shared address space (100.64.0.0/10) -fn is_shared_address(ip: &Ipv4Addr) -> bool { - let octets = ip.octets(); - octets[0] == 100 && (octets[1] & 0xC0) == 64 -} - -/// Check for IPv4 documentation addresses -fn is_documentation_v4(ip: &Ipv4Addr) -> bool { - let octets = ip.octets(); - // 192.0.2.0/24 (TEST-NET-1) - (octets[0] == 192 && octets[1] == 0 && octets[2] == 2) - // 198.51.100.0/24 (TEST-NET-2) - || (octets[0] == 198 && octets[1] == 51 && octets[2] == 100) - // 203.0.113.0/24 (TEST-NET-3) - || (octets[0] == 203 && octets[1] == 0 && octets[2] == 113) -} - -/// Check for cloud metadata endpoint (169.254.169.254) -fn is_cloud_metadata(ip: &Ipv4Addr) -> bool { - ip.octets() == [169, 254, 169, 254] -} - -/// Check for IPv6 unique local addresses (fc00::/7) -fn is_unique_local(ip: &Ipv6Addr) -> bool { - let segments = ip.segments(); - (segments[0] & 0xfe00) == 0xfc00 -} - -/// Check for IPv6 link-local addresses (fe80::/10) -fn is_link_local_v6(ip: &Ipv6Addr) -> bool { - let segments = ip.segments(); - (segments[0] & 0xffc0) == 0xfe80 -} - -/// Validate a URL for SSRF protection. -/// Returns the validated URL or an error. -/// -/// For testing purposes, HTTP is allowed for localhost/127.0.0.1 only. -/// In production, only HTTPS should be used. -pub(crate) fn validate_url(url_str: &str) -> Result { - let url = Url::parse(url_str) - .map_err(|e| JwksError::InvalidUrl(format!("Failed to parse URL: {}", e)))?; - - // Check if this is a localhost URL (allowed for testing with HTTP) - let is_localhost = match url.host_str() { - Some("localhost") | Some("127.0.0.1") => true, - Some(host) => host == "::1" || host == "[::1]", - None => false, - }; - - // Only allow HTTPS (except for localhost in tests) - match url.scheme() { - "https" => {} - "http" => { - if !is_localhost { - return Err(JwksError::SsrfBlocked( - "Only HTTPS URLs are allowed for JWKS endpoints".to_string(), - )); - } - // HTTP localhost is allowed - skip further checks and return early - return Ok(url); - } - scheme => { - return Err(JwksError::SsrfBlocked(format!( - "Invalid URL scheme '{}'. Only HTTPS is allowed", - scheme - ))); - } - } - - // For HTTPS URLs, check for private/internal addresses - if let Some(host) = url.host() { - match host { - url::Host::Ipv4(ip) => { - if is_private_ip(&IpAddr::V4(ip)) { - return Err(JwksError::SsrfBlocked(format!( - "Private/internal IP addresses are not allowed: {}", - ip - ))); - } - } - url::Host::Ipv6(ip) => { - if is_private_ip(&IpAddr::V6(ip)) { - return Err(JwksError::SsrfBlocked(format!( - "Private/internal IP addresses are not allowed: {}", - ip - ))); - } - } - url::Host::Domain(domain) => { - // Block common internal hostnames (except localhost which was handled above) - let lower = domain.to_lowercase(); - if lower == "metadata" - || lower == "metadata.google.internal" - || lower.ends_with(".internal") - || lower.ends_with(".local") - { - return Err(JwksError::SsrfBlocked(format!( - "Internal hostnames are not allowed: {}", - domain - ))); - } - } - } - } - - Ok(url) -} - -/// OIDC discovery document (subset of fields we need). -#[derive(Debug, serde::Deserialize)] -struct OidcDiscovery { - jwks_uri: String, - #[allow(dead_code)] - issuer: String, -} - -/// Cached JWKS with expiration tracking. -struct CachedJwks { - jwks: JwkSet, - fetched_at: Instant, - ttl: Duration, -} - -impl CachedJwks { - fn is_expired(&self) -> bool { - self.fetched_at.elapsed() > self.ttl - } -} - -/// JWKS provider with caching and automatic refresh. -pub(crate) struct JwksProvider { - /// HTTP client for fetching JWKS - client: reqwest::Client, - /// JWKS endpoint URL (validated) - jwks_uri: String, - /// Cached JWKS - cache: RwLock>, - /// Cache TTL - ttl: Duration, -} - -impl JwksProvider { - /// Create a new JWKS provider with explicit JWKS URI. - /// The URL is validated for SSRF protection. - pub fn new(jwks_uri: impl Into, ttl: Duration) -> Result { - let jwks_uri = jwks_uri.into(); - - // Validate URL for SSRF protection - validate_url(&jwks_uri)?; - - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(10)) - .redirect(reqwest::redirect::Policy::none()) // Prevent SSRF via redirects - .build() - .map_err(|e| JwksError::HttpClientError(e.to_string()))?; - - Ok(Self { - client, - jwks_uri, - cache: RwLock::new(None), - ttl, - }) - } - - /// Create a new JWKS provider using OIDC discovery. - /// Both issuer and discovered JWKS URI are validated for SSRF protection. - pub async fn from_issuer(issuer: &str, ttl: Duration) -> Result { - // Normalize and validate issuer URL - let issuer = issuer.trim_end_matches('/'); - let discovery_url = format!("{}/.well-known/openid-configuration", issuer); - - // Validate discovery URL for SSRF protection - validate_url(&discovery_url)?; - - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(10)) - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(|e| JwksError::HttpClientError(e.to_string()))?; - - info!("Fetching OIDC discovery from: {}", discovery_url); - - let response = client - .get(&discovery_url) - .send() - .await - .map_err(|e| JwksError::DiscoveryFetch(e.to_string()))?; - - if !response.status().is_success() { - return Err(JwksError::DiscoveryFetch(format!( - "HTTP {}", - response.status() - ))); - } - - // Check response size before parsing - if let Some(content_length) = response.content_length() { - if content_length > MAX_JWKS_RESPONSE_SIZE { - return Err(JwksError::ResponseTooLarge( - content_length, - MAX_JWKS_RESPONSE_SIZE, - )); - } - } - - let discovery: OidcDiscovery = response - .json() - .await - .map_err(|e| JwksError::DiscoveryParse(e.to_string()))?; - - // Validate discovered JWKS URI for SSRF protection - validate_url(&discovery.jwks_uri)?; - - info!("Discovered JWKS URI: {}", discovery.jwks_uri); - - Ok(Self { - client, - jwks_uri: discovery.jwks_uri, - cache: RwLock::new(None), - ttl, - }) - } - - /// Get the JWKS URI. - #[allow(dead_code)] - pub fn jwks_uri(&self) -> &str { - &self.jwks_uri - } - - /// Fetch JWKS from the endpoint. - /// Response size is limited to prevent DoS attacks. - async fn fetch_jwks(&self) -> Result { - debug!("Fetching JWKS from: {}", self.jwks_uri); - - let response = self - .client - .get(&self.jwks_uri) - .send() - .await - .map_err(|e| JwksError::JwksFetch(e.to_string()))?; - - if !response.status().is_success() { - return Err(JwksError::JwksFetch(format!("HTTP {}", response.status()))); - } - - // Check response size before reading body - if let Some(content_length) = response.content_length() { - if content_length > MAX_JWKS_RESPONSE_SIZE { - return Err(JwksError::ResponseTooLarge( - content_length, - MAX_JWKS_RESPONSE_SIZE, - )); - } - } - - // Read body with size limit - let bytes = response - .bytes() - .await - .map_err(|e| JwksError::JwksFetch(e.to_string()))?; - - if bytes.len() as u64 > MAX_JWKS_RESPONSE_SIZE { - return Err(JwksError::ResponseTooLarge( - bytes.len() as u64, - MAX_JWKS_RESPONSE_SIZE, - )); - } - - let jwks: JwkSet = - serde_json::from_slice(&bytes).map_err(|e| JwksError::JwksParse(e.to_string()))?; - - debug!("Fetched JWKS with {} keys", jwks.keys.len()); - Ok(jwks) - } - - /// Get the cached JWKS, refreshing if expired or not present. - pub async fn get_jwks(&self) -> Result { - // Check cache first - { - let cache = self.cache.read(); - if let Some(cached) = cache.as_ref() { - if !cached.is_expired() { - return Ok(cached.jwks.clone()); - } - } - } - - // Cache miss or expired, fetch new JWKS - let jwks = self.fetch_jwks().await?; - - // Update cache - { - let mut cache = self.cache.write(); - *cache = Some(CachedJwks { - jwks: jwks.clone(), - fetched_at: Instant::now(), - ttl: self.ttl, - }); - } - - Ok(jwks) - } - - /// Get a specific key by kid (key ID). - pub async fn get_key(&self, kid: &str) -> Result { - let jwks = self.get_jwks().await?; - - // First try to find by kid - if let Some(key) = jwks.find(kid) { - return Ok(key.clone()); - } - - // Key not found - try refreshing the cache in case keys were rotated - warn!("Key {} not found in cached JWKS, refreshing...", kid); - let jwks = self.fetch_jwks().await?; - - // Update cache - { - let mut cache = self.cache.write(); - *cache = Some(CachedJwks { - jwks: jwks.clone(), - fetched_at: Instant::now(), - ttl: self.ttl, - }); - } - - jwks.find(kid) - .cloned() - .ok_or_else(|| JwksError::KeyNotFound(kid.to_string())) - } - - /// Force refresh the JWKS cache. - #[allow(dead_code)] - pub async fn refresh(&self) -> Result<(), JwksError> { - let jwks = self.fetch_jwks().await?; - let mut cache = self.cache.write(); - *cache = Some(CachedJwks { - jwks, - fetched_at: Instant::now(), - ttl: self.ttl, - }); - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_cached_jwks_expiration() { - let jwks = JwkSet { keys: vec![] }; - let cached = CachedJwks { - jwks, - fetched_at: Instant::now() - Duration::from_secs(100), - ttl: Duration::from_secs(60), - }; - assert!(cached.is_expired()); - - let jwks = JwkSet { keys: vec![] }; - let cached = CachedJwks { - jwks, - fetched_at: Instant::now(), - ttl: Duration::from_secs(60), - }; - assert!(!cached.is_expired()); - } - - #[test] - fn test_validate_url_https_required() { - // Valid HTTPS URL - assert!(validate_url("https://example.com/.well-known/jwks.json").is_ok()); - - // HTTP not allowed for non-localhost - assert!(validate_url("http://example.com/.well-known/jwks.json").is_err()); - - // HTTP allowed for localhost (testing) - assert!(validate_url("http://localhost:8080/jwks").is_ok()); - assert!(validate_url("http://127.0.0.1:8080/jwks").is_ok()); - - // Invalid schemes - assert!(validate_url("ftp://example.com/jwks").is_err()); - assert!(validate_url("file:///etc/passwd").is_err()); - } - - #[test] - fn test_validate_url_blocks_private_ips() { - // Private IPv4 ranges - assert!(validate_url("https://10.0.0.1/jwks").is_err()); - assert!(validate_url("https://172.16.0.1/jwks").is_err()); - assert!(validate_url("https://192.168.1.1/jwks").is_err()); - - // Loopback - assert!(validate_url("https://127.0.0.1/jwks").is_err()); - - // Link-local - assert!(validate_url("https://169.254.1.1/jwks").is_err()); - - // Cloud metadata endpoint - assert!(validate_url("https://169.254.169.254/jwks").is_err()); - - // CGNAT - assert!(validate_url("https://100.64.0.1/jwks").is_err()); - } - - #[test] - fn test_validate_url_blocks_internal_hostnames() { - assert!(validate_url("https://metadata/jwks").is_err()); - assert!(validate_url("https://metadata.google.internal/jwks").is_err()); - assert!(validate_url("https://internal.example.internal/jwks").is_err()); - assert!(validate_url("https://printer.local/jwks").is_err()); - } - - #[test] - fn test_is_private_ip() { - use std::net::Ipv4Addr; - - // Private ranges - assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)))); - assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(172, 16, 0, 1)))); - assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)))); - - // Loopback - assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))); - - // Cloud metadata - assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new( - 169, 254, 169, 254 - )))); - - // Public IPs should not be blocked - assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)))); - assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)))); - } -} diff --git a/sgl-model-gateway/src/auth/jwt.rs b/sgl-model-gateway/src/auth/jwt.rs deleted file mode 100644 index edb1b51d6..000000000 --- a/sgl-model-gateway/src/auth/jwt.rs +++ /dev/null @@ -1,536 +0,0 @@ -//! JWT validation for control plane authentication. -//! -//! Supports: -//! - RS256, RS384, RS512 (RSA) -//! - ES256, ES384 (ECDSA) -//! - Audience and issuer validation -//! - Role extraction from claims -//! -//! Security features: -//! - Algorithm verification (token alg must match key algorithm) -//! - Optional JTI tracking for replay protection -//! - SSRF protection via JWKS provider - -use std::{ - collections::HashMap, - num::NonZeroUsize, - sync::Arc, - time::{Duration, Instant}, -}; - -use jsonwebtoken::{ - decode, decode_header, - jwk::{AlgorithmParameters, Jwk}, - Algorithm, DecodingKey, TokenData, Validation, -}; -use lru::LruCache; -use parking_lot::Mutex; -use serde::{Deserialize, Serialize}; -use tracing::{debug, warn}; - -use super::{ - config::{JwtConfig, Role}, - jwks::{JwksError, JwksProvider}, -}; - -/// Default size for JTI cache (number of tokens to track) -const DEFAULT_JTI_CACHE_SIZE: usize = 10_000; - -/// Error types for JWT validation. -#[derive(Debug, thiserror::Error)] -pub enum JwtValidatorError { - #[error("Invalid token format")] - InvalidFormat, - - #[error("Token header missing 'kid' claim")] - MissingKid, - - #[error("Failed to get signing key: {0}")] - KeyError(#[from] JwksError), - - #[error("Unsupported algorithm: {0:?}")] - UnsupportedAlgorithm(Algorithm), - - #[error("Token validation failed: {0}")] - ValidationFailed(String), - - #[error("Failed to decode token: {0}")] - DecodeFailed(#[from] jsonwebtoken::errors::Error), - - #[error("Failed to extract role from claims")] - RoleExtractionFailed, - - #[error("No role mapping found for: {0}")] - NoRoleMapping(String), - - #[error("Algorithm mismatch: token uses {token_alg:?} but key requires {key_alg:?}")] - AlgorithmMismatch { - token_alg: Algorithm, - key_alg: Algorithm, - }, - - #[error("Token replay detected: JTI '{0}' has already been used")] - TokenReplay(String), -} - -/// Standard JWT claims we extract. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub(crate) struct StandardClaims { - /// Subject (user ID) - pub sub: Option, - - /// Issuer - pub iss: Option, - - /// Audience (can be string or array) - #[serde(default)] - pub aud: Audience, - - /// Expiration time - pub exp: Option, - - /// Issued at - pub iat: Option, - - /// Not before - pub nbf: Option, - - /// JWT ID - pub jti: Option, - - /// Email (common claim) - pub email: Option, - - /// Name (common claim) - pub name: Option, - - /// Preferred username (OIDC claim) - pub preferred_username: Option, - - /// All other claims - #[serde(flatten)] - pub extra: HashMap, -} - -/// Audience claim can be a single string or an array. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum Audience { - Single(String), - Multiple(Vec), - #[default] - None, -} - -impl Audience { - #[allow(dead_code)] - pub fn contains(&self, aud: &str) -> bool { - match self { - Audience::Single(s) => s == aud, - Audience::Multiple(v) => v.iter().any(|s| s == aud), - Audience::None => false, - } - } -} - -/// Validated token with extracted claims and role. -#[derive(Debug, Clone)] -pub struct ValidatedToken { - /// Subject (user ID) - pub subject: String, - - /// Issuer - pub issuer: String, - - /// Assigned role - pub role: Role, - - /// Email if present - pub email: Option, - - /// Display name if present - pub name: Option, - - /// Full claims for additional processing - #[allow(dead_code)] - pub(crate) claims: StandardClaims, -} - -/// JTI (JWT ID) cache entry with expiration tracking. -struct JtiCacheEntry { - /// When the token expires (for cleanup) - expires_at: Instant, -} - -/// JWT validator with JWKS integration. -pub struct JwtValidator { - /// JWKS provider for key fetching - jwks_provider: Arc, - - /// JWT configuration - config: JwtConfig, - - /// Pre-configured validation settings - validation: Validation, - - /// JTI cache for replay protection (optional) - /// Maps JTI -> expiration time - jti_cache: Option>>, - - /// Whether to enable JTI replay protection - enable_jti_check: bool, -} - -impl JwtValidator { - /// Create a new JWT validator with explicit JWKS URI. - #[allow(dead_code)] - pub(crate) fn new(config: JwtConfig, jwks_provider: Arc) -> Self { - Self::new_with_options(config, jwks_provider, false) - } - - /// Create a new JWT validator with optional JTI replay protection. - pub(crate) fn new_with_options( - config: JwtConfig, - jwks_provider: Arc, - enable_jti_check: bool, - ) -> Self { - let mut validation = Validation::default(); - - // Set audience - validation.set_audience(&[&config.audience]); - - // Set issuer - validation.set_issuer(&[&config.issuer]); - - // Set leeway for clock skew - validation.leeway = config.leeway_secs; - - // We'll set the algorithm per-token based on the key - validation.algorithms = vec![ - Algorithm::RS256, - Algorithm::RS384, - Algorithm::RS512, - Algorithm::ES256, - Algorithm::ES384, - ]; - - let jti_cache = if enable_jti_check { - Some(Mutex::new(LruCache::new( - NonZeroUsize::new(DEFAULT_JTI_CACHE_SIZE).unwrap(), - ))) - } else { - None - }; - - Self { - jwks_provider, - config, - validation, - jti_cache, - enable_jti_check, - } - } - - /// Create a new JWT validator using OIDC discovery. - pub async fn from_config(config: JwtConfig) -> Result { - Self::from_config_with_options(config, false).await - } - - /// Create a new JWT validator with optional JTI replay protection. - pub async fn from_config_with_options( - config: JwtConfig, - enable_jti_check: bool, - ) -> Result { - let ttl = Duration::from_secs(config.jwks_cache_ttl_secs); - - let jwks_provider = if let Some(jwks_uri) = &config.jwks_uri { - Arc::new(JwksProvider::new(jwks_uri.clone(), ttl)?) - } else { - Arc::new(JwksProvider::from_issuer(&config.issuer, ttl).await?) - }; - - Ok(Self::new_with_options( - config, - jwks_provider, - enable_jti_check, - )) - } - - /// Validate a JWT token and extract claims. - pub async fn validate(&self, token: &str) -> Result { - // Decode header to get kid and algorithm - let header = decode_header(token)?; - let kid = header.kid.ok_or(JwtValidatorError::MissingKid)?; - let token_algorithm = header.alg; - - debug!( - "Validating JWT with kid: {}, alg: {:?}", - kid, token_algorithm - ); - - // Get the signing key - let jwk = self.jwks_provider.get_key(&kid).await?; - - // Determine algorithm from JWK - let key_algorithm = Self::jwk_to_algorithm(&jwk)?; - - // SECURITY: Verify token algorithm matches key algorithm - // This prevents algorithm confusion attacks - if token_algorithm != key_algorithm { - warn!( - "Algorithm mismatch: token uses {:?} but key requires {:?}", - token_algorithm, key_algorithm - ); - return Err(JwtValidatorError::AlgorithmMismatch { - token_alg: token_algorithm, - key_alg: key_algorithm, - }); - } - - // Create decoding key from JWK - let decoding_key = Self::jwk_to_decoding_key(&jwk)?; - - // Create validation with specific algorithm - let mut validation = self.validation.clone(); - validation.algorithms = vec![key_algorithm]; - - // Decode and validate token - let token_data: TokenData = decode(token, &decoding_key, &validation)?; - - let claims = token_data.claims; - - // Check JTI for replay protection if enabled - if self.enable_jti_check { - if let Some(jti) = &claims.jti { - self.check_jti_replay(jti, &claims)?; - } - } - - // Extract subject - let subject = claims - .sub - .clone() - .or_else(|| claims.email.clone()) - .or_else(|| claims.preferred_username.clone()) - .unwrap_or_else(|| "unknown".to_string()); - - // Extract issuer - let issuer = claims - .iss - .clone() - .unwrap_or_else(|| self.config.issuer.clone()); - - // Extract role - let role = self.extract_role(&claims)?; - - debug!( - "JWT validated: subject={}, issuer={}, role={:?}", - subject, issuer, role - ); - - Ok(ValidatedToken { - subject, - issuer, - role, - email: claims.email.clone(), - name: claims.name.clone(), - claims, - }) - } - - /// Check if a JTI has been used before (replay protection). - fn check_jti_replay( - &self, - jti: &str, - claims: &StandardClaims, - ) -> Result<(), JwtValidatorError> { - let Some(cache) = &self.jti_cache else { - return Ok(()); - }; - - let mut cache = cache.lock(); - - // Clean up expired entries first (lazy cleanup) - let now = Instant::now(); - - // Check if JTI exists and is still valid - if let Some(entry) = cache.get(jti) { - if entry.expires_at > now { - // Token is still valid but JTI was already used - return Err(JwtValidatorError::TokenReplay(jti.to_string())); - } - // Entry expired, remove it - cache.pop(jti); - } - - // Calculate expiration time from claims - let expires_at = if let Some(exp) = claims.exp { - // exp is Unix timestamp - let exp_duration = Duration::from_secs(exp); - let now_unix = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default(); - - if exp_duration > now_unix { - now + (exp_duration - now_unix) - } else { - // Already expired, but we still track it briefly - now + Duration::from_secs(60) - } - } else { - // No exp claim, use default TTL - now + Duration::from_secs(3600) - }; - - // Record this JTI - cache.put(jti.to_string(), JtiCacheEntry { expires_at }); - - Ok(()) - } - - /// Extract role from claims using configured role claim and mapping. - fn extract_role(&self, claims: &StandardClaims) -> Result { - // Try to get the role claim value - let role_value = claims.extra.get(&self.config.role_claim); - - let role_strings: Vec = match role_value { - Some(serde_json::Value::String(s)) => vec![s.clone()], - Some(serde_json::Value::Array(arr)) => arr - .iter() - .filter_map(|v| v.as_str().map(String::from)) - .collect(), - None => { - // Try alternate claim names - let alternates = ["role", "roles", "groups", "group"]; - let mut found = Vec::new(); - for alt in alternates { - if let Some(v) = claims.extra.get(alt) { - match v { - serde_json::Value::String(s) => found.push(s.clone()), - serde_json::Value::Array(arr) => found - .extend(arr.iter().filter_map(|v| v.as_str().map(String::from))), - _ => {} - } - } - } - found - } - _ => Vec::new(), - }; - - // If no role mapping configured, check for direct "admin" or "user" values - if self.config.role_mapping.is_empty() { - for role_str in &role_strings { - if let Ok(role) = role_str.parse::() { - return Ok(role); - } - } - // Default to User if no explicit role found - warn!("No role found in JWT claims, defaulting to User"); - return Ok(Role::User); - } - - // Use role mapping - for role_str in &role_strings { - if let Some(role) = self.config.role_mapping.get(role_str) { - return Ok(*role); - } - } - - // Check if any mapped role is admin - if so, we need explicit mapping - // Otherwise default to User for safety - warn!( - "No matching role mapping found for {:?}, defaulting to User", - role_strings - ); - Ok(Role::User) - } - - /// Convert a JWK to a DecodingKey. - fn jwk_to_decoding_key(jwk: &Jwk) -> Result { - match &jwk.algorithm { - AlgorithmParameters::RSA(rsa) => Ok(DecodingKey::from_rsa_components(&rsa.n, &rsa.e)?), - AlgorithmParameters::EllipticCurve(ec) => { - Ok(DecodingKey::from_ec_components(&ec.x, &ec.y)?) - } - AlgorithmParameters::OctetKey(_) => { - Err(JwtValidatorError::UnsupportedAlgorithm(Algorithm::HS256)) - } - AlgorithmParameters::OctetKeyPair(_) => { - Err(JwtValidatorError::UnsupportedAlgorithm(Algorithm::EdDSA)) - } - } - } - - /// Determine the algorithm from a JWK. - fn jwk_to_algorithm(jwk: &Jwk) -> Result { - // First check if algorithm is explicitly specified - if let Some(alg) = &jwk.common.key_algorithm { - return Ok(match alg { - jsonwebtoken::jwk::KeyAlgorithm::RS256 => Algorithm::RS256, - jsonwebtoken::jwk::KeyAlgorithm::RS384 => Algorithm::RS384, - jsonwebtoken::jwk::KeyAlgorithm::RS512 => Algorithm::RS512, - jsonwebtoken::jwk::KeyAlgorithm::ES256 => Algorithm::ES256, - jsonwebtoken::jwk::KeyAlgorithm::ES384 => Algorithm::ES384, - other => { - return Err(JwtValidatorError::ValidationFailed(format!( - "Unsupported key algorithm: {:?}", - other - ))) - } - }); - } - - // Infer from key type - match &jwk.algorithm { - AlgorithmParameters::RSA(_) => Ok(Algorithm::RS256), // Default RSA to RS256 - AlgorithmParameters::EllipticCurve(ec) => { - use jsonwebtoken::jwk::EllipticCurve; - match ec.curve { - EllipticCurve::P256 => Ok(Algorithm::ES256), - EllipticCurve::P384 => Ok(Algorithm::ES384), - // Other curves not supported for ECDSA - _ => Err(JwtValidatorError::ValidationFailed(format!( - "Unsupported EC curve: {:?}", - ec.curve - ))), - } - } - _ => Err(JwtValidatorError::ValidationFailed( - "Cannot determine algorithm from key".to_string(), - )), - } - } - - /// Get a reference to the JWKS provider. - #[allow(dead_code)] - pub(crate) fn jwks_provider(&self) -> &Arc { - &self.jwks_provider - } - - /// Get a reference to the JWT config. - pub fn config(&self) -> &JwtConfig { - &self.config - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_audience_contains() { - let single = Audience::Single("api://test".to_string()); - assert!(single.contains("api://test")); - assert!(!single.contains("other")); - - let multiple = - Audience::Multiple(vec!["api://test".to_string(), "api://other".to_string()]); - assert!(multiple.contains("api://test")); - assert!(multiple.contains("api://other")); - assert!(!multiple.contains("unknown")); - - let none = Audience::None; - assert!(!none.contains("anything")); - } -} diff --git a/sgl-model-gateway/src/auth/middleware.rs b/sgl-model-gateway/src/auth/middleware.rs deleted file mode 100644 index 5e82e1b3e..000000000 --- a/sgl-model-gateway/src/auth/middleware.rs +++ /dev/null @@ -1,399 +0,0 @@ -//! Control plane authentication middleware. -//! -//! Provides middleware for authenticating and authorizing access to control plane APIs. -//! Supports both JWT/OIDC tokens and API keys. - -use std::sync::Arc; - -use axum::{ - body::Body, - extract::{Request, State}, - http::{header, StatusCode}, - middleware::Next, - response::{IntoResponse, Response}, -}; -use tracing::{debug, error, info, warn}; - -use super::{ - audit::{AuditContext, AuditLogger}, - config::{ControlPlaneAuthConfig, Role}, - jwt::JwtValidator, -}; -use crate::middleware::RequestId; - -/// Authenticated principal information. -#[derive(Debug, Clone)] -pub struct Principal { - /// Subject identifier (user ID, email, or API key ID) - pub id: String, - - /// Display name if available - pub name: Option, - - /// Authentication method used - pub auth_method: AuthMethod, - - /// Assigned role - pub role: Role, -} - -/// Authentication method used to authenticate the principal. -#[derive(Debug, Clone)] -pub enum AuthMethod { - /// JWT/OIDC token from external IDP - Jwt { issuer: String }, - /// API key for service accounts - ApiKey { key_id: String }, -} - -impl std::fmt::Display for AuthMethod { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - AuthMethod::Jwt { issuer } => write!(f, "jwt:{}", issuer), - AuthMethod::ApiKey { key_id } => write!(f, "api_key:{}", key_id), - } - } -} - -/// Extension trait for extracting Principal from request extensions. -pub trait PrincipalExt { - fn principal(&self) -> Option<&Principal>; -} - -impl PrincipalExt for http::Request { - fn principal(&self) -> Option<&Principal> { - self.extensions().get::() - } -} - -/// State for the control plane authentication middleware. -#[derive(Clone)] -pub struct ControlPlaneAuthState { - /// Authentication configuration - pub config: ControlPlaneAuthConfig, - - /// JWT validator (if JWT auth is configured) - pub jwt_validator: Option>, - - /// Audit logger - pub audit_logger: AuditLogger, -} - -impl ControlPlaneAuthState { - /// Create a new control plane auth state. - pub fn new(config: ControlPlaneAuthConfig, jwt_validator: Option>) -> Self { - let audit_logger = AuditLogger::new(config.audit_enabled); - Self { - config, - jwt_validator, - audit_logger, - } - } - - /// Create from config, initializing JWT validator if needed. - pub async fn from_config( - config: ControlPlaneAuthConfig, - ) -> Result> { - let jwt_validator = if let Some(jwt_config) = &config.jwt { - Some(Arc::new( - JwtValidator::from_config(jwt_config.clone()).await?, - )) - } else { - None - }; - - Ok(Self::new(config, jwt_validator)) - } - - /// Try to initialize control plane auth from config. - /// - /// Returns `Some(state)` if auth is configured and initialized successfully. - /// Returns `None` if auth is not configured or initialization fails (with error logged). - pub async fn try_init(config: Option<&ControlPlaneAuthConfig>) -> Option { - let config = config.filter(|c| c.is_enabled())?; - - info!("Initializing control plane authentication..."); - match Self::from_config(config.clone()).await { - Ok(state) => { - if config.has_jwt() { - info!("Control plane JWT/OIDC authentication enabled"); - } - if config.has_api_keys() { - info!( - "Control plane API key authentication enabled ({} keys)", - config.api_keys.len() - ); - } - if config.audit_enabled { - info!("Control plane audit logging enabled"); - } - Some(state) - } - Err(e) => { - error!( - "Failed to initialize control plane auth: {}. Falling back to simple API key auth.", - e - ); - None - } - } - } - - /// Check if authentication is required. - pub fn is_auth_required(&self) -> bool { - self.config.is_enabled() - } -} - -/// Check admin role and log denial if not admin. -/// Returns Some(Response) if denied, None if allowed. -fn check_admin_role( - principal_id: &str, - auth_method: &str, - role: Role, - method: &str, - path: &str, - request_id: Option<&str>, - audit_logger: &AuditLogger, -) -> Option { - if role.is_admin() { - return None; - } - - warn!( - "{} {} has role {:?} but admin is required for control plane access", - auth_method, principal_id, role - ); - let ctx = AuditContext::new(principal_id, auth_method, role, method, path, request_id); - audit_logger.log_denied(&ctx, "Admin role required for control plane access"); - - Some( - ( - StatusCode::FORBIDDEN, - "Admin role required for control plane access", - ) - .into_response(), - ) -} - -/// Log successful authentication. -fn log_auth_success( - principal: &Principal, - auth_method: &str, - method: &str, - path: &str, - request_id: Option<&str>, - audit_logger: &AuditLogger, -) { - debug!( - "{} authentication successful for {} with role {:?}", - auth_method, principal.id, principal.role - ); - let ctx = AuditContext::new( - &principal.id, - auth_method, - principal.role, - method, - path, - request_id, - ); - audit_logger.log_success(&ctx, None); -} - -/// Control plane authentication middleware. -/// -/// This middleware: -/// 1. Extracts the Bearer token from the Authorization header -/// 2. Attempts JWT validation first (if configured) -/// 3. Falls back to API key validation (if configured) -/// 4. Checks if the authenticated principal has admin role -/// 5. Logs audit events for control plane access -/// -/// Returns 401 Unauthorized if authentication fails. -/// Returns 403 Forbidden if the user doesn't have admin role. -pub async fn control_plane_auth_middleware( - State(auth_state): State, - mut request: Request, - next: Next, -) -> Response { - // If no authentication is configured, allow through (backward compatibility) - if !auth_state.is_auth_required() { - return next.run(request).await; - } - - let method = request.method().to_string(); - let path = request.uri().path().to_string(); - let request_id = request.extensions().get::().map(|r| r.0.clone()); - - // Extract Bearer token from Authorization header - let token = request - .headers() - .get(header::AUTHORIZATION) - .and_then(|h| h.to_str().ok()) - .and_then(|h| h.strip_prefix("Bearer ")); - - let Some(token) = token else { - debug!("Missing or invalid Authorization header for control plane API"); - auth_state.audit_logger.log_auth_failure( - &method, - &path, - "Missing or invalid Authorization header", - request_id.as_deref(), - ); - return ( - StatusCode::UNAUTHORIZED, - [("WWW-Authenticate", "Bearer realm=\"control-plane\"")], - "Missing or invalid Authorization header", - ) - .into_response(); - }; - - // Try JWT validation first - if let Some(jwt_validator) = &auth_state.jwt_validator { - match jwt_validator.validate(token).await { - Ok(validated_token) => { - if let Some(resp) = check_admin_role( - &validated_token.subject, - "jwt", - validated_token.role, - &method, - &path, - request_id.as_deref(), - &auth_state.audit_logger, - ) { - return resp; - } - - let principal = Principal { - id: validated_token.subject.clone(), - name: validated_token.name.clone(), - auth_method: AuthMethod::Jwt { - issuer: validated_token.issuer.clone(), - }, - role: validated_token.role, - }; - - log_auth_success( - &principal, - "jwt", - &method, - &path, - request_id.as_deref(), - &auth_state.audit_logger, - ); - request.extensions_mut().insert(principal); - return next.run(request).await; - } - Err(e) => { - // If the token looks like a JWT (3 parts separated by dots), it's likely - // an invalid JWT, not an API key. Fail fast with a specific error - // instead of silently falling back. This provides better feedback. - if token.split('.').count() == 3 { - warn!("Invalid JWT provided: {}. Not falling back to API key.", e); - auth_state.audit_logger.log_auth_failure( - &method, - &path, - &format!("Invalid JWT: {}", e), - request_id.as_deref(), - ); - return ( - StatusCode::UNAUTHORIZED, - [("WWW-Authenticate", "Bearer realm=\"control-plane\"")], - format!("Invalid JWT: {}", e), - ) - .into_response(); - } - debug!("JWT validation failed: {}, trying API key", e); - } - } - } - - // Try API key validation - if let Some(api_key_entry) = auth_state.config.find_api_key(token) { - if let Some(resp) = check_admin_role( - &api_key_entry.id, - "api_key", - api_key_entry.role, - &method, - &path, - request_id.as_deref(), - &auth_state.audit_logger, - ) { - return resp; - } - - let principal = Principal { - id: api_key_entry.id.clone(), - name: Some(api_key_entry.name.clone()), - auth_method: AuthMethod::ApiKey { - key_id: api_key_entry.id.clone(), - }, - role: api_key_entry.role, - }; - - log_auth_success( - &principal, - "api_key", - &method, - &path, - request_id.as_deref(), - &auth_state.audit_logger, - ); - request.extensions_mut().insert(principal); - return next.run(request).await; - } - - // Authentication failed - debug!("Control plane authentication failed: invalid token"); - auth_state.audit_logger.log_auth_failure( - &method, - &path, - "Invalid token", - request_id.as_deref(), - ); - - ( - StatusCode::UNAUTHORIZED, - [("WWW-Authenticate", "Bearer realm=\"control-plane\"")], - "Invalid authentication token", - ) - .into_response() -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_auth_method_display() { - let jwt = AuthMethod::Jwt { - issuer: "https://example.com".to_string(), - }; - assert_eq!(jwt.to_string(), "jwt:https://example.com"); - - let api_key = AuthMethod::ApiKey { - key_id: "key-123".to_string(), - }; - assert_eq!(api_key.to_string(), "api_key:key-123"); - } - - #[test] - fn test_control_plane_auth_state_no_config() { - let config = ControlPlaneAuthConfig::default(); - let state = ControlPlaneAuthState::new(config, None); - assert!(!state.is_auth_required()); - } - - #[test] - fn test_control_plane_auth_state_with_api_keys() { - use super::super::config::ApiKeyEntry; - - let config = ControlPlaneAuthConfig { - jwt: None, - api_keys: vec![ApiKeyEntry::new("test", "Test Key", "secret", Role::Admin)], - audit_enabled: true, - }; - let state = ControlPlaneAuthState::new(config, None); - assert!(state.is_auth_required()); - } -} diff --git a/sgl-model-gateway/src/auth/mod.rs b/sgl-model-gateway/src/auth/mod.rs deleted file mode 100644 index 947a2e895..000000000 --- a/sgl-model-gateway/src/auth/mod.rs +++ /dev/null @@ -1,20 +0,0 @@ -//! Authentication and authorization module for control plane APIs. -//! -//! This module provides: -//! - JWT/OIDC authentication for external IDP integration -//! - API key authentication with role-based access -//! - Audit logging for control plane operations -//! - Middleware for securing admin and worker routes - -mod audit; -mod config; -mod jwks; -mod jwt; -mod middleware; - -pub use audit::{AuditEvent, AuditLogger, AuditOutcome}; -pub use config::{ApiKeyEntry, ControlPlaneAuthConfig, JwtConfig, Role}; -pub use jwt::{JwtValidator, JwtValidatorError}; -pub use middleware::{ - control_plane_auth_middleware, ControlPlaneAuthState, Principal, PrincipalExt, -}; diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs index da8bb985e..2693b364e 100644 --- a/sgl-model-gateway/src/lib.rs +++ b/sgl-model-gateway/src/lib.rs @@ -1,5 +1,5 @@ pub mod app_context; -pub mod auth; +pub use smg_auth as auth; pub mod config; pub mod core; pub mod data_connector;