arcly_http_core/auth/

jwt.rs

//! JWT authentication service — sign, decode, and validate JSON Web Tokens.
//!
//! ## Usage
//!
//! Provide a `JwtService` instance in an `ArclyPlugin::on_init`:
//!
//! ```ignore
//! ctx.provide(JwtService::new(JwtConfig {
//!     secret: "change-in-prod".to_string(),
//!     access_ttl_secs:  900,
//!     refresh_ttl_secs: 604_800,
//!     ..Default::default()
//! }));
//! ```
//!
//! Once provided, the HTTP and WebSocket boundaries automatically decode the
//! `Authorization: Bearer <token>` header and populate `RequestContext::claims()`
//! on every request — no per-handler boilerplate needed. Protect routes with
//! `JWT_AUTH.check(&ctx)?` or `RoleGuard::require("admin").check(&ctx)?`.

use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use serde::{Deserialize, Serialize};

use crate::web::context::Claims;

// ─── Configuration ────────────────────────────────────────────────────────────

/// Configuration for `JwtService`. Build once at startup and provide via DI.
/// Token signing failed — malformed key material (typically a bad rotation
/// payload). Map to a 500 at the route; the process must keep serving.
#[derive(Debug)]
pub struct JwtSignError(pub jsonwebtoken::errors::Error);

impl std::fmt::Display for JwtSignError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "JWT signing failed: {}", self.0)
    }
}
impl std::error::Error for JwtSignError {}

pub struct JwtConfig {
    /// For HMAC families (HS256/384/512): the raw shared secret.
    /// For RSA (RS*/PS*) and ECDSA (ES*) families: the **PKCS#8 / SEC1 PEM
    /// private key** used for signing.
    pub secret: String,
    /// PEM-encoded **public key** — required for RSA/ECDSA verification and for
    /// publishing JWKS (`JwtService::jwks_json`). Ignored for HMAC families,
    /// where the shared secret both signs and verifies.
    pub public_key_pem: Option<String>,
    /// The `kid` (key ID) advertised in the JWT header and JWKS document.
    /// Required for OIDC so relying parties can select the right JWK.
    pub key_id: Option<String>,
    /// Signing algorithm. Defaults to `HS256`.
    pub algorithm: Algorithm,
    /// Lifetime of access tokens in seconds. Defaults to 900 (15 min).
    pub access_ttl_secs: u64,
    /// Lifetime of refresh tokens in seconds. Defaults to 604 800 (7 days).
    pub refresh_ttl_secs: u64,
}

impl Default for JwtConfig {
    fn default() -> Self {
        Self {
            secret: "change-me-in-production".to_string(),
            public_key_pem: None,
            key_id: None,
            algorithm: Algorithm::HS256,
            access_ttl_secs: 900,
            refresh_ttl_secs: 604_800,
        }
    }
}

/// True for RSA / RSA-PSS / ECDSA families — anything that signs with a private
/// key and verifies with a distinct public key.
fn is_asymmetric(alg: Algorithm) -> bool {
    matches!(
        alg,
        Algorithm::RS256
            | Algorithm::RS384
            | Algorithm::RS512
            | Algorithm::PS256
            | Algorithm::PS384
            | Algorithm::PS512
            | Algorithm::ES256
            | Algorithm::ES384
            | Algorithm::EdDSA
    )
}

// ─── Internal claims struct ───────────────────────────────────────────────────

/// Private claims struct used for `encode` / `decode`.
/// Decoded into a `serde_json::Map` before being stored on `RequestContext`.
#[derive(Debug, Serialize, Deserialize)]
struct JwtClaims {
    sub: String,
    /// Omitted from refresh tokens (always empty there — no point encoding them).
    #[serde(skip_serializing_if = "String::is_empty", default)]
    role: String,
    #[serde(skip_serializing_if = "String::is_empty", default)]
    email: String,
    /// "access" or "refresh"
    #[serde(rename = "type")]
    kind: String,
    /// JWT ID — unique token identifier, used for refresh token rotation.
    jti: String,
    iat: u64,
    exp: u64,
    /// Fine-grained permissions (e.g. `["users:*", "orders:read"]`).
    /// Omitted from refresh tokens and when no permissions are set.
    #[serde(skip_serializing_if = "Vec::is_empty", default)]
    perms: Vec<String>,
    /// Home tenant of the principal. `TenantGuard` cross-checks this against
    /// the request's resolved tenant, which (a) blocks forged tenant headers
    /// and (b) makes dropping the header useless: a token bound to tenant A
    /// can never act as the fallback tenant.
    #[serde(skip_serializing_if = "String::is_empty", default)]
    tenant: String,
    /// RFC 9449 confirmation claim for sender-constrained (DPoP) tokens:
    /// `{"jkt": "<JWK SHA-256 thumbprint>"}`. Present only for bound tokens; a
    /// resource server rejects the token unless the request carries a DPoP proof
    /// whose key hashes to this `jkt`.
    #[serde(skip_serializing_if = "Option::is_none", default)]
    cnf: Option<Cnf>,
}

