key_token/
execution_plan.rs

1use serde::{Deserialize, Serialize};
2use std::fmt;
3use uuid::Uuid;
4
5/// 敏感字符串：用于保护 API Key 等敏感信息
6///
7/// - 序列化时会隐藏内容（显示为 ***REDACTED***）
8/// - Debug 时隐藏内容
9/// - Display 时隐藏内容
10#[derive(Clone, Default, Deserialize, PartialEq)]
11pub struct SensitiveString(String);
12
13impl SensitiveString {
14    /// 创建新的敏感字符串
15    pub fn new(s: impl Into<String>) -> Self {
16        Self(s.into())
17    }
18
19    /// 获取内部值（谨慎使用）
20    pub fn expose(&self) -> &str {
21        &self.0
22    }
23
24    /// 检查是否为空
25    pub fn is_empty(&self) -> bool {
26        self.0.is_empty()
27    }
28
29    /// 获取长度
30    pub fn len(&self) -> usize {
31        self.0.len()
32    }
33}
34
35impl fmt::Debug for SensitiveString {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        write!(f, "***REDACTED***")
38    }
39}
40
41impl fmt::Display for SensitiveString {
42    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43        write!(f, "***REDACTED***")
44    }
45}
46
47impl Serialize for SensitiveString {
48    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
49    where
50        S: serde::Serializer,
51    {
52        serializer.serialize_str("***REDACTED***")
53    }
54}
55
56impl From<String> for SensitiveString {
57    fn from(s: String) -> Self {
58        Self(s)
59    }
60}
61
62impl From<&str> for SensitiveString {
63    fn from(s: &str) -> Self {
64        Self(s.to_string())
65    }
66}
67
68/// 执行计划：包含主目标和回退链
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct ExecutionPlan {
71    pub primary: ExecutionTarget,
72    pub fallback_chain: Vec<ExecutionTarget>,
73}
74
75impl ExecutionPlan {
76    pub fn new(primary: ExecutionTarget) -> Self {
77        Self {
78            primary,
79            fallback_chain: Vec::new(),
80        }
81    }
82
83    pub fn with_fallback(mut self, fallback: ExecutionTarget) -> Self {
84        self.fallback_chain.push(fallback);
85        self
86    }
87
88    pub fn with_fallbacks(mut self, fallbacks: Vec<ExecutionTarget>) -> Self {
89        self.fallback_chain.extend(fallbacks);
90        self
91    }
92
93    /// 获取所有执行目标（主目标 + 回退链）
94    pub fn all_targets(&self) -> impl Iterator<Item = &ExecutionTarget> {
95        std::iter::once(&self.primary).chain(self.fallback_chain.iter())
96    }
97}
98
99/// 执行目标：指定具体的 provider 和账号
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct ExecutionTarget {
102    pub provider: String,
103    pub account_id: Uuid,
104    pub endpoint: String,
105    pub upstream_api_key: SensitiveString, // 已解密的上游 Provider API Key（敏感信息自动隐藏）
106}
107
108impl ExecutionTarget {
109    pub fn new(
110        provider: impl Into<String>,
111        account_id: Uuid,
112        endpoint: impl Into<String>,
113        upstream_api_key: impl Into<SensitiveString>,
114    ) -> Self {
115        Self {
116            provider: provider.into(),
117            account_id,
118            endpoint: endpoint.into(),
119            upstream_api_key: upstream_api_key.into(),
120        }
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn test_sensitive_string_creation() {
130        let secret = SensitiveString::new("my-secret-key");
131        assert_eq!(secret.expose(), "my-secret-key");
132        assert_eq!(secret.len(), 13);
133        assert!(!secret.is_empty());
134    }
135
136    #[test]
137    fn test_sensitive_string_debug() {
138        let secret = SensitiveString::new("my-secret-key");
139        let debug_str = format!("{:?}", secret);
140        assert_eq!(debug_str, "***REDACTED***");
141        assert!(!debug_str.contains("my-secret-key"));
142    }
143
144    #[test]
145    fn test_sensitive_string_display() {
146        let secret = SensitiveString::new("my-secret-key");
147        let display_str = format!("{}", secret);
148        assert_eq!(display_str, "***REDACTED***");
149        assert!(!display_str.contains("my-secret-key"));
150    }
151
152    #[test]
153    fn test_sensitive_string_serialize() {
154        let secret = SensitiveString::new("my-secret-key");
155        let json = serde_json::to_string(&secret).unwrap();
156        assert_eq!(json, "\"***REDACTED***\"");
157        assert!(!json.contains("my-secret-key"));
158    }
159
160    #[test]
161    fn test_sensitive_string_deserialize() {
162        let json = "\"my-secret-key\"";
163        let secret: SensitiveString = serde_json::from_str(json).unwrap();
164        assert_eq!(secret.expose(), "my-secret-key");
165    }
166
167    #[test]
168    fn test_sensitive_string_default() {
169        let secret = SensitiveString::default();
170        assert!(secret.is_empty());
171        assert_eq!(secret.len(), 0);
172    }
173
174    #[test]
175    fn test_sensitive_string_partial_eq() {
176        let s1 = SensitiveString::new("key");
177        let s2 = SensitiveString::new("key");
178        let s3 = SensitiveString::new("different");
179        assert_eq!(s1, s2);
180        assert_ne!(s1, s3);
181    }
182
183    #[test]
184    fn test_sensitive_string_partial_eq_no_leak() {
185        // 验证 PartialEq 比较不会泄露敏感信息
186        let s1 = SensitiveString::new("secret-key-123");
187        let s2 = SensitiveString::new("secret-key-123");
188        let s3 = SensitiveString::new("different-key");
189
190        // 相等比较应该正常工作
191        assert!(s1 == s2);
192        assert!(s1 != s3);
193
194        // 验证 Debug 输出不包含原始值
195        let debug = format!("{:?}", s1);
196        assert!(!debug.contains("secret-key-123"));
197
198        // 验证 Display 输出不包含原始值
199        let display = format!("{}", s1);
200        assert!(!display.contains("secret-key-123"));
201
202        // 只有通过 expose() 才能获取原始值
203        assert_eq!(s1.expose(), "secret-key-123");
204    }
205
206    #[test]
207    fn test_sensitive_string_from_str() {
208        let s1: SensitiveString = "test-key".into();
209        let s2 = SensitiveString::new("test-key");
210        assert_eq!(s1, s2);
211        assert_eq!(s1.expose(), "test-key");
212    }
213
214    #[test]
215    fn test_sensitive_string_from_string() {
216        let s1: SensitiveString = String::from("test-key").into();
217        let s2 = SensitiveString::new("test-key");
218        assert_eq!(s1, s2);
219    }
220
221    #[test]
222    fn test_execution_plan_new() {
223        let target = ExecutionTarget::new(
224            "openai",
225            Uuid::new_v4(),
226            "https://api.openai.com",
227            "sk-test-key",
228        );
229        let plan = ExecutionPlan::new(target);
230        assert_eq!(plan.primary.provider, "openai");
231        assert!(plan.fallback_chain.is_empty());
232    }
233
234    #[test]
235    fn test_execution_plan_with_fallback() {
236        let primary = ExecutionTarget::new(
237            "openai",
238            Uuid::new_v4(),
239            "https://api.openai.com",
240            "sk-primary-key",
241        );
242        let fallback = ExecutionTarget::new(
243            "claude",
244            Uuid::new_v4(),
245            "https://api.anthropic.com",
246            "sk-fallback-key",
247        );
248        let plan = ExecutionPlan::new(primary).with_fallback(fallback);
249        assert_eq!(plan.fallback_chain.len(), 1);
250        assert_eq!(plan.fallback_chain[0].provider, "claude");
251    }
252
253    #[test]
254    fn test_execution_plan_all_targets() {
255        let primary = ExecutionTarget::new(
256            "openai",
257            Uuid::new_v4(),
258            "https://api.openai.com",
259            "sk-primary-key",
260        );
261        let fallback1 = ExecutionTarget::new(
262            "claude",
263            Uuid::new_v4(),
264            "https://api.anthropic.com",
265            "sk-fallback1-key",
266        );
267        let fallback2 = ExecutionTarget::new(
268            "gemini",
269            Uuid::new_v4(),
270            "https://api.gemini.com",
271            "sk-fallback2-key",
272        );
273        let plan = ExecutionPlan::new(primary)
274            .with_fallback(fallback1)
275            .with_fallback(fallback2);
276
277        let targets: Vec<_> = plan.all_targets().collect();
278        assert_eq!(targets.len(), 3);
279        assert_eq!(targets[0].provider, "openai");
280        assert_eq!(targets[1].provider, "claude");
281        assert_eq!(targets[2].provider, "gemini");
282    }
283
284    #[test]
285    fn test_execution_target_api_key_hidden_in_debug() {
286        let target = ExecutionTarget::new(
287            "openai",
288            Uuid::new_v4(),
289            "https://api.openai.com",
290            "sk-secret-key",
291        );
292        let debug_str = format!("{:?}", target);
293        assert!(debug_str.contains("***REDACTED***"));
294        assert!(!debug_str.contains("sk-secret-key"));
295    }
296
297    #[test]
298    fn test_execution_target_api_key_hidden_in_serialize() {
299        let target = ExecutionTarget::new(
300            "openai",
301            Uuid::new_v4(),
302            "https://api.openai.com",
303            "sk-secret-key",
304        );
305        let json = serde_json::to_string(&target).unwrap();
306        assert!(json.contains("***REDACTED***"));
307        assert!(!json.contains("sk-secret-key"));
308    }
309
310    #[test]
311    fn test_execution_target_api_key_expose() {
312        let target = ExecutionTarget::new(
313            "openai",
314            Uuid::new_v4(),
315            "https://api.openai.com",
316            "sk-secret-key",
317        );
318        // expose() 方法可以获取原始值（用于实际请求）
319        assert_eq!(target.upstream_api_key.expose(), "sk-secret-key");
320    }
321}
key_token/execution_plan.rs

key_token/
execution_plan.rs
