micromegas-auth 0.18.0

use crate::types::{AuthContext, AuthProvider, AuthType};
use anyhow::{Result, anyhow};
use base64::Engine;
use chrono::{DateTime, TimeDelta, Utc};
use jsonwebtoken::{Algorithm, Validation, decode, decode_header};
use moka::future::Cache;
use openidconnect::core::{CoreJsonWebKeySet, CoreProviderMetadata};
use openidconnect::{IssuerUrl, JsonWebKey};
use rsa::pkcs1::EncodeRsaPublicKey;
use rsa::{BigUint, RsaPublicKey};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

/// Create HTTP client for OIDC operations with security best practices
///
/// This client is configured with SSRF protection:
/// - No redirects allowed (prevents open redirect attacks)
/// - Suitable for OIDC discovery and token exchange operations
///
/// # Security
///
/// The no-redirect policy is important for OIDC operations because:
/// - Prevents attackers from redirecting requests to internal services
/// - Ensures OIDC endpoints are accessed directly without intermediate hops
/// - Protects against SSRF (Server-Side Request Forgery) attacks
///
/// # Example
///
/// ```rust
/// use micromegas_auth::oidc::create_http_client;
///
/// # async fn example() -> anyhow::Result<()> {
/// let client = create_http_client()?;
/// let response = client.get("https://accounts.google.com/.well-known/openid-configuration")
///     .send()
///     .await?;
/// # Ok(())
/// # }
/// ```
pub fn create_http_client() -> Result<reqwest::Client> {
    reqwest::ClientBuilder::new()
        .redirect(reqwest::redirect::Policy::none())
        .build()
        .map_err(|e| anyhow!("Failed to create HTTP client: {e:?}"))
}

/// Fetch JWKS from the OIDC provider using openidconnect's built-in discovery
async fn fetch_jwks(issuer_url: &IssuerUrl) -> Result<Arc<CoreJsonWebKeySet>> {
    let http_client = create_http_client()?;

    // Use openidconnect's built-in OIDC discovery
    let metadata = CoreProviderMetadata::discover_async(issuer_url.clone(), &http_client)
        .await
        .map_err(|e| {
            anyhow!(
                "Failed to discover OIDC metadata from {}: {e:?}",
                issuer_url
            )
        })?;

    // Fetch JWKS from jwks_uri
    let jwks_uri = metadata.jwks_uri();
    let jwks: CoreJsonWebKeySet = http_client
        .get(jwks_uri.url().as_str())
        .send()
        .await
        .map_err(|e| anyhow!("Failed to fetch JWKS from {}: {e:?}", jwks_uri))?
        .json()
        .await
        .map_err(|e| anyhow!("Failed to parse JWKS: {e:?}"))?;

    Ok(Arc::new(jwks))
}

/// JWKS cache for an OIDC issuer
///
/// Caches JSON Web Key Sets with automatic TTL expiration.
/// Uses moka for thread-safe caching with atomic cache miss handling.
struct JwksCache {
    issuer_url: IssuerUrl,
    cache: Cache<String, Arc<CoreJsonWebKeySet>>,
}

impl JwksCache {
    /// Create a new JWKS cache
    fn new(issuer_url: IssuerUrl, ttl: Duration) -> Self {
        let cache = Cache::builder().time_to_live(ttl).build();

        Self { issuer_url, cache }
    }

    /// Get the JWKS, fetching from the issuer if not cached
    async fn get(&self) -> Result<Arc<CoreJsonWebKeySet>> {
        let issuer_url = self.issuer_url.clone();

        self.cache
            .try_get_with(
                "jwks".to_string(),
                async move { fetch_jwks(&issuer_url).await },
            )
            .await
            .map_err(|e| anyhow!("Failed to fetch JWKS: {e:?}"))
    }
}

/// Configuration for a single OIDC issuer
#[derive(Debug, Clone, Deserialize)]
pub struct OidcIssuer {
    /// Issuer URL (e.g., <https://accounts.google.com>)
    pub issuer: String,
    /// Expected audience
    /// - For access tokens: API audience (e.g., "<https://api.example.com>")
    /// - For ID tokens: client ID
    pub audience: String,
}

const DEFAULT_JWKS_REFRESH_INTERVAL_SECS: u64 = 3600;
const DEFAULT_TOKEN_CACHE_SIZE: u64 = 1000;
const DEFAULT_TOKEN_CACHE_TTL_SECS: u64 = 300;