/// The `cnf` (confirmation) claim body — currently only the JWK thumbprint.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Cnf {
    pub jkt: String,
}

// ─── JwtService ───────────────────────────────────────────────────────────────

/// Live signing/verification keys. Swapped atomically as one bundle on
/// rotation so readers never observe a half-rotated state. `verify` keeps the
/// previous key so tokens signed before rotation stay valid through their TTL.
struct JwtKeyMaterial {
    encoding: EncodingKey,
    verify: Vec<DecodingKey>, // [current, previous?]
    version: u64,
}

impl JwtKeyMaterial {
    /// HMAC construction — signing and verification share one secret.
    fn from_secret(secret: &[u8], version: u64, previous: Option<DecodingKey>) -> Self {
        let mut verify = vec![DecodingKey::from_secret(secret)];
        verify.extend(previous);
        Self {
            encoding: EncodingKey::from_secret(secret),
            verify,
            version,
        }
    }

    /// Algorithm-aware construction. For HMAC families `signing` is the shared
    /// secret and `public_pem` is ignored; for RSA/ECDSA `signing` is the
    /// private-key PEM and `public_pem` (required) is the verification key.
    fn build(
        algorithm: Algorithm,
        signing: &[u8],
        public_pem: Option<&str>,
        version: u64,
        previous: Option<DecodingKey>,
    ) -> Result<Self, JwtSignError> {
        if !is_asymmetric(algorithm) {
            return Ok(Self::from_secret(signing, version, previous));
        }
        let pub_pem = public_pem
            .ok_or_else(|| JwtSignError(jwt_err()))?
            .as_bytes();
        let (encoding, decoding) = match algorithm {
            Algorithm::ES256 | Algorithm::ES384 => (
                EncodingKey::from_ec_pem(signing).map_err(JwtSignError)?,
                DecodingKey::from_ec_pem(pub_pem).map_err(JwtSignError)?,
            ),
            Algorithm::EdDSA => (
                EncodingKey::from_ed_pem(signing).map_err(JwtSignError)?,
                DecodingKey::from_ed_pem(pub_pem).map_err(JwtSignError)?,
            ),
            // RS* / PS*
            _ => (
                EncodingKey::from_rsa_pem(signing).map_err(JwtSignError)?,
                DecodingKey::from_rsa_pem(pub_pem).map_err(JwtSignError)?,
            ),
        };
        let mut verify = vec![decoding];
        verify.extend(previous);
        Ok(Self {
            encoding,
            verify,
            version,
        })
    }
}

fn jwt_err() -> jsonwebtoken::errors::Error {
    // Stable, matchable error kind: an asymmetric algorithm was configured
    // without the required public key.
    jsonwebtoken::errors::ErrorKind::InvalidKeyFormat.into()
}

/// Signs and validates JWTs. Provide this into the DI container so the framework
/// boundaries (`boundary.rs`, `ws.rs`) can auto-populate `RequestContext::claims()`
/// on every incoming request.
///
/// ## Secret rotation
///
/// Keys live behind [`Rotating`](crate::auth::secrets::Rotating) (an
/// `ArcSwap`): the request path pays one atomic pointer load, while
/// [`rotate_secret`](Self::rotate_secret) — typically driven by a
/// `SecretSource` watcher — swaps in a new bundle with no restart. The
/// previous key is retained for verification, so live tokens (≤ TTL old)
/// keep validating through the grace window.
///
/// Token *signing* returns `Result<_, JwtSignError>` — a malformed key from
/// a bad rotation payload must surface as a 500 on the affected request,
/// never as a process panic.
pub struct JwtService {
    keys: crate::auth::secrets::Rotating<JwtKeyMaterial>,
    header: Header,
    validation: Validation,
    config: JwtConfig,
}

