vtcode_auth/credentials/
mod.rs1mod encryption;
14pub(crate) mod keyring;
15mod legacy;
16mod mode;
17mod storage;
18
19pub use mode::AuthCredentialsStoreMode;
20pub use storage::CredentialStorage;
21
22use std::collections::BTreeMap;
23
24use anyhow::{Context, Result, bail};
25
26pub struct CustomApiKeyStorage {
31 provider: String,
32 storage: CredentialStorage,
33}
34
35impl CustomApiKeyStorage {
36 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 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 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 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
108pub 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
134pub 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
150pub 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}