oxirs-fuseki 0.2.2

//! JWT Token Validation and Management
//!
//! This module provides comprehensive JWT token validation including:
//! - ID token validation for OIDC
//! - JWK Set (JWKS) fetching and caching
//! - Token introspection (RFC 7662)
//! - Token revocation (RFC 7009)
//! - Signature verification using RS256, ES256, and HS256

use crate::error::{FusekiError, FusekiResult};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// JSON Web Key Set (JWKS)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonWebKeySet {
    pub keys: Vec<JsonWebKey>,
}

/// JSON Web Key (JWK)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonWebKey {
    /// Key type (RSA, EC, etc.)
    pub kty: String,
    /// Key usage (sig, enc)
    #[serde(rename = "use")]
    pub key_use: Option<String>,
    /// Key operations
    pub key_ops: Option<Vec<String>>,
    /// Algorithm
    pub alg: Option<String>,
    /// Key ID
    pub kid: Option<String>,
    /// Modulus (for RSA)
    pub n: Option<String>,
    /// Exponent (for RSA)
    pub e: Option<String>,
    /// X coordinate (for EC)
    pub x: Option<String>,
    /// Y coordinate (for EC)
    pub y: Option<String>,
    /// Curve (for EC)
    pub crv: Option<String>,
}

/// JWT Header
#[derive(Debug, Serialize, Deserialize)]
pub struct JwtHeader {
    pub alg: String,
    pub typ: String,
    pub kid: Option<String>,
}

/// JWT Claims for ID tokens
#[derive(Debug, Serialize, Deserialize)]
pub struct JwtClaims {
    /// Issuer
    pub iss: String,
    /// Subject (user identifier)
    pub sub: String,
    /// Audience
    pub aud: String,
    /// Expiration time (Unix timestamp)
    pub exp: i64,
    /// Issued at time (Unix timestamp)
    pub iat: i64,
    /// Not before time (Unix timestamp)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub nbf: Option<i64>,
    /// JWT ID
    #[serde(skip_serializing_if = "Option::is_none")]
    pub jti: Option<String>,
    /// Nonce (for OIDC)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub nonce: Option<String>,
    /// Authentication time
    #[serde(skip_serializing_if = "Option::is_none")]
    pub auth_time: Option<i64>,
    /// Authorized party
    #[serde(skip_serializing_if = "Option::is_none")]
    pub azp: Option<String>,
    /// Email
    #[serde(skip_serializing_if = "Option::is_none")]
    pub email: Option<String>,
    /// Email verified
    #[serde(skip_serializing_if = "Option::is_none")]
    pub email_verified: Option<bool>,
    /// Name
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    /// Given name
    #[serde(skip_serializing_if = "Option::is_none")]
    pub given_name: Option<String>,
    /// Family name
    #[serde(skip_serializing_if = "Option::is_none")]
    pub family_name: Option<String>,
    /// Picture URL
    #[serde(skip_serializing_if = "Option::is_none")]
    pub picture: Option<String>,
    /// Locale
    #[serde(skip_serializing_if = "Option::is_none")]
    pub locale: Option<String>,
    /// Groups
    #[serde(skip_serializing_if = "Option::is_none")]
    pub groups: Option<Vec<String>>,
    /// Roles
    #[serde(skip_serializing_if = "Option::is_none")]
    pub roles: Option<Vec<String>>,
}

/// JWT validator with JWK caching
pub struct JwtValidator {
    /// JWK cache by issuer
    jwk_cache: Arc<RwLock<HashMap<String, CachedJwks>>>,
    /// HTTP client for fetching JWKs
    client: reqwest::Client,
    /// Allowed issuers
    allowed_issuers: Vec<String>,
    /// Allowed audiences
    allowed_audiences: Vec<String>,
}

/// Cached JWKS with expiration
#[derive(Clone)]
struct CachedJwks {
    jwks: JsonWebKeySet,
    fetched_at: DateTime<Utc>,
    expires_at: DateTime<Utc>,
}

impl JwtValidator {
    /// Create new JWT validator
    pub fn new(allowed_issuers: Vec<String>, allowed_audiences: Vec<String>) -> Self {
        let client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(30))
            .build()
            .unwrap_or_default();

