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(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        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    /// Return the identity for a key-scoped handle, or `None` for a legacy
116    /// provider-only handle.
117    pub fn identity(&self) -> Option<&CredentialIdentity> {
118        self.identity.as_ref()
119    }
120
121    /// Store an API key securely.
122    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    /// Retrieve a stored API key.
132    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    /// Load a key-scoped credential and, when explicitly allowed, lazily
146    /// migrate the provider-only legacy entry into this identity.
147    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        // This is the obsolete provider-only secure-storage entry, not the
167        // rollback-safe plaintext file under the legacy directory. Remove the
168        // former after the key-scoped copy succeeds, but preserve the latter.
169        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    /// Clear (delete) a stored API key.
177    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    /// Clear this key-scoped credential and optionally its provider-only
186    /// legacy fallback.
187    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
226/// Migrate plain-text API keys from a config map into secure storage.
227///
228/// Returns a map of provider → success/failure.
229pub 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
252/// Load all custom API keys from secure storage.
253///
254/// Returns a map of provider → API key for those that have stored keys.
255pub 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
268/// Clear all custom API keys from secure storage.
269pub 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}