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/// Custom API Key storage for provider-specific keys.
27///
28/// Provides secure storage and retrieval of API keys for custom providers
29/// using the OS keyring or encrypted file storage.
30pub struct CustomApiKeyStorage {
31    provider: String,
32    storage: CredentialStorage,
33}
34
35impl CustomApiKeyStorage {
36    /// Create a new custom API key storage for a specific provider.
37    pub fn new(provider: &str) -> Self {
38        let normalized_provider = provider.trim().to_lowercase();
39        Self {
40            provider: normalized_provider.clone(),
41            storage: CredentialStorage::new("vtcode", format!("api_key_{normalized_provider}")),
42        }
43    }
44
45    /// Store an API key securely.
46    pub fn store(&self, api_key: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
47        let api_key = api_key.trim();
48        if api_key.is_empty() {
49            bail!("API key cannot be empty");
50        }
51        self.storage.store_with_mode(api_key, mode)?;
52        let persisted = self
53            .storage
54            .load_with_mode(mode)
55            .context("failed to verify persisted API key")?;
56        if persisted.as_deref() != Some(api_key) {
57            bail!("secure storage did not return the API key after saving");
58        }
59        legacy::clear_for_provider(&self.provider)
60            .context("failed to remove legacy plaintext credential after secure save")?;
61        Ok(())
62    }
63
64    /// Retrieve a stored API key.
65    pub fn load(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
66        if let Some(key) = self.storage.load_with_mode(mode)? {
67            let key = key.trim();
68            return Ok((!key.is_empty()).then(|| key.to_owned()));
69        }
70
71        self.load_legacy_auth_json(mode)
72    }
73
74    /// Clear (delete) a stored API key.
75    pub fn clear(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
76        self.storage.clear_with_mode(mode)?;
77        legacy::clear_for_provider(&self.provider).context("failed to remove legacy plaintext credential")?;
78        Ok(())
79    }
80
81    fn load_legacy_auth_json(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
82        let Some(legacy_entry) = legacy::load_for_provider(&self.provider)? else {
83            return Ok(None);
84        };
85
86        if let Err(err) = self.store(&legacy_entry.api_key, mode) {
87            tracing::warn!(
88                "Failed to migrate legacy plaintext auth.json entry for provider '{}' into secure storage: {}",
89                self.provider,
90                err
91            );
92            return Err(err).context("failed to migrate legacy API key into secure storage");
93        }
94
95        let path = crate::storage_paths::legacy_auth_storage_path().ok();
96        if let Some(p) = path {
97            legacy::delete_file(&p).context("failed to remove migrated plaintext auth file")?;
98        }
99
100        tracing::warn!(
101            "Migrated legacy plaintext auth.json entry for provider '{}' into secure storage",
102            self.provider
103        );
104        self.load(mode)
105    }
106}
107
108/// Migrate plain-text API keys from a config map into secure storage.
109///
110/// Returns a map of provider → success/failure.
111pub fn migrate_custom_api_keys(
112    custom_api_keys: &BTreeMap<String, String>,
113    mode: AuthCredentialsStoreMode,
114) -> Result<BTreeMap<String, bool>> {
115    let mut results = BTreeMap::new();
116
117    for (provider, api_key) in custom_api_keys {
118        let storage = CustomApiKeyStorage::new(provider);
119        match storage.store(api_key, mode) {
120            Ok(()) => {
121                tracing::info!("Migrated API key for provider '{provider}' to secure storage");
122                results.insert(provider.clone(), true);
123            }
124            Err(e) => {
125                tracing::warn!("Failed to migrate API key for provider '{provider}': {e}");
126                results.insert(provider.clone(), false);
127            }
128        }
129    }
130
131    Ok(results)
132}
133
134/// Load all custom API keys from secure storage.
135///
136/// Returns a map of provider → API key for those that have stored keys.
137pub fn load_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<BTreeMap<String, String>> {
138    let mut api_keys = BTreeMap::new();
139
140    for provider in providers {
141        let storage = CustomApiKeyStorage::new(provider);
142        if let Some(key) = storage.load(mode)? {
143            api_keys.insert(provider.clone(), key);
144        }
145    }
146
147    Ok(api_keys)
148}
149
150/// Clear all custom API keys from secure storage.
151pub fn clear_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<()> {
152    for provider in providers {
153        let storage = CustomApiKeyStorage::new(provider);
154        if let Err(e) = storage.clear(mode) {
155            tracing::warn!("Failed to clear API key for provider '{provider}': {e}");
156        }
157    }
158    Ok(())
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use serial_test::serial;
165    use std::path::PathBuf;
166    use tempfile::TempDir;
167
168    struct TestAuthDirGuard {
169        temp_dir: Option<TempDir>,
170        previous: Option<PathBuf>,
171    }
172
173    impl TestAuthDirGuard {
174        fn new() -> Self {
175            let temp_dir = TempDir::new().expect("create temp auth dir");
176            let previous = crate::storage_paths::auth_storage_dir_override_for_tests().expect("read auth dir override");
177            crate::storage_paths::set_auth_storage_dir_override_for_tests(Some(temp_dir.path().to_path_buf()))
178                .expect("set temp auth dir override");
179            Self { temp_dir: Some(temp_dir), previous }
180        }
181    }
182
183    impl Drop for TestAuthDirGuard {
184        fn drop(&mut self) {
185            crate::storage_paths::set_auth_storage_dir_override_for_tests(self.previous.clone())
186                .expect("restore auth dir override");
187            if let Some(temp_dir) = self.temp_dir.take() {
188                temp_dir.close().expect("remove temp auth dir");
189            }
190        }
191    }
192
193    #[test]
194    #[serial]
195    #[cfg(unix)]
196    fn file_api_key_storage_round_trips_with_private_permissions() {
197        use std::fs;
198        use std::os::unix::fs::PermissionsExt;
199
200        let guard = TestAuthDirGuard::new();
201        let storage = CustomApiKeyStorage::new("stepfun");
202        storage
203            .store("test-stepfun-key", AuthCredentialsStoreMode::File)
204            .expect("store API key");
205        assert_eq!(
206            storage.load(AuthCredentialsStoreMode::File).expect("load API key").as_deref(),
207            Some("test-stepfun-key")
208        );
209
210        let auth_dir = guard.temp_dir.as_ref().expect("test auth dir").path();
211        assert_eq!(fs::metadata(auth_dir).expect("auth dir metadata").permissions().mode() & 0o777, 0o700);
212        let credential_file = fs::read_dir(auth_dir)
213            .expect("read auth dir")
214            .map(|entry| entry.expect("credential entry").path())
215            .find(|path| path.extension().and_then(|extension| extension.to_str()) == Some("json"))
216            .expect("credential file");
217        assert_eq!(fs::metadata(credential_file).expect("credential metadata").permissions().mode() & 0o777, 0o600);
218    }
219
220    #[test]
221    #[serial]
222    fn keyring_mode_falls_back_to_encrypted_file_when_keyring_is_unavailable() {
223        let _guard = TestAuthDirGuard::new();
224        let storage = CustomApiKeyStorage::new("stepfun");
225
226        storage
227            .store("test-stepfun-key", AuthCredentialsStoreMode::Keyring)
228            .expect("keyring mode should fall back to encrypted file storage");
229        assert_eq!(
230            storage
231                .load(AuthCredentialsStoreMode::Keyring)
232                .expect("load API key")
233                .as_deref(),
234            Some("test-stepfun-key")
235        );
236    }
237}