zclaw_openfang/crates/zclaw-saas/src/auth/jwt.rs

//! JWT Token 创建与验证

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

use crate::error::SaasResult;

/// JWT Claims
#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    /// JWT ID — 唯一标识，用于 token 追踪和吊销
    pub jti: Option<String>,
    pub sub: String,
    pub role: String,
    pub permissions: Vec<String>,
    /// token 类型: "access" 或 "refresh"
    #[serde(default = "default_token_type")]
    pub token_type: String,
    pub iat: i64,
    pub exp: i64,
}

fn default_token_type() -> String {
    "access".to_string()
}

impl Claims {
    pub fn new_access(account_id: &str, role: &str, permissions: Vec<String>, expiration_hours: i64) -> Self {
        let now = Utc::now();
        Self {
            jti: Some(uuid::Uuid::new_v4().to_string()),
            sub: account_id.to_string(),
            role: role.to_string(),
            permissions,
            token_type: "access".to_string(),
            iat: now.timestamp(),
            exp: (now + Duration::hours(expiration_hours)).timestamp(),
        }
    }

    /// 创建 refresh token claims (有效期更长，用于一次性刷新)
    pub fn new_refresh(account_id: &str, role: &str, permissions: Vec<String>, refresh_hours: i64) -> Self {
        let now = Utc::now();
        Self {
            jti: Some(uuid::Uuid::new_v4().to_string()),
            sub: account_id.to_string(),
            role: role.to_string(),
            permissions,
            token_type: "refresh".to_string(),
            iat: now.timestamp(),
            exp: (now + Duration::hours(refresh_hours)).timestamp(),
        }
    }
}

/// 创建 Access JWT Token
pub fn create_token(
    account_id: &str,
    role: &str,
    permissions: Vec<String>,
    secret: &str,
    expiration_hours: i64,
) -> SaasResult<String> {
    let claims = Claims::new_access(account_id, role, permissions, expiration_hours);
    let token = encode(
        &Header::default(),
        &claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    )?;
    Ok(token)
}

/// 创建 Refresh JWT Token (独立 jti，有效期更长)
pub fn create_refresh_token(
    account_id: &str,
    role: &str,
    permissions: Vec<String>,
    secret: &str,
    refresh_hours: i64,
) -> SaasResult<String> {
    let claims = Claims::new_refresh(account_id, role, permissions, refresh_hours);
    let token = encode(
        &Header::default(),
        &claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    )?;
    Ok(token)
}

/// 验证 JWT Token
pub fn verify_token(token: &str, secret: &str) -> SaasResult<Claims> {
    let token_data = decode::<Claims>(
        token,
        &DecodingKey::from_secret(secret.as_bytes()),
        &Validation::default(),
    )?;
    Ok(token_data.claims)
}

/// 验证 JWT Token 但跳过过期检查（仅用于 refresh token 刷新）
/// 限制: 原始 token 的 iat 必须在 7 天内
pub fn verify_token_skip_expiry(token: &str, secret: &str) -> SaasResult<Claims> {
    let mut validation = Validation::default();
    validation.validate_exp = false;
    let token_data = decode::<Claims>(
        token,
        &DecodingKey::from_secret(secret.as_bytes()),
        &validation,
    )?;
    let claims = &token_data.claims;

    // 限制刷新窗口: token 签发时间必须在 7 天内
    let now = Utc::now().timestamp();
    let max_refresh_window = 7 * 24 * 3600; // 7 天
    if now - claims.iat > max_refresh_window {
        return Err(jsonwebtoken::errors::Error::from(
            jsonwebtoken::errors::ErrorKind::ExpiredSignature
        ).into());
    }

    Ok(token_data.claims)
}

/// Token 对: access token + refresh token
#[derive(Debug, serde::Serialize)]
pub struct TokenPair {
    pub access_token: String,
    pub refresh_token: String,
}

/// 创建 access + refresh token 对
pub fn create_token_pair(
    account_id: &str,
    role: &str,
    permissions: Vec<String>,
    secret: &str,
    access_hours: i64,
    refresh_hours: i64,
) -> SaasResult<TokenPair> {
    Ok(TokenPair {
        access_token: create_token(account_id, role, permissions.clone(), secret, access_hours)?,
        refresh_token: create_refresh_token(account_id, role, permissions, secret, refresh_hours)?,
    })
}

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

    const TEST_SECRET: &str = "test-secret-key";

    #[test]
    fn test_create_and_verify_token() {
        let token = create_token(
            "account-123", "admin",
            vec!["model:read".to_string()],
            TEST_SECRET, 24,
        ).unwrap();

        let claims = verify_token(&token, TEST_SECRET).unwrap();
        assert_eq!(claims.sub, "account-123");
        assert_eq!(claims.role, "admin");
        assert_eq!(claims.permissions, vec!["model:read"]);
        assert!(claims.jti.is_some());
        assert_eq!(claims.token_type, "access");
    }

    #[test]
    fn test_invalid_token() {
        let result = verify_token("invalid.token.here", TEST_SECRET);
        assert!(result.is_err());
    }

    #[test]
    fn test_wrong_secret() {
        let token = create_token("account-123", "admin", vec![], TEST_SECRET, 24).unwrap();
        let result = verify_token(&token, "wrong-secret");
        assert!(result.is_err());
    }

    #[test]
    fn test_refresh_token_has_different_jti() {
        let access = create_token("acct-1", "user", vec![], TEST_SECRET, 1).unwrap();
        let refresh = create_refresh_token("acct-1", "user", vec![], TEST_SECRET, 168).unwrap();

        let access_claims = verify_token(&access, TEST_SECRET).unwrap();
        let refresh_claims = verify_token(&refresh, TEST_SECRET).unwrap();

        assert_ne!(access_claims.jti, refresh_claims.jti);
        assert_eq!(access_claims.token_type, "access");
        assert_eq!(refresh_claims.token_type, "refresh");
    }
}