hms/crates/erp-auth/src/service/token_service.rs

use chrono::Utc;
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use uuid::Uuid;

use crate::entity::{permission, role, role_permission, user_role, user_token};
use crate::error::AuthError;

use crate::error::AuthResult;

/// JWT claims embedded in access and refresh tokens.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Claims {
    /// Subject — the user ID
    pub sub: Uuid,
    /// Tenant ID
    pub tid: Uuid,
    /// Role codes assigned to this user
    pub roles: Vec<String>,
    /// Permission codes granted to this user
    pub permissions: Vec<String>,
    /// Expiry (unix timestamp)
    pub exp: i64,
    /// Issued at (unix timestamp)
    pub iat: i64,
    /// Token type: "access" or "refresh"
    pub token_type: String,
}

/// Stateless service for JWT token signing, validation, and revocation.
pub struct TokenService;

impl TokenService {
    /// Sign a short-lived access token containing roles and permissions.
    pub fn sign_access_token(
        user_id: Uuid,
        tenant_id: Uuid,
        roles: Vec<String>,
        permissions: Vec<String>,
        secret: &str,
        ttl_secs: i64,
    ) -> AuthResult<String> {
        let now = Utc::now();
        let claims = Claims {
            sub: user_id,
            tid: tenant_id,
            roles,
            permissions,
            exp: now.timestamp() + ttl_secs,
            iat: now.timestamp(),
            token_type: "access".to_string(),
        };
        let header = jsonwebtoken::Header::default();
        let encoded = jsonwebtoken::encode(
            &header,
            &claims,
            &jsonwebtoken::EncodingKey::from_secret(secret.as_bytes()),
        )?;
        Ok(encoded)
    }

