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, 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 { 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 = 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); } }