impl JwtService {
    /// Build from config. **Panics** if an asymmetric algorithm is configured
    /// with malformed / missing PEM key material — this is a boot-time
    /// misconfiguration, not a per-request condition. Use
    /// [`try_new`](Self::try_new) to handle it gracefully.
    pub fn new(config: JwtConfig) -> Self {
        Self::try_new(config).expect("JwtService: invalid key material")
    }

    /// Fallible constructor — returns `Err` instead of panicking on bad key
    /// material (e.g. an ES256 config with an unparseable private-key PEM).
    pub fn try_new(config: JwtConfig) -> Result<Self, JwtSignError> {
        let material = JwtKeyMaterial::build(
            config.algorithm,
            config.secret.as_bytes(),
            config.public_key_pem.as_deref(),
            1,
            None,
        )?;
        let keys = crate::auth::secrets::Rotating::new(material);
        let mut header = Header::new(config.algorithm);
        header.kid = config.key_id.clone();
        let mut validation = Validation::new(config.algorithm);
        validation.validate_exp = true;
        Ok(Self {
            keys,
            header,
            validation,
            config,
        })
    }

    /// Hot-swap the signing secret — no restart, no token mass-invalidation.
    ///
    /// New tokens sign with the new key immediately; tokens signed with the
    /// previous key keep verifying until natural expiry. Versions are
    /// monotonic: a stale (≤ current) version is ignored, making concurrent
    /// watchers and duplicate delivery harmless.
    ///
    /// For HMAC families `new_secret` is the new shared secret. For asymmetric
    /// families it is the new **private-key PEM**; the paired public key must be
    /// supplied via [`rotate_keypair`](Self::rotate_keypair) instead — this
    /// method assumes the configured public key still verifies.
    pub fn rotate_secret(&self, new_secret: &[u8], version: u64) {
        self.rotate_keypair(new_secret, self.config.public_key_pem.as_deref(), version);
    }

    /// Rotate an asymmetric key **pair** (private PEM + public PEM). Retains the
    /// previous public key for the verification grace window, exactly like the
    /// HMAC path. No-op on stale versions.
    pub fn rotate_keypair(&self, new_signing: &[u8], new_public_pem: Option<&str>, version: u64) {
        let current = self.keys.load();
        if version <= current.version {
            tracing::warn!(
                current = current.version,
                offered = version,
                "ignoring stale JWT secret rotation",
            );
            return;
        }
        let previous = current.verify.first().cloned();
        match JwtKeyMaterial::build(
            self.config.algorithm,
            new_signing,
            new_public_pem,
            version,
            previous,
        ) {
            Ok(material) => {
                self.keys.store(material);
                tracing::info!(version, "JwtService signing key rotated");
            }
            Err(e) => {
                tracing::error!(version, error = %e, "JWT key rotation rejected — keeping current key");
            }
        }
    }

    /// Publish the current public verification key as a JWKS document (JSON).
    ///
    /// Returns `None` for HMAC families (a shared secret must never be exposed)
    /// or when no `public_key_pem` was configured. The OIDC provider serves this
    /// at `/.well-known/jwks.json`. The JWK is derived from the configured
    /// public PEM by the identity crate's OIDC layer; core stores the raw PEM
    /// and `kid` so that layer can convert without re-parsing config.
    pub fn public_key_pem(&self) -> Option<&str> {
        if is_asymmetric(self.config.algorithm) {
            self.config.public_key_pem.as_deref()
        } else {
            None
        }
    }

    /// The advertised key id (`kid`), if configured.
    pub fn key_id(&self) -> Option<&str> {
        self.config.key_id.as_deref()
    }

