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