/// OIDC configuration
#[derive(Debug, Clone, Deserialize)]
#[serde(default)]
pub struct OidcConfig {
    /// List of configured OIDC issuers
    pub issuers: Vec<OidcIssuer>,
    /// JWKS refresh interval in seconds (default: 3600 = 1 hour)
    pub jwks_refresh_interval_secs: u64,
    /// Token cache size (default: 1000)
    pub token_cache_size: u64,
    /// Token cache TTL in seconds (default: 300 = 5 min)
    pub token_cache_ttl_secs: u64,
}

impl Default for OidcConfig {
    fn default() -> Self {
        Self {
            issuers: Vec::new(),
            jwks_refresh_interval_secs: DEFAULT_JWKS_REFRESH_INTERVAL_SECS,
            token_cache_size: DEFAULT_TOKEN_CACHE_SIZE,
            token_cache_ttl_secs: DEFAULT_TOKEN_CACHE_TTL_SECS,
        }
    }
}

impl OidcConfig {
    /// Load OIDC configuration from environment variable
    pub fn from_env() -> Result<Self> {
        let json = std::env::var("MICROMEGAS_OIDC_CONFIG")
            .map_err(|_| anyhow!("MICROMEGAS_OIDC_CONFIG environment variable not set"))?;
        let config: OidcConfig = serde_json::from_str(&json)
            .map_err(|e| anyhow!("Failed to parse MICROMEGAS_OIDC_CONFIG: {e:?}"))?;
        Ok(config)
    }
}

/// Audience can be either a string or an array of strings in OIDC tokens
#[derive(Debug, Serialize, Deserialize)]
#[serde(untagged)]
enum Audience {
    Single(String),
    Multiple(Vec<String>),
}

impl Audience {
    fn contains(&self, aud: &str) -> bool {
        match self {
            Audience::Single(s) => s == aud,
            Audience::Multiple(v) => v.iter().any(|a| a == aud),
        }
    }
}

/// JWT Claims from OIDC ID token or access token
///
/// This struct supports multiple OIDC providers by including various email claim fields.
/// Different providers use different claim names for email addresses:
///
/// Email extraction priority (see `get_email()` method):
/// 1. `verified_primary_email` - Azure AD optional claim (most reliable, verified by provider)
/// 2. `email` - Standard OIDC claim (Google, some Azure AD configurations)
/// 3. `namespaced_email` - Custom namespaced claims (Auth0: `https://micromegas.io/email`)
/// 4. `preferred_username` - Azure AD standard claim (often contains email)
/// 5. `upn` - User Principal Name (Azure AD enterprise)
/// 6. `unique_name` - Legacy Azure AD claim
#[derive(Debug, Serialize, Deserialize)]
struct Claims {
    /// Issuer - identifies the principal that issued the JWT
    iss: String,
    /// Subject - identifies the principal that is the subject of the JWT
    sub: String,
    /// Audience - identifies the recipients that the JWT is intended for
    /// Can be either a single string or an array of strings
    aud: Audience,
    /// Expiration time - identifies the expiration time on or after which the JWT must not be accepted
    exp: i64,
    /// Email address of the user (standard OIDC claim, used by Google and some Azure AD configurations)
    #[serde(skip_serializing_if = "Option::is_none")]
    email: Option<String>,
    /// Verified primary email (Azure AD optional claim - most reliable, verified by provider)
    #[serde(skip_serializing_if = "Option::is_none")]
    verified_primary_email: Option<String>,
    /// Preferred username (Azure AD standard claim, often contains email)
    #[serde(skip_serializing_if = "Option::is_none")]
    preferred_username: Option<String>,
    /// User Principal Name (Azure AD enterprise claim)
    #[serde(skip_serializing_if = "Option::is_none")]
    upn: Option<String>,
    /// Unique name (legacy Azure AD claim from older tokens)
    #[serde(skip_serializing_if = "Option::is_none")]
    unique_name: Option<String>,
    /// Custom namespaced email claim (Auth0 custom claims via Actions)
    #[serde(rename = "https://micromegas.io/email")]
    #[serde(skip_serializing_if = "Option::is_none")]
    namespaced_email: Option<String>,
    /// Custom namespaced name claim (Auth0 custom claims via Actions)
    #[serde(rename = "https://micromegas.io/name")]
    #[serde(skip_serializing_if = "Option::is_none")]
    namespaced_name: Option<String>,
}

