kirino 0.7.9

use anyhow::{anyhow, Result};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, sync::Arc};

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

/// JWT claims issued by [`JwtManager`] (root-crate wire format).
///
/// # `sub` semantics — read before mixing claim sets
///
/// This claim set binds the RFC 7519 `sub` to the **username**, with the
/// stable user UUID in [`Claims::user_id`]. The session crate's
/// `kirino_session::TokenClaims` does the opposite: its `sub` **is** the
/// user UUID (username in a dedicated field). Copying `sub` verbatim
/// between the two claim sets silently swaps identity semantics; map
/// through [`Claims::sub_username`] / [`Claims::sub_uuid`] instead.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Claims {
    /// JWT `sub` — the **username** in this claim set. Diverges from
    /// `kirino_session::TokenClaims`, whose `sub` is the user UUID; see the
    /// struct-level docs before migrating claims across crates.
    pub sub: String,
    /// Stable user identifier (UUID string). This is what `sub` carries in
    /// `kirino_session::TokenClaims`; prefer [`Claims::sub_uuid`] for
    /// migration-safe reads.
    pub user_id: String,
    pub roles: Vec<String>,
    #[serde(default)]
    pub permissions: Vec<String>,
    #[serde(default)]
    pub session_id: Option<String>,
    #[serde(default)]
    pub jti: Option<String>,
    pub iat: i64,
    pub exp: i64,
}

#[derive(Clone)]
pub struct JwtManager {
    encoding_key: EncodingKey,
    decoding_key: DecodingKey,
    expiration_hours: i64,
    revocation: Arc<dyn RevocationBackend>,
}

/// Audit kirino-R1-SESS-4: the revocation-table abstraction. The default
/// backend ([`MemoryRevocationBackend`]) is PROCESS-LOCAL — multi-replica
/// deployments must supply a shared backend (see
/// [`crate::database::db_revocation::DbRevocationStore`], feature
/// `rbac-db-session`) or revocations will not propagate across replicas.
#[async_trait::async_trait]
pub trait RevocationBackend: Send + Sync {
    /// Record `not_before` for `key` (user id, jti, or `session:<id>`).
    async fn revoke(&self, key: &str, not_before: i64);
    /// The recorded `not_before` for `key`, if any.
    async fn get_not_before(&self, key: &str) -> Option<i64>;
    /// Drop entries whose `not_before` predates `cutoff` (maintenance —
    /// entries older than the longest token lifetime can never match).
    async fn cleanup_before(&self, cutoff: i64);
}

/// The default process-local backend (audit kirino-R1-SESS-4): single-
/// replica semantics only, with the documented 50k-entry cap and 7-day
/// auto-cleanup.
#[derive(Default)]
pub struct MemoryRevocationBackend {
    map: tokio::sync::RwLock<HashMap<String, i64>>,
}

#[async_trait::async_trait]
impl RevocationBackend for MemoryRevocationBackend {
    async fn revoke(&self, key: &str, not_before: i64) {
        let mut map = self.map.write().await;
        map.insert(key.to_string(), not_before);
        if map.len() > REVOCATION_MAX_ENTRIES {
            let cutoff = not_before - REVOCATION_AUTO_CLEANUP_WINDOW_SECS;
            map.retain(|_, &mut nb| nb >= cutoff);
        }
    }

    async fn get_not_before(&self, key: &str) -> Option<i64> {
        self.map.read().await.get(key).copied()
    }

    async fn cleanup_before(&self, cutoff: i64) {
        self.map.write().await.retain(|_, &mut nb| nb >= cutoff);
    }
}

const MIN_JWT_SECRET_LENGTH: usize = 32;
const MAX_JWT_EXPIRATION_HOURS: i64 = 87600;
const REVOCATION_MAX_ENTRIES: usize = 50_000;
const REVOCATION_AUTO_CLEANUP_WINDOW_SECS: i64 = 86400 * 7;

impl Claims {
    /// The authenticated user's username.
    ///
    /// Migration-safe counterpart of `kirino_session::TokenClaims::sub_username`:
    /// returns the login name regardless of which claim set carries it.
    #[must_use]
    pub fn sub_username(&self) -> &str {
        &self.sub
    }

