1pub(crate) mod encryption;
14pub(crate) mod keyring;
15pub(crate) mod legacy;
16pub(crate) mod mode;
17pub(crate) mod storage;
18
19pub use mode::AuthCredentialsStoreMode;
20pub use storage::CredentialStorage;
21
22use std::collections::BTreeMap;
23
24use anyhow::{Context, Result, bail};
25
26#[derive(Debug, Clone, PartialEq, Eq, Hash)]
33pub struct CredentialIdentity {
34 provider: String,
35 key_name: String,
36}
37
38impl CredentialIdentity {
39 pub fn new(provider: &str, key_name: &str) -> Result<Self> {
42 let provider = provider.trim().to_ascii_lowercase();
43 if provider.is_empty() || !provider.chars().all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_') {
44 bail!("credential provider must contain only letters, digits, '-' or '_'");
45 }
46
47 let key_name = key_name.trim().to_ascii_uppercase();
48 let mut chars = key_name.chars();
49 let valid_start = chars.next().is_some_and(|ch| ch.is_ascii_uppercase() || ch == '_');
50 if !valid_start || !chars.all(|ch| ch.is_ascii_uppercase() || ch.is_ascii_digit() || ch == '_') {
51 bail!("credential key name must be a valid environment variable name");
52 }
53
54 Ok(Self { provider, key_name })
55 }
56
57 pub fn provider(&self) -> &str {
59 &self.provider
60 }
61
62 pub fn key_name(&self) -> &str {
64 &self.key_name
65 }
66
67 pub fn uses_default_key_name(&self, default_key_name: &str) -> bool {
69 self.key_name.eq_ignore_ascii_case(default_key_name.trim())
70 }
71}
72
73pub struct CustomApiKeyStorage {
78 provider: String,
79 identity: Option<CredentialIdentity>,
80 storage: CredentialStorage,
81}
82
83impl CustomApiKeyStorage {
84 pub fn new(provider: &str) -> Self {
90 let normalized_provider = provider.trim().to_lowercase();
91 Self {
92 provider: normalized_provider.clone(),
93 identity: None,
94 storage: CredentialStorage::new("vtcode", format!("api_key_{normalized_provider}")),
95 }
96 }
97
98 pub fn for_provider_key(provider: &str, key_name: &str) -> Result<Self> {
100 Self::for_identity(CredentialIdentity::new(provider, key_name)?)
101 }
102
103 pub fn for_identity(identity: CredentialIdentity) -> Result<Self> {
105 let provider = identity.provider().to_owned();
106 let user = format!("api_key_{}_{}", identity.provider(), identity.key_name());
107 Ok(Self {
108 provider,
109 identity: Some(identity),
110 storage: CredentialStorage::new("vtcode", user),
111 })
112 }
113
114 pub fn identity(&self) -> Option<&CredentialIdentity> {
117 self.identity.as_ref()
118 }
119
120 pub fn store(&self, api_key: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
122 let api_key = api_key.trim();
123 if api_key.is_empty() {
124 bail!("API key cannot be empty");
125 }
126 self.store_value(api_key, mode)?;
127 Ok(())
128 }
129
130 pub fn load(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
132 if let Some(key) = self.storage.load_with_mode(mode)? {
133 let key = key.trim();
134 return Ok((!key.is_empty()).then(|| key.to_owned()));
135 }
136
137 if self.identity.is_none() {
138 self.load_legacy_auth_json(mode)
139 } else {
140 Ok(None)
141 }
142 }
143
144 pub fn load_with_legacy_fallback(
147 &self,
148 mode: AuthCredentialsStoreMode,
149 allow_legacy: bool,
150 ) -> Result<Option<String>> {
151 if let Some(key) = self.load(mode)? {
152 return Ok(Some(key));
153 }
154 if !allow_legacy || self.identity.is_none() {
155 return Ok(None);
156 }
157
158 let legacy_storage = Self::new(&self.provider);
159 let Some(key) = legacy_storage.load(mode)? else {
160 return Ok(None);
161 };
162
163 self.store_value(&key, mode)
164 .context("failed to migrate provider-only credential to key-scoped storage")?;
165 legacy_storage
169 .storage
170 .clear_with_mode(mode)
171 .context("failed to remove provider-only credential after migration")?;
172 self.load(mode)
173 }
174
175 pub fn clear(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
177 self.storage.clear_with_mode(mode)?;
178 if self.identity.is_none() {
179 legacy::clear_for_provider(&self.provider).context("failed to remove legacy plaintext credential")?;
180 }
181 Ok(())
182 }
183
184 pub fn clear_with_legacy_fallback(&self, mode: AuthCredentialsStoreMode, clear_legacy: bool) -> Result<()> {
187 self.clear(mode)?;
188 if clear_legacy && self.identity.is_some() {
189 Self::new(&self.provider).clear(mode)?;
190 }
191 Ok(())
192 }
193
194 fn store_value(&self, api_key: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
195 self.storage.store_with_mode(api_key, mode)?;
196 let persisted = self
197 .storage
198 .load_with_mode(mode)
199 .context("failed to verify persisted API key")?;
200 if persisted.as_deref().map(str::trim) != Some(api_key) {
201 bail!("secure storage did not return the API key after saving");
202 }
203 Ok(())
204 }
205
206 fn load_legacy_auth_json(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
207 let Some(legacy_entry) = legacy::load_for_provider(&self.provider)? else {
208 return Ok(None);
209 };
210
211 if let Err(err) = self.store(&legacy_entry.credentials.api_key, mode) {
212 tracing::warn!(
213 "Failed to migrate legacy plaintext auth.json entry for provider '{}' into secure storage: {}",
214 self.provider,
215 err
216 );
217 return Err(err).context("failed to migrate legacy API key into secure storage");
218 }
219
220 tracing::warn!(self.provider);
221 self.load(mode)
222 }
223}
224
225pub fn migrate_custom_api_keys(
229 custom_api_keys: &BTreeMap<String, String>,
230 mode: AuthCredentialsStoreMode,
231) -> Result<BTreeMap<String, bool>> {
232 let mut results = BTreeMap::new();
233
234 for (provider, api_key) in custom_api_keys {
235 let storage = CustomApiKeyStorage::new(provider);
236 match storage.store(api_key, mode) {
237 Ok(()) => {
238 tracing::info!("Migrated API key for provider '{provider}' to secure storage");
239 let _ignored = results.insert(provider.clone(), true);
240 }
241 Err(e) => {
242 tracing::warn!("Failed to migrate API key for provider '{provider}': {e}");
243 let _ignored = results.insert(provider.clone(), false);
244 }
245 }
246 }
247
248 Ok(results)
249}
250
251pub fn load_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<BTreeMap<String, String>> {
255 let mut api_keys = BTreeMap::new();
256
257 for provider in providers {
258 let storage = CustomApiKeyStorage::new(provider);
259 if let Some(key) = storage.load(mode)? {
260 drop(api_keys.insert(provider.clone(), key));
261 }
262 }
263
264 Ok(api_keys)
265}
266
267pub fn clear_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<()> {
269 for provider in providers {
270 let storage = CustomApiKeyStorage::new(provider);
271 if let Err(e) = storage.clear(mode) {
272 tracing::warn!("Failed to clear API key for provider '{provider}': {e}");
273 }
274 }
275 Ok(())
276}
277
278#[cfg(test)]
279mod tests {
280 use super::*;
281 use serial_test::serial;
282 use std::path::PathBuf;
283 use tempfile::TempDir;
284
285 struct TestAuthDirGuard {
286 temp_dir: Option<TempDir>,
287 previous: Option<PathBuf>,
288 }
289
290 impl TestAuthDirGuard {
291 fn new() -> Self {
292 let temp_dir = TempDir::new().expect("create temp auth dir");
293 let previous = crate::storage_paths::auth_storage_dir_override_for_tests().expect("read auth dir override");
294 crate::storage_paths::set_auth_storage_dir_override_for_tests(Some(temp_dir.path().to_path_buf()))
295 .expect("set temp auth dir override");
296 Self { temp_dir: Some(temp_dir), previous }
297 }
298 }
299
300 impl Drop for TestAuthDirGuard {
301 fn drop(&mut self) {
302 crate::storage_paths::set_auth_storage_dir_override_for_tests(self.previous.clone())
303 .expect("restore auth dir override");
304 if let Some(temp_dir) = self.temp_dir.take() {
305 temp_dir.close().expect("remove temp auth dir");
306 }
307 }
308 }
309
310 #[test]
311 fn credential_identity_normalizes_provider_and_key_name() {
312 let identity = CredentialIdentity::new(" MiMo ", "mimo_token_plan_key").expect("valid identity");
313
314 assert_eq!(identity.provider(), "mimo");
315 assert_eq!(identity.key_name(), "MIMO_TOKEN_PLAN_KEY");
316 assert!(identity.uses_default_key_name("MIMO_TOKEN_PLAN_KEY"));
317 assert!(!identity.uses_default_key_name("MIMO_API_KEY"));
318 }
319
320 #[test]
321 fn credential_identity_rejects_invalid_names() {
322 assert!(CredentialIdentity::new("my corp", "MYCORP_API_KEY").is_err());
323 assert!(CredentialIdentity::new("mycorp", "MY-CORP-API-KEY").is_err());
324 assert!(CredentialIdentity::new("mycorp", "1MYCORP_API_KEY").is_err());
325 }
326
327 #[test]
328 #[serial]
329 #[cfg(unix)]
330 fn file_api_key_storage_round_trips_with_private_permissions() {
331 use std::fs;
332 use std::os::unix::fs::PermissionsExt;
333
334 let guard = TestAuthDirGuard::new();
335 let storage = CustomApiKeyStorage::new("stepfun");
336 storage
337 .store("test-stepfun-key", AuthCredentialsStoreMode::File)
338 .expect("store API key");
339 assert_eq!(
340 storage.load(AuthCredentialsStoreMode::File).expect("load API key").as_deref(),
341 Some("test-stepfun-key")
342 );
343
344 let auth_dir = guard.temp_dir.as_ref().expect("test auth dir").path();
345 assert_eq!(fs::metadata(auth_dir).expect("auth dir metadata").permissions().mode() & 0o777, 0o700);
346 let credential_file = fs::read_dir(auth_dir)
347 .expect("read auth dir")
348 .map(|entry| entry.expect("credential entry").path())
349 .find(|path| path.extension().and_then(|extension| extension.to_str()) == Some("json"))
350 .expect("credential file");
351 assert_eq!(fs::metadata(credential_file).expect("credential metadata").permissions().mode() & 0o777, 0o600);
352 }
353
354 #[test]
355 #[serial]
356 fn keyring_mode_falls_back_to_encrypted_file_when_keyring_is_unavailable() {
357 let _guard = TestAuthDirGuard::new();
358 let storage = CustomApiKeyStorage::new("stepfun");
359
360 storage
361 .store("test-stepfun-key", AuthCredentialsStoreMode::Keyring)
362 .expect("keyring mode should fall back to encrypted file storage");
363 assert_eq!(
364 storage
365 .load(AuthCredentialsStoreMode::Keyring)
366 .expect("load API key")
367 .as_deref(),
368 Some("test-stepfun-key")
369 );
370 }
371
372 #[test]
373 #[serial]
374 fn key_scoped_storage_keeps_provider_profiles_isolated() {
375 let _guard = TestAuthDirGuard::new();
376 let payg = CustomApiKeyStorage::for_provider_key(" MiMo ", "mimo_api_key").expect("payg storage");
377 let token_plan =
378 CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
379
380 payg.store("sk-payg", AuthCredentialsStoreMode::File).expect("store payg");
381 token_plan
382 .store("tp-token-plan", AuthCredentialsStoreMode::File)
383 .expect("store token plan");
384
385 assert_eq!(payg.load(AuthCredentialsStoreMode::File).expect("load payg").as_deref(), Some("sk-payg"));
386 assert_eq!(
387 token_plan
388 .load(AuthCredentialsStoreMode::File)
389 .expect("load token plan")
390 .as_deref(),
391 Some("tp-token-plan")
392 );
393 }
394
395 #[test]
396 #[serial]
397 fn default_identity_lazily_migrates_provider_only_storage() {
398 let _guard = TestAuthDirGuard::new();
399 let legacy = CustomApiKeyStorage::new("mimo");
400 let target = CustomApiKeyStorage::for_provider_key("mimo", "MIMO_API_KEY").expect("target storage");
401
402 legacy
403 .store("legacy-key", AuthCredentialsStoreMode::File)
404 .expect("store legacy key");
405 assert_eq!(
406 target
407 .load_with_legacy_fallback(AuthCredentialsStoreMode::File, true)
408 .expect("migrate legacy key")
409 .as_deref(),
410 Some("legacy-key")
411 );
412 assert_eq!(legacy.load(AuthCredentialsStoreMode::File).expect("legacy should be cleared"), None);
413 assert_eq!(
414 target
415 .load(AuthCredentialsStoreMode::File)
416 .expect("load migrated key")
417 .as_deref(),
418 Some("legacy-key")
419 );
420 }
421
422 #[test]
423 #[serial]
424 fn non_default_identity_does_not_reuse_provider_only_storage() {
425 let _guard = TestAuthDirGuard::new();
426 let legacy = CustomApiKeyStorage::new("mimo");
427 let token_plan =
428 CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
429
430 legacy
431 .store("legacy-payg", AuthCredentialsStoreMode::File)
432 .expect("store legacy key");
433 assert_eq!(
434 token_plan
435 .load_with_legacy_fallback(AuthCredentialsStoreMode::File, false)
436 .expect("load token-plan key"),
437 None
438 );
439 assert_eq!(
440 legacy.load(AuthCredentialsStoreMode::File).expect("legacy remains"),
441 Some("legacy-payg".to_string())
442 );
443 }
444
445 #[test]
446 #[serial]
447 fn clearing_one_identity_does_not_clear_another() {
448 let _guard = TestAuthDirGuard::new();
449 let payg = CustomApiKeyStorage::for_provider_key("mimo", "MIMO_API_KEY").expect("payg storage");
450 let token_plan =
451 CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
452
453 payg.store("sk-payg", AuthCredentialsStoreMode::File).expect("store payg");
454 token_plan
455 .store("tp-token-plan", AuthCredentialsStoreMode::File)
456 .expect("store token plan");
457 payg.clear_with_legacy_fallback(AuthCredentialsStoreMode::File, true)
458 .expect("clear payg");
459
460 assert_eq!(payg.load(AuthCredentialsStoreMode::File).expect("payg cleared"), None);
461 assert_eq!(
462 token_plan
463 .load(AuthCredentialsStoreMode::File)
464 .expect("token plan remains")
465 .as_deref(),
466 Some("tp-token-plan")
467 );
468 }
469}