impl Claims {
    /// Get email from various possible claim fields
    /// Priority order: verified_primary_email (most reliable) → email → namespaced_email → preferred_username → upn → unique_name
    fn get_email(&self) -> Option<String> {
        self.verified_primary_email
            .clone()
            .or_else(|| self.email.clone())
            .or_else(|| self.namespaced_email.clone())
            .or_else(|| self.preferred_username.clone())
            .or_else(|| self.upn.clone())
            .or_else(|| self.unique_name.clone())
    }
}

/// OIDC issuer client for token validation
struct OidcIssuerClient {
    issuer: String,
    audience: String,
    jwks_cache: JwksCache,
}

impl OidcIssuerClient {
    fn new(issuer: String, audience: String, jwks_ttl: Duration) -> Result<Self> {
        let issuer_url = IssuerUrl::new(issuer.clone())
            .map_err(|e| anyhow!("Invalid issuer URL '{}': {e:?}", issuer))?;

        Ok(Self {
            issuer,
            audience,
            jwks_cache: JwksCache::new(issuer_url, jwks_ttl),
        })
    }
}

/// Load admin users from environment variable
fn load_admin_users() -> Vec<String> {
    match std::env::var("MICROMEGAS_ADMINS") {
        Ok(json) => serde_json::from_str::<Vec<String>>(&json).unwrap_or_default(),
        Err(_) => vec![],
    }
}

/// Convert a JWK to a DecodingKey for jsonwebtoken
fn jwk_to_decoding_key(
    jwk: &openidconnect::core::CoreJsonWebKey,
) -> Result<jsonwebtoken::DecodingKey> {
    // Serialize the JWK to JSON to extract parameters
    let jwk_json =
        serde_json::to_value(jwk).map_err(|e| anyhow!("Failed to serialize JWK: {e:?}"))?;

    // Extract n and e parameters
    let n = jwk_json
        .get("n")
        .and_then(|v| v.as_str())
        .ok_or_else(|| anyhow!("JWK missing 'n' parameter"))?;
    let e = jwk_json
        .get("e")
        .and_then(|v| v.as_str())
        .ok_or_else(|| anyhow!("JWK missing 'e' parameter"))?;

    // Decode base64url encoded parameters
    let n_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(n.as_bytes())
        .map_err(|e| anyhow!("Failed to decode 'n': {e:?}"))?;
    let e_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(e.as_bytes())
        .map_err(|e| anyhow!("Failed to decode 'e': {e:?}"))?;

    // Create RSA public key
    let n_bigint = BigUint::from_bytes_be(&n_bytes);
    let e_bigint = BigUint::from_bytes_be(&e_bytes);

    let public_key = RsaPublicKey::new(n_bigint, e_bigint)
        .map_err(|e| anyhow!("Failed to create RSA public key: {e:?}"))?;

    // Convert to PEM format
    let pem = public_key
        .to_pkcs1_pem(rsa::pkcs1::LineEnding::LF)
        .map_err(|e| anyhow!("Failed to encode public key as PEM: {e:?}"))?;

    // Create DecodingKey
    jsonwebtoken::DecodingKey::from_rsa_pem(pem.as_bytes())
        .map_err(|e| anyhow!("Failed to create decoding key: {e:?}"))
}

/// OIDC authentication provider
///
/// Validates JWT tokens (access or ID tokens) from configured OIDC providers.
/// Caches both JWKS and validated tokens for performance.
pub struct OidcAuthProvider {
    /// Map from issuer URL to list of clients (supports multiple audiences per issuer)
    clients: HashMap<String, Vec<Arc<OidcIssuerClient>>>,
    /// Cache for validated tokens
    token_cache: Cache<String, Arc<AuthContext>>,
    /// Admin users (by email or subject)
    admin_users: Vec<String>,
}

impl std::fmt::Debug for OidcAuthProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OidcAuthProvider")
            .field("num_clients", &self.clients.len())
            .field("admin_users", &"(not printed)")
            .finish()
    }
}

