vtcode_auth/credentials/
mod.rs

1//! Credential storage — keyring and encrypted-file backends.
2//!
3//! # Module Structure
4//!
5//! | Submodule | Responsibility |
6//! |---|---|
7//! | mode | Backend selection enum (`Keyring` / `File` / `Auto`) |
8//! | keyring | OS keyring creation, liveness, disable detection |
9//! | encryption | AES-256-GCM encrypt/decrypt (pure, no IO) |
10//! | storage | `CredentialStorage` — orchestrates backends |
11//! | legacy | Legacy `auth.json` migration |
12
13pub(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/// Validated identity for a stored provider credential.
28///
29/// Provider names are normalized to lowercase and environment-variable names
30/// are normalized to uppercase. Keeping both dimensions in a value object
31/// prevents a credential for one provider profile from being reused for
32/// another profile that happens to share the provider name.
33#[derive(Debug, Clone, PartialEq, Eq, Hash)]
34pub struct CredentialIdentity {
35    provider: String,
36    key_name: String,
37}
38
39impl CredentialIdentity {
40    /// Create a credential identity from a provider name and environment
41    /// variable-style key name.
42    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    /// Normalized provider name.
59    pub fn provider(&self) -> &str {
60        &self.provider
61    }
62
63    /// Normalized credential key name.
64    pub fn key_name(&self) -> &str {
65        &self.key_name
66    }
67
68    /// Whether this identity uses the provider's default key name.
69    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
74/// Custom API Key storage for provider-specific keys.
75///
76/// Provides secure storage and retrieval of API keys for custom providers
77/// using the OS keyring or encrypted file storage.
78pub struct CustomApiKeyStorage {
79    provider: String,
80    identity: Option<CredentialIdentity>,
81    storage: CredentialStorage,
82}
83
84impl CustomApiKeyStorage {
85    /// Create a provider-only legacy storage handle.
86    ///
87    /// New callers should use [`Self::for_provider_key`]. This constructor is
88    /// retained so older integrations can still read and clear the legacy
89    /// `api_key_<provider>` entry.
90    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    /// Create storage scoped to a provider and credential key name.
100    pub fn for_provider_key(provider: &str, key_name: &str) -> Result<Self> {
101        Self::for_identity(CredentialIdentity::new(provider, key_name)?)
102    }
103
104    /// Create storage scoped to a validated credential identity.
105    pub fn for_identity(identity: CredentialIdentity) -> Result<Self> {
106        let provider = identity.provider().to_owned();
107        // The old underscore-separated account could represent multiple
108        // provider/key pairs. Do not read it without a verifiable identity.
109        let user = format!("api_key:v2:{}:{}", identity.provider(), identity.key_name());
110        Ok(Self {
111            provider,
112            identity: Some(identity),
113            storage: CredentialStorage::new("vtcode", user),
114        })
115    }
116
117    /// Return the identity for a key-scoped handle, or `None` for a legacy
118    /// provider-only handle.
119    pub fn identity(&self) -> Option<&CredentialIdentity> {
120        self.identity.as_ref()
121    }
122
123    /// Store an API key securely.
124    pub fn store(&self, api_key: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
125        let api_key = api_key.trim();
126        if api_key.is_empty() {
127            bail!("API key cannot be empty");
128        }
129        self.store_value(api_key, mode)?;
130        Ok(())
131    }
132
133    /// Retrieve a stored API key.
134    pub fn load(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
135        if let Some(key) = self.storage.load_with_mode(mode)? {
136            let key = key.trim();
137            return Ok((!key.is_empty()).then(|| key.to_owned()));
138        }
139
140        if self.identity.is_none() {
141            self.load_legacy_auth_json(mode)
142        } else {
143            Ok(None)
144        }
145    }
146
147    /// Load a key-scoped credential and, when explicitly allowed, lazily
148    /// migrate the provider-only legacy entry into this identity.
149    pub fn load_with_legacy_fallback(
150        &self,
151        mode: AuthCredentialsStoreMode,
152        allow_legacy: bool,
153    ) -> Result<Option<String>> {
154        if let Some(key) = self.load(mode)? {
155            return Ok(Some(key));
156        }
157        if !allow_legacy || self.identity.is_none() || legacy_provider_account_is_ambiguous(&self.provider) {
158            return Ok(None);
159        }
160
161        let legacy_storage = Self::new(&self.provider);
162        let Some(key) = legacy_storage.load(mode)? else {
163            return Ok(None);
164        };
165
166        self.store_value(&key, mode)
167            .context("failed to migrate provider-only credential to key-scoped storage")?;
168        // This is the obsolete provider-only secure-storage entry, not the
169        // rollback-safe plaintext file under the legacy directory. Remove the
170        // former after the key-scoped copy succeeds, but preserve the latter.
171        legacy_storage
172            .storage
173            .clear_with_mode(mode)
174            .context("failed to remove provider-only credential after migration")?;
175        self.load(mode)
176    }
177
178    /// Clear (delete) a stored API key.
179    pub fn clear(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
180        self.storage.clear_with_mode(mode)?;
181        if self.identity.is_none() {
182            legacy::clear_for_provider(&self.provider).context("failed to remove legacy plaintext credential")?;
183        }
184        Ok(())
185    }
186
187    /// Clear this key-scoped credential and optionally its provider-only
188    /// legacy fallback.
189    pub fn clear_with_legacy_fallback(&self, mode: AuthCredentialsStoreMode, clear_legacy: bool) -> Result<()> {
190        self.clear(mode)?;
191        if clear_legacy && self.identity.is_some() && !legacy_provider_account_is_ambiguous(&self.provider) {
192            Self::new(&self.provider).clear(mode)?;
193        }
194        Ok(())
195    }
196
197    fn store_value(&self, api_key: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
198        self.storage.store_with_mode(api_key, mode)?;
199        let persisted = self
200            .storage
201            .load_with_mode(mode)
202            .context("failed to verify persisted API key")?;
203        if persisted.as_deref().map(str::trim) != Some(api_key) {
204            bail!("secure storage did not return the API key after saving");
205        }
206        Ok(())
207    }
208
209    fn load_legacy_auth_json(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
210        let Some(legacy_entry) = legacy::load_for_provider(&self.provider)? else {
211            return Ok(None);
212        };
213
214        if let Err(err) = self.store(&legacy_entry.credentials.api_key, mode) {
215            tracing::warn!(
216                "Failed to migrate legacy plaintext auth.json entry for provider '{}' into secure storage: {}",
217                self.provider,
218                err
219            );
220            return Err(err).context("failed to migrate legacy API key into secure storage");
221        }
222
223        tracing::warn!(self.provider);
224        self.load(mode)
225    }
226}
227
228// A legacy provider-only account `api_key_<provider>` can also be an old
229// scoped account when the provider ends in `__` followed only by digits or
230// underscores: the final portion could be a valid underscore-prefixed key.
231fn legacy_provider_account_is_ambiguous(provider: &str) -> bool {
232    provider.rsplit_once("__").is_some_and(|(prefix, suffix)| {
233        !prefix.is_empty() && suffix.bytes().all(|byte| byte == b'_' || byte.is_ascii_digit())
234    })
235}
236
237/// Migrate plain-text API keys from a config map into secure storage.
238///
239/// Returns a map of provider → success/failure.
240pub fn migrate_custom_api_keys(
241    custom_api_keys: &BTreeMap<String, String>,
242    mode: AuthCredentialsStoreMode,
243) -> Result<BTreeMap<String, bool>> {
244    let mut results = BTreeMap::new();
245
246    for (provider, api_key) in custom_api_keys {
247        let storage = CustomApiKeyStorage::new(provider);
248        match storage.store(api_key, mode) {
249            Ok(()) => {
250                tracing::info!("Migrated API key for provider '{provider}' to secure storage");
251                let _ignored = results.insert(provider.clone(), true);
252            }
253            Err(e) => {
254                tracing::warn!("Failed to migrate API key for provider '{provider}': {e}");
255                let _ignored = results.insert(provider.clone(), false);
256            }
257        }
258    }
259
260    Ok(results)
261}
262
263/// Load all custom API keys from secure storage.
264///
265/// Returns a map of provider → API key for those that have stored keys.
266pub fn load_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<BTreeMap<String, String>> {
267    let mut api_keys = BTreeMap::new();
268
269    for provider in providers {
270        let storage = CustomApiKeyStorage::new(provider);
271        if let Some(key) = storage.load(mode)? {
272            drop(api_keys.insert(provider.clone(), key));
273        }
274    }
275
276    Ok(api_keys)
277}
278
279/// Clear all custom API keys from secure storage.
280pub fn clear_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<()> {
281    for provider in providers {
282        let storage = CustomApiKeyStorage::new(provider);
283        if let Err(e) = storage.clear(mode) {
284            tracing::warn!("Failed to clear API key for provider '{provider}': {e}");
285        }
286    }
287    Ok(())
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293    use serial_test::serial;
294    use std::path::PathBuf;
295    use tempfile::TempDir;
296
297    struct TestAuthDirGuard {
298        temp_dir: Option<TempDir>,
299        previous: Option<PathBuf>,
300    }
301
302    impl TestAuthDirGuard {
303        fn new() -> Self {
304            let temp_dir = TempDir::new().expect("create temp auth dir");
305            let previous = crate::storage_paths::auth_storage_dir_override_for_tests().expect("read auth dir override");
306            crate::storage_paths::set_auth_storage_dir_override_for_tests(Some(temp_dir.path().to_path_buf()))
307                .expect("set temp auth dir override");
308            Self { temp_dir: Some(temp_dir), previous }
309        }
310    }
311
312    impl Drop for TestAuthDirGuard {
313        fn drop(&mut self) {
314            crate::storage_paths::set_auth_storage_dir_override_for_tests(self.previous.clone())
315                .expect("restore auth dir override");
316            if let Some(temp_dir) = self.temp_dir.take() {
317                temp_dir.close().expect("remove temp auth dir");
318            }
319        }
320    }
321
322    #[test]
323    fn credential_identity_normalizes_provider_and_key_name() {
324        let identity = CredentialIdentity::new(" MiMo ", "mimo_token_plan_key").expect("valid identity");
325
326        assert_eq!(identity.provider(), "mimo");
327        assert_eq!(identity.key_name(), "MIMO_TOKEN_PLAN_KEY");
328        assert!(identity.uses_default_key_name("MIMO_TOKEN_PLAN_KEY"));
329        assert!(!identity.uses_default_key_name("MIMO_API_KEY"));
330    }
331
332    #[test]
333    fn credential_identity_rejects_invalid_names() {
334        assert!(CredentialIdentity::new("my corp", "MYCORP_API_KEY").is_err());
335        assert!(CredentialIdentity::new("mycorp", "MY-CORP-API-KEY").is_err());
336        assert!(CredentialIdentity::new("mycorp", "1MYCORP_API_KEY").is_err());
337    }
338
339    #[test]
340    #[serial]
341    #[cfg(unix)]
342    fn file_api_key_storage_round_trips_with_private_permissions() {
343        use std::fs;
344        use std::os::unix::fs::PermissionsExt;
345
346        let guard = TestAuthDirGuard::new();
347        let storage = CustomApiKeyStorage::new("stepfun");
348        storage
349            .store("test-stepfun-key", AuthCredentialsStoreMode::File)
350            .expect("store API key");
351        assert_eq!(
352            storage.load(AuthCredentialsStoreMode::File).expect("load API key").as_deref(),
353            Some("test-stepfun-key")
354        );
355
356        let auth_dir = guard.temp_dir.as_ref().expect("test auth dir").path();
357        assert_eq!(fs::metadata(auth_dir).expect("auth dir metadata").permissions().mode() & 0o777, 0o700);
358        let credential_file = fs::read_dir(auth_dir)
359            .expect("read auth dir")
360            .map(|entry| entry.expect("credential entry").path())
361            .find(|path| path.extension().and_then(|extension| extension.to_str()) == Some("json"))
362            .expect("credential file");
363        assert_eq!(fs::metadata(credential_file).expect("credential metadata").permissions().mode() & 0o777, 0o600);
364    }
365
366    #[test]
367    #[serial]
368    fn keyring_mode_falls_back_to_encrypted_file_when_keyring_is_unavailable() {
369        let _guard = TestAuthDirGuard::new();
370        let storage = CustomApiKeyStorage::new("stepfun");
371
372        storage
373            .store("test-stepfun-key", AuthCredentialsStoreMode::Keyring)
374            .expect("keyring mode should fall back to encrypted file storage");
375        assert_eq!(
376            storage
377                .load(AuthCredentialsStoreMode::Keyring)
378                .expect("load API key")
379                .as_deref(),
380            Some("test-stepfun-key")
381        );
382    }
383
384    #[test]
385    #[serial]
386    fn key_scoped_storage_keeps_provider_profiles_isolated() {
387        let _guard = TestAuthDirGuard::new();
388        let payg = CustomApiKeyStorage::for_provider_key(" MiMo ", "mimo_api_key").expect("payg storage");
389        let token_plan =
390            CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
391
392        payg.store("sk-payg", AuthCredentialsStoreMode::File).expect("store payg");
393        token_plan
394            .store("tp-token-plan", AuthCredentialsStoreMode::File)
395            .expect("store token plan");
396
397        assert_eq!(payg.load(AuthCredentialsStoreMode::File).expect("load payg").as_deref(), Some("sk-payg"));
398        assert_eq!(
399            token_plan
400                .load(AuthCredentialsStoreMode::File)
401                .expect("load token plan")
402                .as_deref(),
403            Some("tp-token-plan")
404        );
405    }
406
407    #[test]
408    #[serial]
409    fn ambiguous_old_accounts_are_not_assigned_to_either_identity() {
410        let _guard = TestAuthDirGuard::new();
411        let old = CredentialStorage::new("vtcode", "api_key_a___B".to_owned());
412        old.store_with_mode("ambiguous", AuthCredentialsStoreMode::File)
413            .expect("store old account");
414        let first = CustomApiKeyStorage::for_provider_key("a_", "_B").expect("first identity");
415        let second = CustomApiKeyStorage::for_provider_key("a", "__B").expect("second identity");
416        assert_eq!(first.load(AuthCredentialsStoreMode::File).expect("first load"), None);
417        assert_eq!(second.load(AuthCredentialsStoreMode::File).expect("second load"), None);
418        first.store("first", AuthCredentialsStoreMode::File).expect("store first");
419        second.store("second", AuthCredentialsStoreMode::File).expect("store second");
420        assert_eq!(first.load(AuthCredentialsStoreMode::File).expect("first reload").as_deref(), Some("first"));
421        assert_eq!(second.load(AuthCredentialsStoreMode::File).expect("second reload").as_deref(), Some("second"));
422        assert_eq!(
423            old.load_with_mode(AuthCredentialsStoreMode::File)
424                .expect("old remains")
425                .as_deref(),
426            Some("ambiguous")
427        );
428    }
429
430    #[test]
431    #[serial]
432    fn default_identity_lazily_migrates_provider_only_storage() {
433        let _guard = TestAuthDirGuard::new();
434        let legacy = CustomApiKeyStorage::new("mimo");
435        let target = CustomApiKeyStorage::for_provider_key("mimo", "MIMO_API_KEY").expect("target storage");
436
437        legacy
438            .store("legacy-key", AuthCredentialsStoreMode::File)
439            .expect("store legacy key");
440        assert_eq!(
441            target
442                .load_with_legacy_fallback(AuthCredentialsStoreMode::File, true)
443                .expect("migrate legacy key")
444                .as_deref(),
445            Some("legacy-key")
446        );
447        assert_eq!(legacy.load(AuthCredentialsStoreMode::File).expect("legacy should be cleared"), None);
448        assert_eq!(
449            target
450                .load(AuthCredentialsStoreMode::File)
451                .expect("load migrated key")
452                .as_deref(),
453            Some("legacy-key")
454        );
455    }
456
457    #[test]
458    #[serial]
459    fn ambiguous_provider_only_fallback_is_neither_loaded_nor_cleared() {
460        let _guard = TestAuthDirGuard::new();
461        // This could be provider `a` with key `_1`, or provider-only `a__1`.
462        let old = CredentialStorage::new("vtcode", "api_key_a__1");
463        old.store_with_mode("other-identity", AuthCredentialsStoreMode::File)
464            .expect("store ambiguous account");
465        let target = CustomApiKeyStorage::for_provider_key("a__1", "A_1_API_KEY").expect("target identity");
466        assert_eq!(
467            target
468                .load_with_legacy_fallback(AuthCredentialsStoreMode::File, true)
469                .expect("load"),
470            None
471        );
472        target
473            .clear_with_legacy_fallback(AuthCredentialsStoreMode::File, true)
474            .expect("clear target");
475        assert_eq!(
476            old.load_with_mode(AuthCredentialsStoreMode::File)
477                .expect("old remains")
478                .as_deref(),
479            Some("other-identity")
480        );
481    }
482
483    #[test]
484    fn legacy_provider_ambiguity_distinguishes_valid_account_boundaries() {
485        for provider in ["a__", "a__1", "a___1", "___"] {
486            assert!(legacy_provider_account_is_ambiguous(provider), "{provider}");
487        }
488        for provider in ["a", "a_b", "a__b", "a__1b", "__"] {
489            assert!(!legacy_provider_account_is_ambiguous(provider), "{provider}");
490        }
491    }
492
493    #[test]
494    #[serial]
495    fn non_default_identity_does_not_reuse_provider_only_storage() {
496        let _guard = TestAuthDirGuard::new();
497        let legacy = CustomApiKeyStorage::new("mimo");
498        let token_plan =
499            CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
500
501        legacy
502            .store("legacy-payg", AuthCredentialsStoreMode::File)
503            .expect("store legacy key");
504        assert_eq!(
505            token_plan
506                .load_with_legacy_fallback(AuthCredentialsStoreMode::File, false)
507                .expect("load token-plan key"),
508            None
509        );
510        assert_eq!(
511            legacy.load(AuthCredentialsStoreMode::File).expect("legacy remains"),
512            Some("legacy-payg".to_string())
513        );
514    }
515
516    #[test]
517    #[serial]
518    fn clearing_one_identity_does_not_clear_another() {
519        let _guard = TestAuthDirGuard::new();
520        let payg = CustomApiKeyStorage::for_provider_key("mimo", "MIMO_API_KEY").expect("payg storage");
521        let token_plan =
522            CustomApiKeyStorage::for_provider_key("mimo", "MIMO_TOKEN_PLAN_KEY").expect("token-plan storage");
523
524        payg.store("sk-payg", AuthCredentialsStoreMode::File).expect("store payg");
525        token_plan
526            .store("tp-token-plan", AuthCredentialsStoreMode::File)
527            .expect("store token plan");
528        payg.clear_with_legacy_fallback(AuthCredentialsStoreMode::File, true)
529            .expect("clear payg");
530
531        assert_eq!(payg.load(AuthCredentialsStoreMode::File).expect("payg cleared"), None);
532        assert_eq!(
533            token_plan
534                .load(AuthCredentialsStoreMode::File)
535                .expect("token plan remains")
536                .as_deref(),
537            Some("tp-token-plan")
538        );
539    }
540}
vtcode_auth/credentials/mod.rs

vtcode_auth/credentials/
mod.rs