    /// The authenticated user's stable identifier (UUID string).
    ///
    /// Migration-safe counterpart of `kirino_session::TokenClaims::sub_uuid`:
    /// in this claim set the UUID lives in `user_id`, while in kirino-session
    /// it is the `sub` claim — this accessor returns it uniformly.
    #[must_use]
    pub fn sub_uuid(&self) -> &str {
        &self.user_id
    }
}

impl JwtManager {
    pub fn new(secret: &str, expiration_hours: i64) -> Result<Self> {
        if secret.len() < MIN_JWT_SECRET_LENGTH {
            return Err(anyhow!(
                "JWT secret must be at least {MIN_JWT_SECRET_LENGTH} bytes for HS256 security"
            ));
        }
        let expiration_hours = expiration_hours.clamp(1, MAX_JWT_EXPIRATION_HOURS);
        // Audit kirino-R1-SESS-3: revocations auto-clean after 7 days —
        // a token whose TTL exceeds that window has a replay gap between
        // revocation-entry expiry and token expiry. Make the mismatch
        // runtime-visible instead of silently accepted.
        if expiration_hours * 3600 > REVOCATION_AUTO_CLEANUP_WINDOW_SECS {
            tracing::warn!(
                expiration_hours,
                revocation_window_secs = REVOCATION_AUTO_CLEANUP_WINDOW_SECS,
                "token TTL exceeds the revocation auto-cleanup window; \
                 revoked tokens may be replayed after the entry expires"
            );
        }
        // Audit kirino-R1-SESS-4: make the process-local revocation
        // semantics visible at startup — multi-replica deployments that
        // need cross-replica revocation must supply a shared backend via
        // [`Self::with_revocation_backend`].
        tracing::info!(
            backend = "memory",
            "JWT revocation table is process-local; multi-replica deployments \
             must supply a shared RevocationBackend for cross-replica revocation"
        );
        Ok(Self {
            encoding_key: EncodingKey::from_secret(secret.as_bytes()),
            decoding_key: DecodingKey::from_secret(secret.as_bytes()),
            expiration_hours,
            revocation: Arc::new(MemoryRevocationBackend::default()),
        })
    }

    /// Issues a JWT for the given user.
    ///
    /// # Errors
    ///
    /// Returns an error if JWT encoding fails.
    pub fn issue(&self, user_id: &str, username: &str, roles: Vec<String>) -> Result<String> {
        self.issue_with_options(user_id, username, roles, vec![], None)
    }

    /// Issues a JWT with optional permissions and session ID.
    ///
    /// # Errors
    ///
    /// Returns an error if JWT encoding fails.
    pub fn issue_with_options(
        &self,
        user_id: &str,
        username: &str,
        roles: Vec<String>,
        permissions: Vec<String>,
        session_id: Option<String>,
    ) -> Result<String> {
        let now = Utc::now();
        let claims = Claims {
            sub: username.to_string(),
            user_id: user_id.to_string(),
            roles,
            permissions,
            session_id,
            jti: Some(uuid::Uuid::now_v7().to_string()),
            iat: now.timestamp(),
            exp: (now + chrono::Duration::hours(self.expiration_hours)).timestamp(),
        };
        encode(&Header::default(), &claims, &self.encoding_key)
            .map_err(|e| anyhow!("JWT encode failed: {e}"))
    }

    /// Verifies a JWT and returns the decoded claims.
    ///
    /// # Errors
    ///
    /// Returns an error if the token is invalid, expired, or has an invalid signature.
    pub fn verify(&self, token: &str) -> Result<Claims> {
        let data = decode::<Claims>(
            token,
            &self.decoding_key,
            &Validation::new(Algorithm::HS256),
        )
        .map_err(|e| anyhow!("JWT verify failed: {e}"))?;
        Ok(data.claims)
    }