impl OidcAuthProvider {
    /// Create a new OIDC authentication provider
    pub async fn new(config: OidcConfig) -> Result<Self> {
        if config.issuers.is_empty() {
            return Err(anyhow!("At least one OIDC issuer must be configured"));
        }

        micromegas_tracing::info!("Configuring OIDC with {} issuer(s)", config.issuers.len());
        for (idx, issuer_config) in config.issuers.iter().enumerate() {
            micromegas_tracing::info!(
                "  Issuer {}: {} (audience: {})",
                idx + 1,
                issuer_config.issuer,
                issuer_config.audience
            );
        }

        let jwks_ttl = Duration::from_secs(config.jwks_refresh_interval_secs);
        let mut clients: HashMap<String, Vec<Arc<OidcIssuerClient>>> = HashMap::new();

        // Initialize a client for each configured issuer
        // Supports multiple audiences for the same issuer (e.g., access tokens + ID tokens)
        for issuer_config in config.issuers {
            let client = OidcIssuerClient::new(
                issuer_config.issuer.clone(),
                issuer_config.audience,
                jwks_ttl,
            )?;

            clients
                .entry(issuer_config.issuer)
                .or_default()
                .push(Arc::new(client));
        }

        // Create token cache
        let token_cache = Cache::builder()
            .max_capacity(config.token_cache_size)
            .time_to_live(Duration::from_secs(config.token_cache_ttl_secs))
            .build();

        // Load admin users from environment
        let admin_users = load_admin_users();

        Ok(Self {
            clients,
            token_cache,
            admin_users,
        })
    }

    fn is_admin(&self, subject: &str, email: Option<&str>) -> bool {
        self.admin_users
            .iter()
            .any(|admin| admin == subject || email.map(|e| admin == e).unwrap_or(false))
    }

    /// Decode JWT payload without validation to extract issuer
    ///
    /// JWTs are structured as: header.payload.signature
    /// Both header and payload are base64url-encoded JSON
    fn decode_payload_unsafe(&self, token: &str) -> Result<Claims> {
        let parts: Vec<&str> = token.split('.').collect();
        if parts.len() != 3 {
            return Err(anyhow!("Invalid JWT format"));
        }

        // Decode the payload (second part)
        let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .decode(parts[1].as_bytes())
            .map_err(|e| anyhow!("Failed to decode JWT payload: {e:?}"))?;

        let claims: Claims = serde_json::from_slice(&payload_bytes)
            .map_err(|e| anyhow!("Failed to parse JWT claims: {e:?}"))?;

        Ok(claims)
    }

    /// Validate a JWT token (access token or ID token) and return authentication context
    ///
    /// This implementation follows OAuth 2.0 best practices by:
    /// 1. Extracting kid from JWT header for direct key lookup (falls back to trying all keys)
    /// 2. Extracting issuer from JWT payload for direct client lookup
    /// 3. Using O(1) lookups instead of O(n*m) iteration
    /// 4. Eliminating timing side-channels
    ///
    /// Supports both:
    /// - Access tokens (with API audience, used by Auth0 and similar)
    /// - ID tokens (with client_id audience, used by Google/Azure AD)
    async fn validate_jwt_token(&self, token: &str) -> Result<AuthContext> {
        // Step 1: Decode header (unsigned) to get key ID (kid)
        let header = decode_header(token).map_err(|e| anyhow!("Invalid JWT header: {e:?}"))?;

        // kid is optional - some providers omit it when there's only one key
        let kid = header.kid;

        // Step 2: Decode payload (unsigned) to get issuer and expiration
        let unverified_claims = self.decode_payload_unsafe(token)?;

        // Step 3: Look up clients for this issuer
        let issuer_clients = self
            .clients
            .get(&unverified_claims.iss)
            .ok_or_else(|| anyhow!("Unknown issuer: {}", unverified_claims.iss))?;

        // Step 4: Fetch JWKS once (all clients for the same issuer share the same JWKS)
        // Use the first client's cache - they all point to the same issuer's JWKS endpoint
        let first_client = issuer_clients
            .first()
            .ok_or_else(|| anyhow!("No clients configured for issuer"))?;

        let jwks = first_client
            .jwks_cache
            .get()
            .await
            .map_err(|e| anyhow!("Failed to fetch JWKS: {e:?}"))?;

        // Step 5: Find the key(s) to try
        // If kid is present, look up directly; otherwise try all keys
        let keys_to_try: Vec<_> = if let Some(ref kid_value) = kid {
            // Direct lookup by kid
            jwks.keys()
                .iter()
                .filter(|k| k.key_id().map(|id| id.as_str()) == Some(kid_value.as_str()))
                .collect()
        } else {
            // No kid provided - try all keys (common for single-key providers)
            jwks.keys().iter().collect()
        };

        if keys_to_try.is_empty() {
            return Err(if let Some(kid_value) = kid {
                anyhow!("Key with kid '{}' not found in JWKS", kid_value)
            } else {
                anyhow!("No keys found in JWKS")
            });
        }

        // Step 6: Try each key until one works
        let mut key_error = anyhow!("No valid key found");
        for key in keys_to_try {
            match jwk_to_decoding_key(key) {
                Ok(decoding_key) => {
                    match self
                        .try_validate_with_key(token, &decoding_key, issuer_clients)
                        .await
                    {
                        Ok(auth_ctx) => return Ok(auth_ctx),
                        Err(e) => key_error = e,
                    }
                }
                Err(e) => key_error = e,
            }
        }

        Err(key_error)
    }