    /// The signing algorithm name as it appears in a JWK / JOSE header
    /// (`"HS256"`, `"RS256"`, `"ES256"`, …).
    pub fn algorithm_name(&self) -> &'static str {
        match self.config.algorithm {
            Algorithm::HS256 => "HS256",
            Algorithm::HS384 => "HS384",
            Algorithm::HS512 => "HS512",
            Algorithm::RS256 => "RS256",
            Algorithm::RS384 => "RS384",
            Algorithm::RS512 => "RS512",
            Algorithm::PS256 => "PS256",
            Algorithm::PS384 => "PS384",
            Algorithm::PS512 => "PS512",
            Algorithm::ES256 => "ES256",
            Algorithm::ES384 => "ES384",
            Algorithm::EdDSA => "EdDSA",
        }
    }

    fn now() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0)
    }

    /// Sign an **arbitrary claim set** with the current key and configured
    /// algorithm (header carries the `kid`). Use for tokens the typed helpers
    /// don't cover — notably OIDC `id_token`s, which need `aud`/`iss`/`nonce`.
    ///
    /// The caller owns the full claim set (including `iat`/`exp`); nothing is
    /// injected. Verify later with [`decode`](Self::decode).
    pub fn sign(&self, claims: &serde_json::Value) -> Result<String, JwtSignError> {
        encode(&self.header, claims, &self.keys.load().encoding).map_err(JwtSignError)
    }

    /// Current unix time in seconds — handy for callers building claim sets to
    /// pass to [`sign`](Self::sign).
    pub fn unix_now() -> u64 {
        Self::now()
    }

    /// Issue a signed **access token**.
    ///
    /// Claims: `sub` = user ID, `role`, `email`, `type = "access"`, `jti`, `iat`, `exp`.
    ///
    /// Signing can only fail on malformed key material (e.g. a bad rotation
    /// payload) — propagate the error instead of panicking mid-traffic.
    pub fn issue_access(&self, sub: &str, role: &str, email: &str) -> Result<String, JwtSignError> {
        self.issue_access_with_perms(sub, role, email, &[])
    }

    /// Like [`Self::issue_access`] but embeds a `perms` claim (array of permission strings).
    ///
    /// Use this when the app maintains a permission map so that `PermissionGuard`
    /// can do a zero-latency lookup without hitting the store on each request.
    pub fn issue_access_with_perms(
        &self,
        sub: &str,
        role: &str,
        email: &str,
        perms: &[String],
    ) -> Result<String, JwtSignError> {
        self.issue_access_bound(sub, role, email, perms, None)
    }

    /// Like [`Self::issue_access_with_perms`] but additionally **binds the token to
    /// a tenant** via the `tenant` claim. `TenantGuard` then enforces that
    /// requests carrying this token resolve to the same tenant — omitting or
    /// forging the tenant header yields `403`, so a suspended tenant's users
    /// cannot ride the fallback pool by dropping the header.
    pub fn issue_access_bound(
        &self,
        sub: &str,
        role: &str,
        email: &str,
        perms: &[String],
        tenant: Option<&str>,
    ) -> Result<String, JwtSignError> {
        self.issue_access_bound_cnf(sub, role, email, perms, tenant, None)
    }

    /// Like [`Self::issue_access_bound`] but additionally **sender-constrains**
    /// the token (RFC 9449 DPoP): `cnf_jkt` is the SHA-256 JWK thumbprint of the
    /// client's proof-of-possession key. A resource server then accepts the
    /// token only when the request carries a matching DPoP proof, so a stolen
    /// bearer token is useless without the client's private key.
    pub fn issue_access_bound_cnf(
        &self,
        sub: &str,
        role: &str,
        email: &str,
        perms: &[String],
        tenant: Option<&str>,
        cnf_jkt: Option<&str>,
    ) -> Result<String, JwtSignError> {
        let now = Self::now();
        let claims = JwtClaims {
            sub: sub.to_owned(),
            role: role.to_owned(),
            email: email.to_owned(),
            kind: "access".to_owned(),
            jti: new_jti(),
            iat: now,
            exp: now + self.config.access_ttl_secs,
            perms: perms.to_vec(),
            tenant: tenant.unwrap_or("").to_owned(),
            cnf: cnf_jkt.map(|jkt| Cnf {
                jkt: jkt.to_owned(),
            }),
        };
        encode(&self.header, &claims, &self.keys.load().encoding).map_err(JwtSignError)
    }

    /// Issue a signed **refresh token** with a unique `jti`.
    ///
    /// Claims: `sub`, `type = "refresh"`, `jti`, `iat`, `exp`.
    /// The `jti` is returned alongside the token so the caller can persist it.
    pub fn issue_refresh(&self, sub: &str) -> Result<(String, String), JwtSignError> {
        let now = Self::now();
        let jti = new_jti();
        let claims = JwtClaims {
            sub: sub.to_owned(),
            role: String::new(),
            email: String::new(),
            kind: "refresh".to_owned(),
            jti: jti.clone(),
            iat: now,
            exp: now + self.config.refresh_ttl_secs,
            perms: Vec::new(),
            tenant: String::new(),
            cnf: None,
        };
        let token =
            encode(&self.header, &claims, &self.keys.load().encoding).map_err(JwtSignError)?;
        Ok((token, jti))
    }

    /// Validate signature + expiry and return the decoded claims as a JSON map.
    ///
    /// Returns `None` for any invalid token (expired, bad signature, malformed).
    /// Does NOT enforce token type — use [`Self::decode_access`] at request boundaries.
    pub fn decode(&self, token: &str) -> Option<Arc<Claims>> {
        // Try current key first, then the retained previous key (rotation
        // grace window). Bundle is one atomic load — keys can't mix versions.
        let keys = self.keys.load();
        let data = keys
            .verify
            .iter()
            .find_map(|k| decode::<serde_json::Value>(token, k, &self.validation).ok())?;
        let obj = data.claims.as_object()?.clone();
        Some(Arc::new(obj))
    }

    /// Like [`decode`] but additionally requires `"type" == "access"`.
    ///
    /// Use this at request boundaries so refresh tokens cannot be passed as
    /// access tokens to authenticate protected routes.
    pub fn decode_access(&self, token: &str) -> Option<Arc<Claims>> {
        let claims = self.decode(token)?;
        if claims.get("type").and_then(|v| v.as_str()) != Some("access") {
            return None;
        }
        Some(claims)
    }

    /// Validate a **refresh token** specifically.
    ///
    /// Returns `(subject, jti)` on success, `None` otherwise.
    /// Callers must verify that the `jti` exists in their token store before
    /// issuing a new pair.
    pub fn validate_refresh(&self, token: &str) -> Option<(String, String)> {
        let claims = self.decode(token)?;
        if claims.get("type")?.as_str()? != "refresh" {
            return None;
        }
        let sub = claims.get("sub")?.as_str()?.to_owned();
        let jti = claims.get("jti")?.as_str()?.to_owned();
        Some((sub, jti))
    }

    /// Lifetime of access tokens in seconds (used in `TokenResponse.expires_in`).
    pub fn access_ttl_secs(&self) -> u64 {
        self.config.access_ttl_secs
    }

    /// Lifetime of refresh tokens (used by token store for TTL).
    pub fn refresh_ttl_secs(&self) -> u64 {
        self.config.refresh_ttl_secs
    }
}