    /// Verifies a JWT and checks it against the revocation list.
    ///
    /// # Errors
    ///
    /// Returns an error if the token is invalid, expired, or has been revoked.
    pub async fn verify_with_revocation(&self, token: &str) -> Result<Claims> {
        let claims = self.verify(token)?;
        if let Some(not_before) = self.revocation.get_not_before(&claims.user_id).await {
            if claims.iat <= not_before {
                return Err(anyhow!("token has been revoked"));
            }
        }
        if let Some(ref jti) = claims.jti {
            if let Some(not_before) = self.revocation.get_not_before(jti).await {
                if claims.iat <= not_before {
                    return Err(anyhow!("token has been revoked"));
                }
            }
        }
        if let Some(ref sid) = claims.session_id {
            let key = format!("session:{sid}");
            if let Some(not_before) = self.revocation.get_not_before(&key).await {
                if claims.iat <= not_before {
                    return Err(anyhow!("token has been revoked"));
                }
            }
        }
        Ok(claims)
    }

    pub async fn revoke_token(&self, jti: &str) {
        let not_before = Utc::now().timestamp();
        self.revocation.revoke(jti, not_before).await;
    }

    pub async fn revoke_all_for_user(&self, user_id: &str) {
        let not_before = Utc::now().timestamp() + 1;
        self.revocation.revoke(user_id, not_before).await;
    }

    /// Audit kirino-R1-SESS-4: supply a shared revocation backend (e.g.
    /// the DB-backed store) — required for multi-replica deployments so
    /// revocations propagate across replicas.
    pub fn with_revocation_backend(mut self, backend: Arc<dyn RevocationBackend>) -> Self {
        self.revocation = backend;
        self
    }

    /// The active revocation backend (exposed so hosts can run
    /// maintenance — e.g. the DB store's window cleanup — or downcast to
    /// a concrete backend for diagnostics).
    pub fn revocation_backend(&self) -> &Arc<dyn RevocationBackend> {
        &self.revocation
    }

    pub async fn revoke_session(&self, session_id: &str) {
        let not_before = Utc::now().timestamp();
        let key = format!("session:{session_id}");
        self.revocation.revoke(&key, not_before).await;
    }

