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 if self.identity.is_none() {
128 legacy::clear_for_provider(&self.provider)
129 .context("failed to remove legacy plaintext credential after secure save")?;
130 }
131 Ok(())
132 }
133
134 pub fn load(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
136 if let Some(key) = self.storage.load_with_mode(mode)? {
137 let key = key.trim();
138 return Ok((!key.is_empty()).then(|| key.to_owned()));
139 }
140
141 if self.identity.is_none() {
142 self.load_legacy_auth_json(mode)
143 } else {
144 Ok(None)
145 }
146 }
147
148 pub fn load_with_legacy_fallback(
151 &self,
152 mode: AuthCredentialsStoreMode,
153 allow_legacy: bool,
154 ) -> Result<Option<String>> {
155 if let Some(key) = self.load(mode)? {
156 return Ok(Some(key));
157 }
158 if !allow_legacy || self.identity.is_none() {
159 return Ok(None);
160 }
161
162 let legacy_storage = Self::new(&self.provider);
163 let Some(key) = legacy_storage.load(mode)? else {
164 return Ok(None);
165 };
166
167 self.store_value(&key, mode)
168 .context("failed to migrate provider-only credential to key-scoped storage")?;
169 legacy_storage
170 .clear(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.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 let path = crate::storage_paths::legacy_auth_storage_path().ok();
221 if let Some(p) = path {
222 legacy::delete_file(&p).context("failed to remove migrated plaintext auth file")?;
223 }
224
225 tracing::warn!(
226 "Migrated legacy plaintext auth.json entry for provider '{}' into secure storage",
227 self.provider
228 );
229 self.load(mode)
230 }
231}
232
233pub fn migrate_custom_api_keys(
237 custom_api_keys: &BTreeMap<String, String>,
238 mode: AuthCredentialsStoreMode,
239) -> Result<BTreeMap<String, bool>> {
240 let mut results = BTreeMap::new();
241
242 for (provider, api_key) in custom_api_keys {
243 let storage = CustomApiKeyStorage::new(provider);
244 match storage.store(api_key, mode) {
245 Ok(()) => {
246 tracing::info!("Migrated API key for provider '{provider}' to secure storage");
247 let _ignored = results.insert(provider.clone(), true);
248 }
249 Err(e) => {
250 tracing::warn!("Failed to migrate API key for provider '{provider}': {e}");
251 let _ignored = results.insert(provider.clone(), false);
252 }
253 }
254 }
255
256 Ok(results)
257}
258
259pub fn load_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<BTreeMap<String, String>> {
263 let mut api_keys = BTreeMap::new();
264
265 for provider in providers {
266 let storage = CustomApiKeyStorage::new(provider);
267 if let Some(key) = storage.load(mode)? {
268 drop(api_keys.insert(provider.clone(), key));
269 }
270 }
271
272 Ok(api_keys)
273}
274
275pub fn clear_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<()> {
277 for provider in providers {
278 let storage = CustomApiKeyStorage::new(provider);
279 if let Err(e) = storage.clear(mode) {
280 tracing::warn!("Failed to clear API key for provider '{provider}': {e}");
281 }
282 }
283 Ok(())
284}
285
286#[cfg(test)]
287mod tests {
288 use super::*;
289 use serial_test::serial;
290 use std::path::PathBuf;
291 use tempfile::TempDir;
292
293 struct TestAuthDirGuard {
294 temp_dir: Option<TempDir>,
295 previous: Option<PathBuf>,
296 }
297
298 impl TestAuthDirGuard {
299 fn new() -> Self {
300 let temp_dir = TempDir::new().expect("create temp auth dir");
301 let previous = crate::storage_paths::auth_storage_dir_override_for_tests().expect("read auth dir override");
302 crate::storage_paths::set_auth_storage_dir_override_for_tests(Some(temp_dir.path().to_path_buf()))
303 .expect("set temp auth dir override");
304 Self { temp_dir: Some(temp_dir), previous }
305 }
306 }
307
308 impl Drop for TestAuthDirGuard {
309 fn drop(&mut self) {
310 crate::storage_paths::set_auth_storage_dir_override_for_tests(self.previous.clone())
311 .expect("restore auth dir override");
312 if let Some(temp_dir) = self.temp_dir.take() {
313 temp_dir.close().expect("remove temp auth dir");
314 }
315 }
316 }
317
318 #[test]
319 fn credential_identity_normalizes_provider_and_key_name() {
320 let identity = CredentialIdentity::new(" MiMo ", "mimo_token_plan_key").expect("valid identity");
321
322 assert_eq!(identity.provider(), "mimo");
323 assert_eq!(identity.key_name(), "MIMO_TOKEN_PLAN_KEY");
324 assert!(identity.uses_default_key_name("MIMO_TOKEN_PLAN_KEY"));
325 assert!(!identity.uses_default_key_name("MIMO_API_KEY"));
326 }
327
328 #[test]
329 fn credential_identity_rejects_invalid_names() {
330 assert!(CredentialIdentity::new("my corp", "MYCORP_API_KEY").is_err());
331 assert!(CredentialIdentity::new("mycorp", "MY-CORP-API-KEY").is_err());
332 assert!(CredentialIdentity::new("mycorp", "1MYCORP_API_KEY").is_err());
333 }
334
335 #[test]
336 #[serial]
337 #[cfg(unix)]
338 fn file_api_key_storage_round_trips_with_private_permissions() {
339 use std::fs;
340 use std::os::unix::fs::PermissionsExt;
341
342 let guard = TestAuthDirGuard::new();
343 let storage = CustomApiKeyStorage::new("stepfun");
344 storage
345 .store("test-stepfun-key", AuthCredentialsStoreMode::File)
346 .expect("store API key");
347 assert_eq!(
348 storage.load(AuthCredentialsStoreMode::File).expect("load API key").as_deref(),
349 Some("test-stepfun-key")
350 );
351
352 let auth_dir = guard.temp_dir.as_ref().expect("test auth dir").path();
353 assert_eq!(fs::metadata(auth_dir).expect("auth dir metadata").permissions().mode() & 0o777, 0o700);
354 let credential_file = fs::read_dir(auth_dir)
355 .expect("read auth dir")
356 .map(|entry| entry.expect("credential entry").path())
357 .find(|path| path.extension().and_then(|extension| extension.to_str()) == Some("json"))
358 .expect("credential file");
359 assert_eq!(fs::metadata(credential_file).expect("credential metadata").permissions().mode() & 0o777, 0o600);
360 }
361
362 #[test]
363 #[serial]
364 fn keyring_mode_falls_back_to_encrypted_file_when_keyring_is_unavailable() {
365 let _guard = TestAuthDirGuard::new();
366 let storage = CustomApiKeyStorage::new("stepfun");
367
368 storage
369 .store("test-stepfun-key", AuthCredentialsStoreMode::Keyring)
370 .expect("keyring mode should fall back to encrypted file storage");
371 assert_eq!(
372 storage
373 .load(AuthCredentialsStoreMode::Keyring)
374 .expect("load API key")
375 .as_deref(),
376 Some("test-stepfun-key")
377 );
378 }
379
380 #[test]
381 #[serial]
382 fn key_scoped_storage_keeps_provider_profiles_isolated() {
383 let _guard = TestAuthDirGuard::new();
384 let payg = CustomApiKeyStorage::for_provider_key(" MiMo ", "mimo_api_key").expect("payg storage");
385 let token_plan =
386 CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
387
388 payg.store("sk-payg", AuthCredentialsStoreMode::File).expect("store payg");
389 token_plan
390 .store("tp-token-plan", AuthCredentialsStoreMode::File)
391 .expect("store token plan");
392
393 assert_eq!(payg.load(AuthCredentialsStoreMode::File).expect("load payg").as_deref(), Some("sk-payg"));
394 assert_eq!(
395 token_plan
396 .load(AuthCredentialsStoreMode::File)
397 .expect("load token plan")
398 .as_deref(),
399 Some("tp-token-plan")
400 );
401 }
402
403 #[test]
404 #[serial]
405 fn default_identity_lazily_migrates_provider_only_storage() {
406 let _guard = TestAuthDirGuard::new();
407 let legacy = CustomApiKeyStorage::new("mimo");
408 let target = CustomApiKeyStorage::for_provider_key("mimo", "MIMO_API_KEY").expect("target storage");
409
410 legacy
411 .store("legacy-key", AuthCredentialsStoreMode::File)
412 .expect("store legacy key");
413 assert_eq!(
414 target
415 .load_with_legacy_fallback(AuthCredentialsStoreMode::File, true)
416 .expect("migrate legacy key")
417 .as_deref(),
418 Some("legacy-key")
419 );
420 assert_eq!(legacy.load(AuthCredentialsStoreMode::File).expect("legacy should be cleared"), None);
421 assert_eq!(
422 target
423 .load(AuthCredentialsStoreMode::File)
424 .expect("load migrated key")
425 .as_deref(),
426 Some("legacy-key")
427 );
428 }
429
430 #[test]
431 #[serial]
432 fn non_default_identity_does_not_reuse_provider_only_storage() {
433 let _guard = TestAuthDirGuard::new();
434 let legacy = CustomApiKeyStorage::new("mimo");
435 let token_plan =
436 CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
437
438 legacy
439 .store("legacy-payg", AuthCredentialsStoreMode::File)
440 .expect("store legacy key");
441 assert_eq!(
442 token_plan
443 .load_with_legacy_fallback(AuthCredentialsStoreMode::File, false)
444 .expect("load token-plan key"),
445 None
446 );
447 assert_eq!(
448 legacy.load(AuthCredentialsStoreMode::File).expect("legacy remains"),
449 Some("legacy-payg".to_string())
450 );
451 }
452
453 #[test]
454 #[serial]
455 fn clearing_one_identity_does_not_clear_another() {
456 let _guard = TestAuthDirGuard::new();
457 let payg = CustomApiKeyStorage::for_provider_key("mimo", "MIMO_API_KEY").expect("payg storage");
458 let token_plan =
459 CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
460
461 payg.store("sk-payg", AuthCredentialsStoreMode::File).expect("store payg");
462 token_plan
463 .store("tp-token-plan", AuthCredentialsStoreMode::File)
464 .expect("store token plan");
465 payg.clear_with_legacy_fallback(AuthCredentialsStoreMode::File, true)
466 .expect("clear payg");
467
468 assert_eq!(payg.load(AuthCredentialsStoreMode::File).expect("payg cleared"), None);
469 assert_eq!(
470 token_plan
471 .load(AuthCredentialsStoreMode::File)
472 .expect("token plan remains")
473 .as_deref(),
474 Some("tp-token-plan")
475 );
476 }
477}