hms/crates/erp-core/src/crypto/key_manager.rs

use std::time::Instant;

use dashmap::DashMap;
use uuid::Uuid;

use crate::error::{AppError, AppResult};

use super::engine;

/// DEK 缓存条目 — Drop 时清零密钥材料
#[derive(Clone)]
struct CachedDek {
    dek: [u8; 32],
    version: u32,
    loaded_at: Instant,
}

impl Drop for CachedDek {
    fn drop(&mut self) {
        self.dek.fill(0);
    }
}

/// DEK 缓存管理 — 每租户独立 DEK，LRU + TTL
#[derive(Clone)]
pub struct DekManager {
    cache: DashMap<Uuid, CachedDek>,
    ttl_secs: u64,
    max_entries: usize,
}

impl DekManager {
    pub fn new(ttl_secs: u64, max_entries: usize) -> Self {
        Self {
            cache: DashMap::new(),
            ttl_secs,
            max_entries,
        }
    }

    /// 获取或创建租户的 DEK
    pub fn get_or_create_dek(
        &self,
        tenant_id: Uuid,
        encrypted_dek: Option<&str>,
        kek: &[u8; 32],
    ) -> AppResult<([u8; 32], u32)> {
        // 检查缓存
        if let Some(entry) = self.cache.get(&tenant_id)
            && entry.loaded_at.elapsed().as_secs() < self.ttl_secs
        {
            return Ok((entry.dek, entry.version));
        }

        // 从加密 DEK 解密
        if let Some(enc_dek) = encrypted_dek {
            let dek_hex = engine::decrypt(kek, enc_dek).map_err(AppError::Internal)?;
            let dek_bytes = hex::decode(&dek_hex).map_err(|e| AppError::Internal(e.to_string()))?;
            if dek_bytes.len() != 32 {
                return Err(AppError::Internal("DEK must be 32 bytes".into()));
            }
            let mut dek = [0u8; 32];
            dek.copy_from_slice(&dek_bytes);

            // 缓存（版本从外部传入时无法确定，使用默认值 1）
            self.evict_if_full();
            self.cache.insert(
                tenant_id,
                CachedDek {
                    dek,
                    version: 1,
                    loaded_at: Instant::now(),
                },
            );
            return Ok((dek, 1));
        }

        // 无现有 DEK → 生成新的
        let dek = Self::generate_dek();
        self.evict_if_full();
        self.cache.insert(
            tenant_id,
            CachedDek {
                dek,
                version: 1,
                loaded_at: Instant::now(),
            },
        );
        Ok((dek, 1))
    }

    /// 使用 KEK 加密 DEK 以便存储
    pub fn encrypt_dek_for_storage(dek: &[u8; 32], kek: &[u8; 32]) -> AppResult<String> {
        let dek_hex = hex::encode(dek);
        engine::encrypt(kek, &dek_hex).map_err(AppError::Internal)
    }

    /// 生成新 DEK 并用 KEK 加密，返回 (新 DEK, 加密后的 DEK)
    pub fn generate_new_dek(kek: &[u8; 32]) -> AppResult<([u8; 32], String)> {
        let dek = Self::generate_dek();
        let encrypted = Self::encrypt_dek_for_storage(&dek, kek)?;
        Ok((dek, encrypted))
    }

    /// 使缓存失效（轮换后调用）
    pub fn invalidate(&self, tenant_id: Uuid) {
        self.cache.remove(&tenant_id);
    }

    fn generate_dek() -> [u8; 32] {
        use rand::RngCore;
        let mut dek = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut dek);
        dek
    }

    fn evict_if_full(&self) {
        if self.cache.len() >= self.max_entries {
            let to_remove: Vec<Uuid> = self
                .cache
                .iter()
                .filter(|e| e.loaded_at.elapsed().as_secs() > self.ttl_secs / 2)
                .map(|e| *e.key())
                .take(self.max_entries / 2)
                .collect();
            for id in to_remove {
                self.cache.remove(&id);
            }
        }
    }
}

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

    fn test_kek() -> [u8; 32] {
        *PiiCrypto::dev_default().kek()
    }

    fn test_uuid(i: u8) -> Uuid {
        let s = format!("00000000-0000-0000-0000-0000000000{:02x}", i);
        Uuid::parse_str(&s).unwrap()
    }

    #[test]
    fn generate_new_dek_returns_32_bytes() {
        let (dek, _enc) = DekManager::generate_new_dek(&test_kek()).unwrap();
        assert_eq!(dek.len(), 32);
    }

    #[test]
    fn generate_new_dek_produces_unique_keys() {
        let (dek1, _) = DekManager::generate_new_dek(&test_kek()).unwrap();
        let (dek2, _) = DekManager::generate_new_dek(&test_kek()).unwrap();
        assert_ne!(dek1, dek2);
    }

    #[test]
    fn encrypt_dek_roundtrip() {
        let kek = test_kek();
        let (original_dek, encrypted) = DekManager::generate_new_dek(&kek).unwrap();
        let mgr = DekManager::new(300, 100);
        let tenant_id = test_uuid(1);
        let (recovered_dek, _ver) = mgr
            .get_or_create_dek(tenant_id, Some(&encrypted), &kek)
            .unwrap();
        assert_eq!(original_dek, recovered_dek);
    }

    #[test]
    fn get_or_create_generates_when_none() {
        let mgr = DekManager::new(300, 100);
        let tenant_id = test_uuid(2);
        let (dek1, ver1) = mgr.get_or_create_dek(tenant_id, None, &test_kek()).unwrap();
        assert_eq!(ver1, 1);
        let (dek2, ver2) = mgr.get_or_create_dek(tenant_id, None, &test_kek()).unwrap();
        assert_eq!(dek1, dek2);
        assert_eq!(ver2, 1);
    }

    #[test]
    fn invalidate_removes_cached_dek() {
        let mgr = DekManager::new(300, 100);
        let tenant_id = test_uuid(3);
        let (dek1, _) = mgr.get_or_create_dek(tenant_id, None, &test_kek()).unwrap();
        mgr.invalidate(tenant_id);
        let (dek2, _) = mgr.get_or_create_dek(tenant_id, None, &test_kek()).unwrap();
        assert_ne!(dek1, dek2);
    }

    #[test]
    fn decrypt_with_wrong_kek_fails() {
        let kek1 = test_kek();
        let kek2 = [0xffu8; 32];
        let (_, encrypted) = DekManager::generate_new_dek(&kek1).unwrap();
        let mgr = DekManager::new(300, 100);
        let tenant_id = test_uuid(4);
        assert!(
            mgr.get_or_create_dek(tenant_id, Some(&encrypted), &kek2)
                .is_err()
        );
    }

    #[test]
    fn expired_entry_not_returned() {
        let mgr = DekManager::new(0, 100);
        let tenant_id = test_uuid(5);
        let (dek1, _) = mgr.get_or_create_dek(tenant_id, None, &test_kek()).unwrap();
        let (dek2, _) = mgr.get_or_create_dek(tenant_id, None, &test_kek()).unwrap();
        assert_ne!(dek1, dek2);
    }

    #[test]
    fn max_entries_eviction() {
        let mgr = DekManager::new(300, 3);
        for i in 0..5u8 {
            let _ = mgr
                .get_or_create_dek(test_uuid(i), None, &test_kek())
                .unwrap();
        }
        assert!(mgr.cache.len() <= 6);
    }
}