    /// Try to validate a token with a specific decoding key against all configured audiences
    async fn try_validate_with_key(
        &self,
        token: &str,
        decoding_key: &jsonwebtoken::DecodingKey,
        issuer_clients: &[Arc<OidcIssuerClient>],
    ) -> Result<AuthContext> {
        let configured_audiences: Vec<String> =
            issuer_clients.iter().map(|c| c.audience.clone()).collect();
        let mut last_error = anyhow!("No matching audience found");

        for client in issuer_clients {
            // Validate token with specific key and issuer
            let mut validation = Validation::new(Algorithm::RS256);
            // Don't validate audience yet - we'll do it manually
            validation.validate_aud = false;
            validation.set_issuer(&[&client.issuer]);

            let claims = match decode::<Claims>(token, decoding_key, &validation) {
                Ok(token_data) => token_data.claims,
                Err(e) => {
                    last_error = anyhow!("Token validation failed: {e:?}");
                    continue;
                }
            };

            // Manually validate audience - if it matches, we found the right client!
            if claims.aud.contains(&client.audience) {
                // Success! Validate expiration and return auth context
                let expires_at = DateTime::from_timestamp(claims.exp, 0)
                    .ok_or_else(|| anyhow!("Invalid expiration timestamp"))?;

                if expires_at < Utc::now() {
                    return Err(anyhow!("Token has expired"));
                }

                let email = claims.get_email();
                let is_admin = self.is_admin(&claims.sub, email.as_deref());

                return Ok(AuthContext {
                    subject: claims.sub,
                    email,
                    issuer: claims.iss,
                    audience: Some(client.audience.clone()),
                    expires_at: Some(expires_at),
                    auth_type: AuthType::Oidc,
                    is_admin,
                    allow_delegation: false,
                });
            } else {
                // Provide detailed error with configured vs actual audiences
                let actual_audiences = match &claims.aud {
                    Audience::Single(s) => vec![s.clone()],
                    Audience::Multiple(v) => v.clone(),
                };
                last_error = anyhow!(
                    "Token audience mismatch - configured audiences: {:?}, token audiences: {:?}",
                    configured_audiences,
                    actual_audiences
                );
            }
        }

        // If we get here, none of the clients matched
        Err(last_error)
    }
}

#[async_trait::async_trait]
impl AuthProvider for OidcAuthProvider {
    async fn validate_request(
        &self,
        parts: &dyn crate::types::RequestParts,
    ) -> Result<AuthContext> {
        let token = parts
            .bearer_token()
            .ok_or_else(|| anyhow!("missing bearer token"))?;

        // Check token cache first, but verify token hasn't expired
        // Use 30-second grace period to account for clock skew and network latency
        if let Some(cached) = self.token_cache.get(token).await {
            let is_expired = cached
                .expires_at
                .map(|exp| exp <= Utc::now() + TimeDelta::seconds(30))
                .unwrap_or(false);

            if is_expired {
                self.token_cache.remove(token).await;
            } else {
                return Ok((*cached).clone());
            }
        }

        // Validate token
        let auth_ctx = self.validate_jwt_token(token).await?;

        // Cache the result
        self.token_cache
            .insert(token.to_string(), Arc::new(auth_ctx.clone()))
            .await;

        Ok(auth_ctx)
    }
}