    /// Sign a long-lived refresh token and persist its SHA-256 hash in the database.
    ///
    /// Returns the raw token string (sent to client) and the database row ID.
    pub async fn sign_refresh_token(
        user_id: Uuid,
        tenant_id: Uuid,
        db: &DatabaseConnection,
        secret: &str,
        ttl_secs: i64,
    ) -> AuthResult<(String, Uuid)> {
        let now = Utc::now();
        let token_id = Uuid::now_v7();

        let claims = Claims {
            sub: user_id,
            tid: tenant_id,
            roles: vec![],
            permissions: vec![],
            exp: now.timestamp() + ttl_secs,
            iat: now.timestamp(),
            token_type: "refresh".to_string(),
        };
        let raw_token = jsonwebtoken::encode(
            &jsonwebtoken::Header::default(),
            &claims,
            &jsonwebtoken::EncodingKey::from_secret(secret.as_bytes()),
        )?;

        // Store the SHA-256 hash — the raw token is never persisted.
        let hash = sha256_hex(&raw_token);

        let token_model = user_token::ActiveModel {
            id: Set(token_id),
            tenant_id: Set(tenant_id),
            user_id: Set(user_id),
            token_hash: Set(hash),
            token_type: Set("refresh".to_string()),
            expires_at: Set(now + chrono::Duration::seconds(ttl_secs)),
            revoked_at: Set(None),
            device_info: Set(None),
            created_at: Set(now),
            updated_at: Set(now),
            created_by: Set(user_id),
            updated_by: Set(user_id),
            deleted_at: Set(None),
            version: Set(1),
        };
        token_model
            .insert(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?;

        Ok((raw_token, token_id))
    }

    /// Validate a refresh token against the database.
    ///
    /// Returns the database row ID and decoded claims.
    pub async fn validate_refresh_token(
        token: &str,
        db: &DatabaseConnection,
        secret: &str,
    ) -> AuthResult<(Uuid, Claims)> {
        let claims = Self::decode_token(token, secret)?;
        if claims.token_type != "refresh" {
            return Err(AuthError::Validation("不是 refresh token".to_string()));
        }

        let hash = sha256_hex(token);
        let token_row = user_token::Entity::find()
            .filter(user_token::Column::TokenHash.eq(hash))
            .filter(user_token::Column::TenantId.eq(claims.tid))
            .filter(user_token::Column::RevokedAt.is_null())
            .one(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?
            .ok_or(AuthError::TokenRevoked)?;

        Ok((token_row.id, claims))
    }

    /// Decode and validate any JWT token, returning the claims.
    pub fn decode_token(token: &str, secret: &str) -> AuthResult<Claims> {
        let data = jsonwebtoken::decode::<Claims>(
            token,
            &jsonwebtoken::DecodingKey::from_secret(secret.as_bytes()),
            &jsonwebtoken::Validation::default(),
        )?;
        Ok(data.claims)
    }

    /// Revoke a specific refresh token by database ID.
    /// Verifies that the token belongs to the specified user for security.
    pub async fn revoke_token(
        token_id: Uuid,
        user_id: Uuid,
        db: &DatabaseConnection,
    ) -> AuthResult<()> {
        let token_row = user_token::Entity::find_by_id(token_id)
            .filter(user_token::Column::UserId.eq(user_id))
            .one(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?
            .ok_or(AuthError::TokenRevoked)?;

        let mut active: user_token::ActiveModel = token_row.into();
        active.revoked_at = Set(Some(Utc::now()));
        active.updated_at = Set(Utc::now());
        active
            .update(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?;
        Ok(())
    }

    /// Revoke all non-revoked refresh tokens for a given user within a tenant.
    pub async fn revoke_all_user_tokens(
        user_id: Uuid,
        tenant_id: Uuid,
        db: &DatabaseConnection,
    ) -> AuthResult<()> {
        let tokens = user_token::Entity::find()
            .filter(user_token::Column::UserId.eq(user_id))
            .filter(user_token::Column::TenantId.eq(tenant_id))
            .filter(user_token::Column::RevokedAt.is_null())
            .all(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?;

        let now = Utc::now();
        for token in tokens {
            let mut active: user_token::ActiveModel = token.into();
            active.revoked_at = Set(Some(now));
            active.updated_at = Set(now);
            active
                .update(db)
                .await
                .map_err(|e| AuthError::Validation(e.to_string()))?;
        }
        Ok(())
    }

    /// Look up a user's permission codes through user_roles -> role_permissions -> permissions.
    pub async fn get_user_permissions(
        user_id: Uuid,
        tenant_id: Uuid,
        db: &DatabaseConnection,
    ) -> AuthResult<Vec<String>> {
        let user_role_rows = user_role::Entity::find()
            .filter(user_role::Column::UserId.eq(user_id))
            .filter(user_role::Column::TenantId.eq(tenant_id))
            .filter(user_role::Column::DeletedAt.is_null())
            .all(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?;

        let role_ids: Vec<Uuid> = user_role_rows.iter().map(|ur| ur.role_id).collect();
        if role_ids.is_empty() {
            return Ok(vec![]);
        }

        let role_perm_rows = role_permission::Entity::find()
            .filter(role_permission::Column::RoleId.is_in(role_ids))
            .filter(role_permission::Column::TenantId.eq(tenant_id))
            .filter(role_permission::Column::DeletedAt.is_null())
            .all(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?;

        let perm_ids: Vec<Uuid> = role_perm_rows.iter().map(|rp| rp.permission_id).collect();
        if perm_ids.is_empty() {
            return Ok(vec![]);
        }

        let perms = permission::Entity::find()
            .filter(permission::Column::Id.is_in(perm_ids))
            .filter(permission::Column::TenantId.eq(tenant_id))
            .all(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?;

        Ok(perms.iter().map(|p| p.code.clone()).collect())
    }

    /// Look up a user's role codes through user_roles -> roles.
    pub async fn get_user_roles(
        user_id: Uuid,
        tenant_id: Uuid,
        db: &DatabaseConnection,
    ) -> AuthResult<Vec<String>> {
        let user_role_rows = user_role::Entity::find()
            .filter(user_role::Column::UserId.eq(user_id))
            .filter(user_role::Column::TenantId.eq(tenant_id))
            .filter(user_role::Column::DeletedAt.is_null())
            .all(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?;

        let role_ids: Vec<Uuid> = user_role_rows.iter().map(|ur| ur.role_id).collect();
        if role_ids.is_empty() {
            return Ok(vec![]);
        }

        let roles = role::Entity::find()
            .filter(role::Column::Id.is_in(role_ids))
            .filter(role::Column::TenantId.eq(tenant_id))
            .all(db)
            .await
            .map_err(|e| AuthError::Validation(e.to_string()))?;

        Ok(roles.iter().map(|r| r.code.clone()).collect())
    }
}

/// Compute a SHA-256 hex digest of the input string.
fn sha256_hex(input: &str) -> String {
    let hash = Sha256::digest(input.as_bytes());
    format!("{:x}", hash)
}