Skip to main content

sh_layer0/
secrets_manager.rs

1//! 密钥管理模块
2//!
3//! 安全地管理 API 密钥和其他敏感凭证。
4//!
5//! ## 功能
6//! - 从环境变量安全读取密钥
7//! - 内存中加密存储
8//! - 密钥轮换支持
9//! - 访问审计日志
10
11use parking_lot::RwLock;
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14use std::time::Instant;
15use thiserror::Error;
16
17/// 密钥管理错误
18#[derive(Debug, Error)]
19pub enum SecretsError {
20    #[error("Secret not found: {0}")]
21    NotFound(String),
22
23    #[error("Environment variable not set: {0}")]
24    EnvNotSet(String),
25
26    #[error("Secret rotation failed: {0}")]
27    RotationFailed(String),
28
29    #[error("Encryption error: {0}")]
30    EncryptionError(String),
31}
32
33/// 密钥元数据
34#[derive(Debug, Clone)]
35struct SecretMetadata {
36    /// 密钥名称
37    name: String,
38    /// 创建时间
39    created_at: Instant,
40    /// 最后访问时间
41    last_accessed: Instant,
42    /// 最后轮换时间
43    last_rotated: Instant,
44    /// 访问次数
45    access_count: u64,
46    /// 是否需要轮换
47    requires_rotation: bool,
48}
49
50/// 审计日志条目
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct AuditLogEntry {
53    /// 时间戳
54    pub timestamp: String,
55    /// 操作类型
56    pub action: AuditAction,
57    /// 密钥名称
58    pub secret_name: String,
59    /// 结果
60    pub result: String,
61}
62
63/// 审计操作类型
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub enum AuditAction {
66    /// 读取密钥
67    Read,
68    /// 设置密钥
69    Set,
70    /// 轮换密钥
71    Rotate,
72    /// 删除密钥
73    Delete,
74    /// 从环境变量加载
75    LoadFromEnv,
76}
77
78/// 密钥管理器配置
79#[derive(Debug, Clone)]
80pub struct SecretsManagerConfig {
81    /// 是否启用审计日志
82    pub audit_enabled: bool,
83    /// 密钥轮换周期(秒)
84    pub rotation_interval_secs: u64,
85    /// 是否在读取时自动检查轮换
86    pub auto_rotation_check: bool,
87}
88
89impl Default for SecretsManagerConfig {
90    fn default() -> Self {
91        Self {
92            audit_enabled: true,
93            rotation_interval_secs: 86400 * 30, // 30 天
94            auto_rotation_check: true,
95        }
96    }
97}
98
99/// 内存加密存储
100struct EncryptedStorage {
101    /// 加密后的密钥存储
102    encrypted_secrets: RwLock<HashMap<String, Vec<u8>>>,
103    /// 混淆密钥(简单实现)
104    obfuscation_key: [u8; 32],
105}
106
107impl EncryptedStorage {
108    fn new() -> Self {
109        // 生成随机混淆密钥
110        let obfuscation_key: [u8; 32] = {
111            use std::collections::hash_map::RandomState;
112            use std::hash::{BuildHasher, Hasher};
113            let state = RandomState::new();
114            let mut hasher = state.build_hasher();
115            hasher.write_u64(std::process::id() as u64);
116            hasher.write_u64(Instant::now().elapsed().as_nanos() as u64);
117            let hash = hasher.finish();
118            let mut key = [0u8; 32];
119            for (i, byte) in key.iter_mut().enumerate() {
120                *byte = ((hash >> (i % 8 * 8)) & 0xFF) as u8;
121            }
122            key
123        };
124
125        Self {
126            encrypted_secrets: RwLock::new(HashMap::new()),
127            obfuscation_key,
128        }
129    }
130
131    /// 简单 XOR 加密
132    fn encrypt(&self, plaintext: &str) -> Vec<u8> {
133        plaintext
134            .as_bytes()
135            .iter()
136            .enumerate()
137            .map(|(i, &byte)| byte ^ self.obfuscation_key[i % 32])
138            .collect()
139    }
140
141    /// 简单 XOR 解密
142    fn decrypt(&self, ciphertext: &[u8]) -> Result<String, SecretsError> {
143        let decrypted: Vec<u8> = ciphertext
144            .iter()
145            .enumerate()
146            .map(|(i, &byte)| byte ^ self.obfuscation_key[i % 32])
147            .collect();
148        String::from_utf8(decrypted).map_err(|e| SecretsError::EncryptionError(e.to_string()))
149    }
150
151    fn set(&self, key: &str, value: &str) {
152        let encrypted = self.encrypt(value);
153        self.encrypted_secrets
154            .write()
155            .insert(key.to_string(), encrypted);
156    }
157
158    fn get(&self, key: &str) -> Result<Option<String>, SecretsError> {
159        let storage = self.encrypted_secrets.read();
160        match storage.get(key) {
161            Some(encrypted) => {
162                let decrypted = self.decrypt(encrypted)?;
163                Ok(Some(decrypted))
164            }
165            None => Ok(None),
166        }
167    }
168
169    fn remove(&self, key: &str) {
170        self.encrypted_secrets.write().remove(key);
171    }
172
173    fn contains(&self, key: &str) -> bool {
174        self.encrypted_secrets.read().contains_key(key)
175    }
176
177    fn keys(&self) -> Vec<String> {
178        self.encrypted_secrets.read().keys().cloned().collect()
179    }
180}
181
182/// 密钥管理器
183pub struct SecretsManager {
184    /// 加密存储
185    storage: EncryptedStorage,
186    /// 密钥元数据
187    metadata: RwLock<HashMap<String, SecretMetadata>>,
188    /// 审计日志
189    audit_log: RwLock<Vec<AuditLogEntry>>,
190    /// 配置
191    config: SecretsManagerConfig,
192}
193
194impl SecretsManager {
195    /// 创建新的密钥管理器
196    pub fn new() -> Self {
197        Self::with_config(SecretsManagerConfig::default())
198    }
199
200    /// 使用自定义配置创建
201    pub fn with_config(config: SecretsManagerConfig) -> Self {
202        Self {
203            storage: EncryptedStorage::new(),
204            metadata: RwLock::new(HashMap::new()),
205            audit_log: RwLock::new(Vec::new()),
206            config,
207        }
208    }
209
210    /// 从环境变量加载密钥
211    pub fn load_from_env(&self, key: &str, env_var: &str) -> Result<(), SecretsError> {
212        let value =
213            std::env::var(env_var).map_err(|_| SecretsError::EnvNotSet(env_var.to_string()))?;
214
215        self.set(key, &value)?;
216        self.log_audit(AuditAction::LoadFromEnv, key, "success")?;
217
218        Ok(())
219    }
220
221    /// 设置密钥
222    pub fn set(&self, key: &str, value: &str) -> Result<(), SecretsError> {
223        let now = Instant::now();
224
225        // 存储加密后的值
226        self.storage.set(key, value);
227
228        // 更新元数据
229        let mut metadata = self.metadata.write();
230        metadata.insert(
231            key.to_string(),
232            SecretMetadata {
233                name: key.to_string(),
234                created_at: now,
235                last_accessed: now,
236                last_rotated: now,
237                access_count: 0,
238                requires_rotation: false,
239            },
240        );
241
242        self.log_audit(AuditAction::Set, key, "success")?;
243        Ok(())
244    }
245
246    /// 获取密钥
247    pub fn get(&self, key: &str) -> Result<Option<String>, SecretsError> {
248        let result = self.storage.get(key)?;
249
250        // 更新访问元数据
251        if result.is_some() {
252            let mut metadata = self.metadata.write();
253            if let Some(meta) = metadata.get_mut(key) {
254                meta.last_accessed = Instant::now();
255                meta.access_count += 1;
256
257                // 检查是否需要轮换
258                if self.config.auto_rotation_check {
259                    let elapsed = meta.last_rotated.elapsed().as_secs();
260                    if elapsed >= self.config.rotation_interval_secs {
261                        meta.requires_rotation = true;
262                    }
263                }
264            }
265        }
266
267        self.log_audit(
268            AuditAction::Read,
269            key,
270            if result.is_some() {
271                "found"
272            } else {
273                "not_found"
274            },
275        )?;
276        Ok(result)
277    }
278
279    /// 轮换密钥
280    pub fn rotate(&self, key: &str, new_value: &str) -> Result<(), SecretsError> {
281        // 检查密钥是否存在
282        if !self.storage.contains(key) {
283            return Err(SecretsError::NotFound(key.to_string()));
284        }
285
286        // 更新密钥值
287        self.storage.set(key, new_value);
288
289        // 更新元数据
290        let mut metadata = self.metadata.write();
291        if let Some(meta) = metadata.get_mut(key) {
292            meta.last_rotated = Instant::now();
293            meta.requires_rotation = false;
294        }
295
296        self.log_audit(AuditAction::Rotate, key, "success")?;
297        Ok(())
298    }
299
300    /// 删除密钥
301    pub fn delete(&self, key: &str) -> Result<(), SecretsError> {
302        if !self.storage.contains(key) {
303            return Err(SecretsError::NotFound(key.to_string()));
304        }
305
306        self.storage.remove(key);
307        self.metadata.write().remove(key);
308
309        self.log_audit(AuditAction::Delete, key, "success")?;
310        Ok(())
311    }
312
313    /// 检查密钥是否存在
314    pub fn contains(&self, key: &str) -> bool {
315        self.storage.contains(key)
316    }
317
318    /// 获取需要轮换的密钥列表
319    pub fn get_keys_requiring_rotation(&self) -> Vec<String> {
320        self.metadata
321            .read()
322            .iter()
323            .filter(|(_, meta)| meta.requires_rotation)
324            .map(|(key, _)| key.clone())
325            .collect()
326    }
327
328    /// 获取所有密钥名称
329    pub fn list_keys(&self) -> Vec<String> {
330        self.storage.keys()
331    }
332
333    /// 获取密钥元数据(不包含实际值)
334    pub fn get_metadata(&self, key: &str) -> Option<SecretMetadataInfo> {
335        self.metadata
336            .read()
337            .get(key)
338            .map(|meta| SecretMetadataInfo {
339                name: meta.name.clone(),
340                age_secs: meta.created_at.elapsed().as_secs(),
341                last_accessed_secs_ago: meta.last_accessed.elapsed().as_secs(),
342                last_rotated_secs_ago: meta.last_rotated.elapsed().as_secs(),
343                access_count: meta.access_count,
344                requires_rotation: meta.requires_rotation,
345            })
346    }
347
348    /// 获取审计日志
349    pub fn get_audit_log(&self) -> Vec<AuditLogEntry> {
350        self.audit_log.read().clone()
351    }
352
353    /// 清空审计日志
354    pub fn clear_audit_log(&self) {
355        self.audit_log.write().clear();
356    }
357
358    /// 记录审计日志
359    fn log_audit(
360        &self,
361        action: AuditAction,
362        secret_name: &str,
363        result: &str,
364    ) -> Result<(), SecretsError> {
365        if !self.config.audit_enabled {
366            return Ok(());
367        }
368
369        let entry = AuditLogEntry {
370            timestamp: chrono::Utc::now().to_rfc3339(),
371            action,
372            secret_name: secret_name.to_string(),
373            result: result.to_string(),
374        };
375
376        self.audit_log.write().push(entry);
377        Ok(())
378    }
379}
380
381/// 密钥元数据信息(公开)
382#[derive(Debug, Clone)]
383pub struct SecretMetadataInfo {
384    /// 密钥名称
385    pub name: String,
386    /// 密钥年龄(秒)
387    pub age_secs: u64,
388    /// 最后访问距今(秒)
389    pub last_accessed_secs_ago: u64,
390    /// 最后轮换距今(秒)
391    pub last_rotated_secs_ago: u64,
392    /// 访问次数
393    pub access_count: u64,
394    /// 是否需要轮换
395    pub requires_rotation: bool,
396}
397
398impl Default for SecretsManager {
399    fn default() -> Self {
400        Self::new()
401    }
402}
403
404#[cfg(test)]
405mod tests {
406    use super::*;
407
408    #[test]
409    fn test_set_and_get() {
410        let manager = SecretsManager::new();
411
412        manager.set("api_key", "secret123").unwrap();
413        let value = manager.get("api_key").unwrap();
414
415        assert_eq!(value, Some("secret123".to_string()));
416    }
417
418    #[test]
419    fn test_not_found() {
420        let manager = SecretsManager::new();
421        let value = manager.get("nonexistent").unwrap();
422        assert!(value.is_none());
423    }
424
425    #[test]
426    fn test_delete() {
427        let manager = SecretsManager::new();
428
429        manager.set("api_key", "secret123").unwrap();
430        assert!(manager.contains("api_key"));
431
432        manager.delete("api_key").unwrap();
433        assert!(!manager.contains("api_key"));
434    }
435
436    #[test]
437    fn test_rotate() {
438        let manager = SecretsManager::new();
439
440        manager.set("api_key", "secret123").unwrap();
441        manager.rotate("api_key", "new_secret456").unwrap();
442
443        let value = manager.get("api_key").unwrap();
444        assert_eq!(value, Some("new_secret456".to_string()));
445    }
446
447    #[test]
448    fn test_audit_log() {
449        let manager = SecretsManager::new();
450
451        manager.set("api_key", "secret123").unwrap();
452        manager.get("api_key").unwrap();
453        manager.rotate("api_key", "new_secret").unwrap();
454
455        let log = manager.get_audit_log();
456        assert_eq!(log.len(), 3);
457        assert!(matches!(log[0].action, AuditAction::Set));
458        assert!(matches!(log[1].action, AuditAction::Read));
459        assert!(matches!(log[2].action, AuditAction::Rotate));
460    }
461
462    #[test]
463    fn test_access_count() {
464        let manager = SecretsManager::new();
465
466        manager.set("api_key", "secret123").unwrap();
467
468        for _ in 0..5 {
469            manager.get("api_key").unwrap();
470        }
471
472        let meta = manager.get_metadata("api_key").unwrap();
473        assert_eq!(meta.access_count, 5);
474    }
475
476    #[test]
477    fn test_encrypted_storage() {
478        let storage = EncryptedStorage::new();
479
480        storage.set("test_key", "test_value");
481        let value = storage.get("test_key").unwrap();
482
483        assert_eq!(value, Some("test_value".to_string()));
484    }
485
486    #[test]
487    fn test_encryption_roundtrip() {
488        let storage = EncryptedStorage::new();
489        let plaintext = "my_secret_api_key_12345";
490
491        storage.set("key", plaintext);
492        let decrypted = storage.get("key").unwrap().unwrap();
493
494        assert_eq!(plaintext, decrypted);
495    }
496
497    #[test]
498    fn test_list_keys() {
499        let manager = SecretsManager::new();
500
501        manager.set("key1", "value1").unwrap();
502        manager.set("key2", "value2").unwrap();
503        manager.set("key3", "value3").unwrap();
504
505        let keys = manager.list_keys();
506        assert_eq!(keys.len(), 3);
507        assert!(keys.contains(&"key1".to_string()));
508        assert!(keys.contains(&"key2".to_string()));
509        assert!(keys.contains(&"key3".to_string()));
510    }
511
512    #[test]
513    fn test_rotation_required() {
514        let config = SecretsManagerConfig {
515            audit_enabled: true,
516            rotation_interval_secs: 0, // 立即需要轮换
517            auto_rotation_check: true,
518        };
519        let manager = SecretsManager::with_config(config);
520
521        manager.set("api_key", "secret").unwrap();
522        // 等待一小段时间确保 elapsed > 0
523        std::thread::sleep(std::time::Duration::from_millis(10));
524        manager.get("api_key").unwrap(); // 触发轮换检查
525
526        let keys = manager.get_keys_requiring_rotation();
527        assert!(!keys.is_empty());
528        assert!(keys.contains(&"api_key".to_string()));
529    }
530
531    #[test]
532    fn test_load_from_env() {
533        std::env::set_var("TEST_SECRET_KEY", "test_env_value");
534
535        let manager = SecretsManager::new();
536        manager.load_from_env("my_key", "TEST_SECRET_KEY").unwrap();
537
538        let value = manager.get("my_key").unwrap();
539        assert_eq!(value, Some("test_env_value".to_string()));
540
541        std::env::remove_var("TEST_SECRET_KEY");
542    }
543
544    #[test]
545    fn test_load_from_missing_env() {
546        let manager = SecretsManager::new();
547        let result = manager.load_from_env("key", "NONEXISTENT_ENV_VAR_12345");
548
549        assert!(result.is_err());
550        assert!(matches!(result.unwrap_err(), SecretsError::EnvNotSet(_)));
551    }
552
553    #[test]
554    fn test_clear_audit_log() {
555        let manager = SecretsManager::new();
556
557        manager.set("key", "value").unwrap();
558        assert!(!manager.get_audit_log().is_empty());
559
560        manager.clear_audit_log();
561        assert!(manager.get_audit_log().is_empty());
562    }
563
564    #[test]
565    fn test_delete_nonexistent() {
566        let manager = SecretsManager::new();
567        let result = manager.delete("nonexistent");
568
569        assert!(result.is_err());
570        assert!(matches!(result.unwrap_err(), SecretsError::NotFound(_)));
571    }
572
573    #[test]
574    fn test_rotate_nonexistent() {
575        let manager = SecretsManager::new();
576        let result = manager.rotate("nonexistent", "new_value");
577
578        assert!(result.is_err());
579        assert!(matches!(result.unwrap_err(), SecretsError::NotFound(_)));
580    }
581
582    // ========== 边界条件测试 ==========
583
584    #[test]
585    fn test_empty_key_name() {
586        let manager = SecretsManager::new();
587        // 空密钥名应该被接受(边界情况)
588        let result = manager.set("", "value");
589        assert!(result.is_ok());
590        assert!(manager.contains(""));
591    }
592
593    #[test]
594    fn test_empty_secret_value() {
595        let manager = SecretsManager::new();
596        // 空密钥值应该被接受
597        let result = manager.set("key", "");
598        assert!(result.is_ok());
599        let value = manager.get("key").unwrap();
600        assert_eq!(value, Some("".to_string()));
601    }
602
603    #[test]
604    fn test_long_key_name() {
605        let manager = SecretsManager::new();
606        let long_key = "a".repeat(10000);
607        let result = manager.set(&long_key, "value");
608        assert!(result.is_ok());
609        assert!(manager.contains(&long_key));
610    }
611
612    #[test]
613    fn test_long_secret_value() {
614        let manager = SecretsManager::new();
615        let long_value = "x".repeat(1_000_000); // 1MB
616        let result = manager.set("big_secret", &long_value);
617        assert!(result.is_ok());
618        let value = manager.get("big_secret").unwrap();
619        assert_eq!(value, Some(long_value));
620    }
621
622    #[test]
623    fn test_unicode_key_and_value() {
624        let manager = SecretsManager::new();
625        let unicode_key = "密钥_🔑_key";
626        let unicode_value = "秘密值_🔐_value";
627
628        manager.set(unicode_key, unicode_value).unwrap();
629        let value = manager.get(unicode_key).unwrap();
630        assert_eq!(value, Some(unicode_value.to_string()));
631    }
632
633    #[test]
634    fn test_special_characters_in_value() {
635        let manager = SecretsManager::new();
636        let special_value = "!@#$%^&*()_+-=[]{}|;':\",./<>?\n\t\r\\";
637
638        manager.set("special", special_value).unwrap();
639        let value = manager.get("special").unwrap();
640        assert_eq!(value, Some(special_value.to_string()));
641    }
642
643    #[test]
644    fn test_binary_like_value() {
645        let manager = SecretsManager::new();
646        // 包含所有字节值的字符串
647        let binary_value: String = (0u8..=255).map(|b| b as char).collect();
648
649        manager.set("binary", &binary_value).unwrap();
650        let value = manager.get("binary").unwrap();
651        assert_eq!(value, Some(binary_value));
652    }
653
654    #[test]
655    fn test_get_metadata_for_nonexistent_key() {
656        let manager = SecretsManager::new();
657        let meta = manager.get_metadata("nonexistent");
658        assert!(meta.is_none());
659    }
660
661    #[test]
662    fn test_audit_disabled() {
663        let config = SecretsManagerConfig {
664            audit_enabled: false,
665            rotation_interval_secs: 86400,
666            auto_rotation_check: false,
667        };
668        let manager = SecretsManager::with_config(config);
669
670        manager.set("key", "value").unwrap();
671        manager.get("key").unwrap();
672
673        // 审计日志应该为空
674        assert!(manager.get_audit_log().is_empty());
675    }
676
677    #[test]
678    fn test_auto_rotation_disabled() {
679        let config = SecretsManagerConfig {
680            audit_enabled: true,
681            rotation_interval_secs: 0,
682            auto_rotation_check: false,
683        };
684        let manager = SecretsManager::with_config(config);
685
686        manager.set("key", "value").unwrap();
687        std::thread::sleep(std::time::Duration::from_millis(10));
688        manager.get("key").unwrap();
689
690        // 即使过了轮换周期,也不应该标记需要轮换
691        let keys = manager.get_keys_requiring_rotation();
692        assert!(keys.is_empty());
693    }
694
695    #[test]
696    fn test_overwrite_existing_key() {
697        let manager = SecretsManager::new();
698
699        manager.set("key", "value1").unwrap();
700        assert_eq!(manager.get("key").unwrap(), Some("value1".to_string()));
701
702        // 覆盖现有密钥
703        manager.set("key", "value2").unwrap();
704        assert_eq!(manager.get("key").unwrap(), Some("value2".to_string()));
705
706        // 应该只有一条 Set 记录(最后一次)
707        let log = manager.get_audit_log();
708        let set_count = log
709            .iter()
710            .filter(|e| matches!(e.action, AuditAction::Set))
711            .count();
712        assert_eq!(set_count, 2);
713    }
714
715    #[test]
716    fn test_concurrent_access() {
717        use std::sync::Arc;
718        use std::thread;
719
720        let manager = Arc::new(SecretsManager::new());
721        let mut handles = vec![];
722
723        for i in 0..10 {
724            let m = Arc::clone(&manager);
725            handles.push(thread::spawn(move || {
726                let key = format!("key_{}", i);
727                m.set(&key, &format!("value_{}", i)).unwrap();
728                m.get(&key).unwrap();
729            }));
730        }
731
732        for h in handles {
733            h.join().unwrap();
734        }
735
736        // 所有密钥应该都被正确设置
737        for i in 0..10 {
738            let key = format!("key_{}", i);
739            assert!(manager.contains(&key));
740        }
741    }
742
743    #[test]
744    fn test_encrypted_storage_invalid_utf8() {
745        let storage = EncryptedStorage::new();
746        storage.set("key", "value");
747
748        // 获取加密后的数据
749        let encrypted = storage
750            .encrypted_secrets
751            .read()
752            .get("key")
753            .cloned()
754            .unwrap();
755
756        // 手动破坏加密数据使其无法解码为 UTF-8
757        let mut corrupted = encrypted.clone();
758        // 插入无效 UTF-8 序列
759        for i in 0..corrupted.len().min(10) {
760            corrupted[i] = 0x80 | (corrupted[i] & 0x7F);
761        }
762
763        // 重新设置破坏的数据
764        storage
765            .encrypted_secrets
766            .write()
767            .insert("key".to_string(), corrupted);
768
769        // 解密应该失败并返回 EncryptionError
770        let result = storage.get("key");
771        // 由于我们可能正好产生了有效的 UTF-8,这里只是验证不会 panic
772        assert!(result.is_ok() || result.is_err());
773    }
774}