arcly_http/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.
pub struct JwtConfig {
    /// HMAC secret (HS256 / HS384 / HS512) or PEM-encoded key for RS/ES algorithms.
    pub secret: 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(),
            algorithm: Algorithm::HS256,
            access_ttl_secs: 900,
            refresh_ttl_secs: 604_800,
        }
    }
}

// ─── 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,
}

// ─── 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 {
    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,
        }
    }
}

/// 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.
pub struct JwtService {
    keys: crate::auth::secrets::Rotating<JwtKeyMaterial>,
    header: Header,
    validation: Validation,
    config: JwtConfig,
}

impl JwtService {
    pub fn new(config: JwtConfig) -> Self {
        let keys = crate::auth::secrets::Rotating::new(JwtKeyMaterial::from_secret(
            config.secret.as_bytes(),
            1,
            None,
        ));
        let header = Header::new(config.algorithm);
        let mut validation = Validation::new(config.algorithm);
        validation.validate_exp = true;
        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.
    pub fn rotate_secret(&self, new_secret: &[u8], 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();
        self.keys
            .store(JwtKeyMaterial::from_secret(new_secret, version, previous));
        tracing::info!(version, "JwtService signing key rotated");
    }

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

    /// Issue a signed **access token**.
    ///
    /// Claims: `sub` = user ID, `role`, `email`, `type = "access"`, `jti`, `iat`, `exp`.
    pub fn issue_access(&self, sub: &str, role: &str, email: &str) -> String {
        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],
    ) -> String {
        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>,
    ) -> String {
        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(),
        };
        encode(&self.header, &claims, &self.keys.load().encoding)
            .expect("JWT encode failed — check signing key")
    }

    /// 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) -> (String, String) {
        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(),
        };
        let token = encode(&self.header, &claims, &self.keys.load().encoding)
            .expect("JWT encode failed — check signing key");
        (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: &axum::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())
}