/// Extract and decode an **access** Bearer token from request headers.
///
/// Shared by all three request boundaries (HTTP macro routes, plugin routes,
/// WebSocket handshake) so a security fix here applies everywhere at once.
///
/// Returns `None` when:
/// - No `JwtService` is registered in the container.
/// - The `Authorization` header is absent or not valid UTF-8.
/// - The token is missing, expired, or has an invalid signature.
/// - The token `"type"` claim is not `"access"` (i.e. refresh tokens are rejected).
pub fn decode_bearer_token(
    headers: &http::HeaderMap,
    container: &crate::core::engine::FrozenDiContainer,
) -> Option<Arc<Claims>> {
    let raw = headers.get("authorization")?.to_str().ok()?;
    let token = raw.strip_prefix("Bearer ").unwrap_or(raw).trim();
    if token.is_empty() {
        return None;
    }
    container.try_get::<JwtService>()?.decode_access(token)
}

/// Generate a collision-resistant JWT ID without external deps.
///
/// Combines a monotonic process counter with current time and thread ID so
/// two tokens issued concurrently on different threads in the same nanosecond
/// still get distinct JTIs.
fn new_jti() -> String {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    use std::sync::atomic::{AtomicU64, Ordering};

    static COUNTER: AtomicU64 = AtomicU64::new(0);
    let seq = COUNTER.fetch_add(1, Ordering::Relaxed);

    let mut h1 = DefaultHasher::new();
    SystemTime::now().hash(&mut h1);
    seq.hash(&mut h1);

    let mut h2 = DefaultHasher::new();
    std::thread::current().id().hash(&mut h2);
    seq.wrapping_add(1).hash(&mut h2);

    format!("{:016x}{:016x}", h1.finish(), h2.finish())
}