    pub async fn cleanup_revocation(&self, max_token_lifetime_secs: i64) {
        let cutoff = Utc::now().timestamp() - max_token_lifetime_secs;
        self.revocation.cleanup_before(cutoff).await;
    }
}

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

    fn make_mgr() -> JwtManager {
        JwtManager::new("test-secret-that-is-at-least-32-bytes-long", 24).unwrap()
    }

    #[test]
    fn test_issue_and_verify() {
        let mgr = make_mgr();
        let token = mgr.issue("user-1", "alice", vec!["admin".into()]).unwrap();
        let claims = mgr.verify(&token).unwrap();
        assert_eq!(claims.sub, "alice");
        assert_eq!(claims.user_id, "user-1");
        assert_eq!(claims.roles, vec!["admin".to_string()]);
    }

    #[test]
    fn test_sub_accessors_pin_cross_crate_semantics() {
        let mgr = make_mgr();
        let uid = "0b7bfc9c-86f5-4e0e-a9bc-8e5a50bb5f9a";
        let token = mgr.issue(uid, "alice", vec![]).unwrap();
        let claims = mgr.verify(&token).unwrap();
        // Semantic accessors agree with kirino-session's TokenClaims
        // accessors regardless of which claim set carries which value.
        assert_eq!(claims.sub_username(), "alice");
        assert_eq!(claims.sub_uuid(), uid);
        // Wire compatibility: this claim set still carries the username in
        // `sub` and the UUID in `user_id` — do not silently change either.
        assert_eq!(claims.sub, "alice");
        assert_eq!(claims.user_id, uid);
    }

    #[test]
    fn test_invalid_token() {
        let mgr = make_mgr();
        assert!(mgr.verify("garbage.token.here").is_err());
    }

    #[test]
    fn test_tampered_token() {
        let mgr = make_mgr();
        let token = mgr.issue("user-1", "alice", vec!["admin".into()]).unwrap();
        let tampered = format!("{}.TAMPERED.{}", &token[..20], &token[token.len() - 20..]);
        assert!(mgr.verify(&tampered).is_err());
    }

    #[test]
    fn test_secret_too_short_returns_error() {
        assert!(JwtManager::new("short", 24).is_err());
    }

    #[test]
    fn test_expiration_hours_minimum_one() {
        let mgr = JwtManager::new("test-secret-that-is-at-least-32-bytes-long", 0).unwrap();
        let token = mgr.issue("u1", "user", vec![]).unwrap();
        let claims = mgr.verify(&token).unwrap();
        assert!(claims.exp > claims.iat);
    }

    #[tokio::test]
    async fn test_revoke_all_for_user() {
        let mgr = make_mgr();
        let token = mgr.issue("user-1", "alice", vec!["admin".into()]).unwrap();
        let claims = mgr.verify_with_revocation(&token).await.unwrap();
        assert_eq!(claims.sub, "alice");

        mgr.revoke_all_for_user("user-1").await;
        assert!(mgr.verify_with_revocation(&token).await.is_err());
    }

    #[tokio::test]
    async fn test_new_token_after_revocation() {
        let mgr = make_mgr();
        let old_token = mgr.issue("user-1", "alice", vec!["admin".into()]).unwrap();
        mgr.revoke_all_for_user("user-1").await;
        assert!(mgr.verify_with_revocation(&old_token).await.is_err());

        // Wait for the not_before window (now + 1) to pass
        tokio::time::sleep(std::time::Duration::from_secs(2)).await;
        let new_token = mgr.issue("user-1", "alice", vec!["admin".into()]).unwrap();
        assert!(mgr.verify_with_revocation(&new_token).await.is_ok());
    }

    #[tokio::test]
    async fn test_revocation_does_not_affect_other_users() {
        let mgr = make_mgr();
        let token_a = mgr.issue("user-a", "alice", vec!["admin".into()]).unwrap();
        let token_b = mgr.issue("user-b", "bob", vec!["viewer".into()]).unwrap();

        mgr.revoke_all_for_user("user-a").await;
        assert!(mgr.verify_with_revocation(&token_a).await.is_err());
        assert!(mgr.verify_with_revocation(&token_b).await.is_ok());
    }

    #[tokio::test]
    async fn test_cleanup_revocation() {
        let mgr = make_mgr();
        mgr.revoke_all_for_user("user-1").await;
        mgr.revoke_all_for_user("user-2").await;

        // Use a small positive lifetime to ensure entries are cleaned up
        // while accounting for the +1 second offset in revoke_all_for_user.
        const CLEANUP_WINDOW: i64 = 3;
        tokio::time::sleep(std::time::Duration::from_secs(CLEANUP_WINDOW as u64)).await;
        mgr.cleanup_revocation(CLEANUP_WINDOW).await;

        let token = mgr.issue("user-1", "alice", vec![]).unwrap();
        assert!(mgr.verify_with_revocation(&token).await.is_ok());
    }

    #[test]
    fn test_issue_with_permissions_and_session() {
        let mgr = make_mgr();
        let token = mgr
            .issue_with_options(
                "user-1",
                "alice",
                vec!["admin".into()],
                vec!["read".into(), "write".into()],
                Some("sess-123".into()),
            )
            .unwrap();
        let claims = mgr.verify(&token).unwrap();
        assert_eq!(claims.permissions, vec!["read", "write"]);
        assert_eq!(claims.session_id, Some("sess-123".to_string()));
    }

    #[test]
    fn test_backward_compat_old_token_deserialize() {
        let mgr = make_mgr();
        let token = mgr.issue("user-1", "alice", vec!["admin".into()]).unwrap();
        let claims = mgr.verify(&token).unwrap();
        assert!(claims.permissions.is_empty());
        assert!(claims.session_id.is_none());
    }

    #[test]
    fn test_claims_serializable() {
        let mgr = make_mgr();
        let token = mgr
            .issue_with_options("u1", "user", vec!["role".into()], vec!["perm".into()], None)
            .unwrap();
        let claims = mgr.verify(&token).unwrap();
        let json = serde_json::to_string(&claims).unwrap();
        assert!(json.contains("\"sub\":\"user\""));
        assert!(json.contains("\"permissions\":[\"perm\"]"));
    }
}