[smg] use already published auth crate for better compilation speed (#17723)

2026-01-25 14:52:42 -08:00
parent 344eeaee90
commit bf139af37a
8 changed files with 2 additions and 2184 deletions
@@ -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"] }
+smg-auth = "1.0.0"
 minijinja-contrib = { version = "2.0", features = ["pycompat"] }
 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"
@@ -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<Utc>,
    /// 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<String>,
    /// Operation outcome
    pub outcome: AuditOutcome,
    /// Request ID for correlation
    #[serde(skip_serializing_if = "Option::is_none")]
    pub request_id: Option<String>,
    /// Additional details or error message
    #[serde(skip_serializing_if = "Option::is_none")]
    pub details: Option<String>,
 }
 /// 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"));
    }
 }
@@ -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<Self, Self::Err> {
        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<String>,
    /// 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<String, Role>,
    /// 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<String>, audience: impl Into<String>) -> 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<String>) -> 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<String>, 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<String>,
        name: impl Into<String>,
        key: impl Into<String>,
        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<JwtConfig>,
    /// 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<ApiKeyEntry>,
    /// 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::<Role>().unwrap(), Role::Admin);
        assert_eq!("ADMIN".parse::<Role>().unwrap(), Role::Admin);
        assert_eq!("user".parse::<Role>().unwrap(), Role::User);
        assert_eq!("USER".parse::<Role>().unwrap(), Role::User);
        assert!("invalid".parse::<Role>().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"));
    }
 }
@@ -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<Url, JwksError> {
    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<Option<CachedJwks>>,
    /// 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<String>, ttl: Duration) -> Result<Self, JwksError> {
        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<Self, JwksError> {
        // 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<JwkSet, JwksError> {
        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<JwkSet, JwksError> {
        // 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<Jwk, JwksError> {
        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))));
    }
 }
@@ -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<String>,
    /// Issuer
    pub iss: Option<String>,
    /// Audience (can be string or array)
    #[serde(default)]
    pub aud: Audience,
    /// Expiration time
    pub exp: Option<u64>,
    /// Issued at
    pub iat: Option<u64>,
    /// Not before
    pub nbf: Option<u64>,
    /// JWT ID
    pub jti: Option<String>,
    /// Email (common claim)
    pub email: Option<String>,
    /// Name (common claim)
    pub name: Option<String>,
    /// Preferred username (OIDC claim)
    pub preferred_username: Option<String>,
    /// All other claims
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
 }
 /// 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<String>),
    #[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<String>,
    /// Display name if present
    pub name: Option<String>,
    /// 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<JwksProvider>,
    /// JWT configuration
    config: JwtConfig,
    /// Pre-configured validation settings
    validation: Validation,
    /// JTI cache for replay protection (optional)
    /// Maps JTI -> expiration time
    jti_cache: Option<Mutex<LruCache<String, JtiCacheEntry>>>,
    /// 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<JwksProvider>) -> 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<JwksProvider>,
        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, JwtValidatorError> {
        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<Self, JwtValidatorError> {
        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<ValidatedToken, JwtValidatorError> {
        // 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<StandardClaims> = 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<Role, JwtValidatorError> {
        // Try to get the role claim value
        let role_value = claims.extra.get(&self.config.role_claim);
        let role_strings: Vec<String> = 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::<Role>() {
                    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<DecodingKey, JwtValidatorError> {
        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<Algorithm, JwtValidatorError> {
        // 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<JwksProvider> {
        &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"));
    }
 }
@@ -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<String>,
    /// 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<B> PrincipalExt for http::Request<B> {
    fn principal(&self) -> Option<&Principal> {
        self.extensions().get::<Principal>()
    }
 }
 /// 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<Arc<JwtValidator>>,
    /// Audit logger
    pub audit_logger: AuditLogger,
 }
 impl ControlPlaneAuthState {
    /// Create a new control plane auth state.
    pub fn new(config: ControlPlaneAuthConfig, jwt_validator: Option<Arc<JwtValidator>>) -> 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<Self, Box<dyn std::error::Error + Send + Sync>> {
        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<Self> {
        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<Response> {
    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<ControlPlaneAuthState>,
    mut request: Request<Body>,
    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::<RequestId>().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());
    }
 }
@@ -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,
 };
@@ -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;