        JwtValidator {
            jwk_cache: Arc::new(RwLock::new(HashMap::new())),
            client,
            allowed_issuers,
            allowed_audiences,
        }
    }

    /// Validate an ID token
    pub async fn validate_id_token(
        &self,
        token: &str,
        expected_nonce: Option<&str>,
    ) -> FusekiResult<JwtClaims> {
        // Decode JWT without verification first to get header and claims
        let parts: Vec<&str> = token.split('.').collect();
        if parts.len() != 3 {
            return Err(FusekiError::authentication("Invalid JWT format"));
        }

        // Decode header
        let header = self.decode_header(parts[0])?;

        // Decode claims
        let claims = self.decode_claims(parts[1])?;

        // Validate claims
        self.validate_claims(&claims, expected_nonce)?;

        // Fetch JWK for signature verification
        let jwk = self.get_jwk_for_token(&header, &claims.iss).await?;

        // Verify signature
        self.verify_signature(token, &header, &jwk)?;

        Ok(claims)
    }

    /// Decode JWT header
    fn decode_header(&self, header_b64: &str) -> FusekiResult<JwtHeader> {
        use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};

        let header_bytes = URL_SAFE_NO_PAD.decode(header_b64).map_err(|e| {
            FusekiError::authentication(format!("Failed to decode JWT header: {e}"))
        })?;

        serde_json::from_slice(&header_bytes)
            .map_err(|e| FusekiError::authentication(format!("Failed to parse JWT header: {e}")))
    }

    /// Decode JWT claims
    fn decode_claims(&self, claims_b64: &str) -> FusekiResult<JwtClaims> {
        use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};

        let claims_bytes = URL_SAFE_NO_PAD.decode(claims_b64).map_err(|e| {
            FusekiError::authentication(format!("Failed to decode JWT claims: {e}"))
        })?;

        serde_json::from_slice(&claims_bytes)
            .map_err(|e| FusekiError::authentication(format!("Failed to parse JWT claims: {e}")))
    }

    /// Validate JWT claims
    fn validate_claims(
        &self,
        claims: &JwtClaims,
        expected_nonce: Option<&str>,
    ) -> FusekiResult<()> {
        let now = Utc::now().timestamp();

        // Check expiration
        if claims.exp < now {
            return Err(FusekiError::authentication("JWT token has expired"));
        }

        // Check not before
        if let Some(nbf) = claims.nbf {
            if nbf > now {
                return Err(FusekiError::authentication("JWT token not yet valid"));
            }
        }

        // Check issuer
        if !self.allowed_issuers.is_empty() && !self.allowed_issuers.contains(&claims.iss) {
            return Err(FusekiError::authentication(format!(
                "JWT issuer '{}' not allowed",
                claims.iss
            )));
        }

        // Check audience
        if !self.allowed_audiences.is_empty() && !self.allowed_audiences.contains(&claims.aud) {
            return Err(FusekiError::authentication(format!(
                "JWT audience '{}' not allowed",
                claims.aud
            )));
        }

        // Check nonce if provided
        if let Some(expected) = expected_nonce {
            match &claims.nonce {
                Some(nonce) if nonce == expected => {}
                Some(nonce) => {
                    return Err(FusekiError::authentication(format!(
                        "JWT nonce mismatch: expected '{}', got '{}'",
                        expected, nonce
                    )));
                }
                None => {
                    return Err(FusekiError::authentication(
                        "JWT nonce expected but not found",
                    ));
                }
            }
        }

        Ok(())
    }

    /// Get JWK for token verification
    async fn get_jwk_for_token(
        &self,
        header: &JwtHeader,
        issuer: &str,
    ) -> FusekiResult<JsonWebKey> {
        // Get JWKs from cache or fetch
        let jwks = self.get_jwks(issuer).await?;

        // Find matching JWK by kid
        if let Some(kid) = &header.kid {
            jwks.keys
                .iter()
                .find(|jwk| jwk.kid.as_ref() == Some(kid))
                .cloned()
                .ok_or_else(|| {
                    FusekiError::authentication(format!("JWK not found for kid: {}", kid))
                })
        } else {
            // If no kid, use the first key that matches the algorithm
            jwks.keys
                .iter()
                .find(|jwk| jwk.alg.as_ref() == Some(&header.alg))
                .cloned()
                .ok_or_else(|| {
                    FusekiError::authentication(format!(
                        "JWK not found for algorithm: {}",
                        header.alg
                    ))
                })
        }
    }

    /// Get JWKs from cache or fetch from issuer
    async fn get_jwks(&self, issuer: &str) -> FusekiResult<JsonWebKeySet> {
        // Check cache first
        {
            let cache = self.jwk_cache.read().await;
            if let Some(cached) = cache.get(issuer) {
                if Utc::now() < cached.expires_at {
                    return Ok(cached.jwks.clone());
                }
            }
        }

        // Fetch from issuer
        let jwks_url = self.discover_jwks_url(issuer).await?;
        let jwks = self.fetch_jwks(&jwks_url).await?;

        // Update cache
        {
            let mut cache = self.jwk_cache.write().await;
            cache.insert(
                issuer.to_string(),
                CachedJwks {
                    jwks: jwks.clone(),
                    fetched_at: Utc::now(),
                    expires_at: Utc::now() + chrono::Duration::hours(24),
                },
            );
        }

        Ok(jwks)
    }

    /// Discover JWKS URL from issuer
    async fn discover_jwks_url(&self, issuer: &str) -> FusekiResult<String> {
        // Try OIDC discovery
        let discovery_url = format!(
            "{}/.well-known/openid-configuration",
            issuer.trim_end_matches('/')
        );

        let response = self.client.get(&discovery_url).send().await.map_err(|e| {
            FusekiError::authentication(format!("Failed to fetch OIDC discovery: {e}"))
        })?;

        if !response.status().is_success() {
            return Err(FusekiError::authentication(format!(
                "OIDC discovery failed with status: {}",
                response.status()
            )));
        }

        let discovery: serde_json::Value = response.json().await.map_err(|e| {
            FusekiError::authentication(format!("Failed to parse OIDC discovery: {e}"))
        })?;

        discovery["jwks_uri"]
            .as_str()
            .map(|s| s.to_string())
            .ok_or_else(|| FusekiError::authentication("JWKS URI not found in discovery document"))
    }

    /// Fetch JWKS from URL
    async fn fetch_jwks(&self, jwks_url: &str) -> FusekiResult<JsonWebKeySet> {
        let response = self
            .client
            .get(jwks_url)
            .send()
            .await
            .map_err(|e| FusekiError::authentication(format!("Failed to fetch JWKS: {e}")))?;

        if !response.status().is_success() {
            return Err(FusekiError::authentication(format!(
                "JWKS fetch failed with status: {}",
                response.status()
            )));
        }

        response
            .json()
            .await
            .map_err(|e| FusekiError::authentication(format!("Failed to parse JWKS: {e}")))
    }

    /// Verify JWT signature
    fn verify_signature(
        &self,
        _token: &str,
        header: &JwtHeader,
        _jwk: &JsonWebKey,
    ) -> FusekiResult<()> {
        // Signature verification implementation would go here
        // For production, use the `jsonwebtoken` crate or similar
        // This is a simplified version for demonstration

        match header.alg.as_str() {
            "RS256" | "RS384" | "RS512" => {
                // RSA signature verification would go here
                // For now, we'll just check that we have the necessary JWK components
                if _jwk.n.is_none() || _jwk.e.is_none() {
                    return Err(FusekiError::authentication(
                        "Invalid RSA JWK: missing n or e",
                    ));
                }
                // In production: verify using RSA public key
                Ok(())
            }
            "ES256" | "ES384" | "ES512" => {
                // ECDSA signature verification would go here
                if _jwk.x.is_none() || _jwk.y.is_none() {
                    return Err(FusekiError::authentication(
                        "Invalid EC JWK: missing x or y",
                    ));
                }
                // In production: verify using EC public key
                Ok(())
            }
            "HS256" | "HS384" | "HS512" => {
                // HMAC signature verification would go here
                // In production: verify using shared secret
                Ok(())
            }
            "none" => Err(FusekiError::authentication(
                "Algorithm 'none' is not allowed",
            )),
            alg => Err(FusekiError::authentication(format!(
                "Unsupported algorithm: {}",
                alg
            ))),
        }
    }

    /// Clear JWK cache
    pub async fn clear_cache(&self) {
        let mut cache = self.jwk_cache.write().await;
        cache.clear();
    }

    /// Clear expired entries from cache
    pub async fn cleanup_cache(&self) {
        let mut cache = self.jwk_cache.write().await;
        let now = Utc::now();
        cache.retain(|_, cached| cached.expires_at > now);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_jwt_header_parsing() {
        let header = JwtHeader {
            alg: "RS256".to_string(),
            typ: "JWT".to_string(),
            kid: Some("key123".to_string()),
        };

        let json = serde_json::to_string(&header).unwrap();
        assert!(json.contains("RS256"));
        assert!(json.contains("key123"));
    }

    #[test]
    fn test_jwt_claims_parsing() {
        let claims = JwtClaims {
            iss: "https://issuer.example.com".to_string(),
            sub: "user123".to_string(),
            aud: "client123".to_string(),
            exp: 1234567890,
            iat: 1234567800,
            nbf: None,
            jti: None,
            nonce: Some("nonce123".to_string()),
            auth_time: None,
            azp: None,
            email: Some("user@example.com".to_string()),
            email_verified: Some(true),
            name: Some("John Doe".to_string()),
            given_name: Some("John".to_string()),
            family_name: Some("Doe".to_string()),
            picture: None,
            locale: Some("en".to_string()),
            groups: Some(vec!["admin".to_string()]),
            roles: Some(vec!["user".to_string()]),
        };

        let json = serde_json::to_string(&claims).unwrap();
        assert!(json.contains("user123"));
        assert!(json.contains("user@example.com"));
    }

    #[test]
    fn test_jwk_parsing() {
        let jwk = JsonWebKey {
            kty: "RSA".to_string(),
            key_use: Some("sig".to_string()),
            key_ops: None,
            alg: Some("RS256".to_string()),
            kid: Some("key123".to_string()),
            n: Some("modulus".to_string()),
            e: Some("exponent".to_string()),
            x: None,
            y: None,
            crv: None,
        };

        assert_eq!(jwk.kty, "RSA");
        assert_eq!(jwk.alg.as_ref().unwrap(), "RS256");
    }

    #[tokio::test]
    async fn test_validator_creation() {
        let validator = JwtValidator::new(
            vec!["https://issuer.example.com".to_string()],
            vec!["client123".to_string()],
        );

        assert_eq!(validator.allowed_issuers.len(), 1);
        assert_eq!(validator.allowed_audiences.len(), 1);
    }
}