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