Skip to main content

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 use storage::CredentialStorage;
21
22use std::collections::BTreeMap;
23
24use anyhow::{Context, Result, bail};
25
26/// Validated identity for a stored provider credential.
27///
28/// Provider names are normalized to lowercase and environment-variable names
29/// are normalized to uppercase. Keeping both dimensions in a value object
30/// prevents a credential for one provider profile from being reused for
31/// another profile that happens to share the provider name.
32#[derive(Debug, Clone, PartialEq, Eq, Hash)]
33pub struct CredentialIdentity {
34    provider: String,
35    key_name: String,
36}
37
38impl CredentialIdentity {
39    /// Create a credential identity from a provider name and environment
40    /// variable-style key name.
41    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    /// Normalized provider name.
58    pub fn provider(&self) -> &str {
59        &self.provider
60    }
61
62    /// Normalized credential key name.
63    pub fn key_name(&self) -> &str {
64        &self.key_name
65    }
66
67    /// Whether this identity uses the provider's default key name.
68    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
73/// Custom API Key storage for provider-specific keys.
74///
75/// Provides secure storage and retrieval of API keys for custom providers
76/// using the OS keyring or encrypted file storage.
77pub struct CustomApiKeyStorage {
78    provider: String,
79    identity: Option<CredentialIdentity>,
80    storage: CredentialStorage,
81}
82
83impl CustomApiKeyStorage {
84    /// Create a provider-only legacy storage handle.
85    ///
86    /// New callers should use [`Self::for_provider_key`]. This constructor is
87    /// retained so older integrations can still read and clear the legacy
88    /// `api_key_<provider>` entry.
89    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    /// Create storage scoped to a provider and credential key name.
99    pub fn for_provider_key(provider: &str, key_name: &str) -> Result<Self> {
100        Self::for_identity(CredentialIdentity::new(provider, key_name)?)
101    }
102
103    /// Create storage scoped to a validated credential identity.
104    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    /// Return the identity for a key-scoped handle, or `None` for a legacy
115    /// provider-only handle.
116    pub fn identity(&self) -> Option<&CredentialIdentity> {
117        self.identity.as_ref()
118    }
119
120    /// Store an API key securely.
121    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    /// Retrieve a stored API key.
135    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    /// Load a key-scoped credential and, when explicitly allowed, lazily
149    /// migrate the provider-only legacy entry into this identity.
150    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    /// Clear (delete) a stored API key.
176    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    /// Clear this key-scoped credential and optionally its provider-only
185    /// legacy fallback.
186    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
233/// Migrate plain-text API keys from a config map into secure storage.
234///
235/// Returns a map of provider → success/failure.
236pub 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
259/// Load all custom API keys from secure storage.
260///
261/// Returns a map of provider → API key for those that have stored keys.
262pub 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
275/// Clear all custom API keys from secure storage.
276pub 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}