1use anyhow::{Context, Result, anyhow};
26use serde::{Deserialize, Serialize};
27use std::collections::BTreeMap;
28
29#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
37#[serde(rename_all = "lowercase")]
38pub enum AuthCredentialsStoreMode {
39 Keyring,
43 File,
46 Auto,
48}
49
50impl Default for AuthCredentialsStoreMode {
51 fn default() -> Self {
54 Self::Keyring
55 }
56}
57
58impl AuthCredentialsStoreMode {
59 pub fn effective_mode(self) -> Self {
61 match self {
62 Self::Auto => {
63 if is_keyring_functional() {
65 Self::Keyring
66 } else {
67 tracing::debug!("Keyring not available, falling back to file storage");
68 Self::File
69 }
70 }
71 mode => mode,
72 }
73 }
74}
75
76pub(crate) fn is_keyring_functional() -> bool {
81 let test_user = format!("test_{}", std::process::id());
83 let entry = match keyring::Entry::new("vtcode", &test_user) {
84 Ok(e) => e,
85 Err(_) => return false,
86 };
87
88 if entry.set_password("test").is_err() {
90 return false;
91 }
92
93 let functional = entry.get_password().is_ok();
95
96 let _ = entry.delete_credential();
98
99 functional
100}
101
102pub struct CredentialStorage {
107 service: String,
108 user: String,
109}
110
111impl CredentialStorage {
112 pub fn new(service: impl Into<String>, user: impl Into<String>) -> Self {
118 Self {
119 service: service.into(),
120 user: user.into(),
121 }
122 }
123
124 pub fn store_with_mode(&self, value: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
130 match mode.effective_mode() {
131 AuthCredentialsStoreMode::Keyring => self.store_keyring(value),
132 AuthCredentialsStoreMode::File => Err(anyhow!(
133 "File storage requires the file_storage feature or custom implementation"
134 )),
135 _ => unreachable!(),
136 }
137 }
138
139 pub fn store(&self, value: &str) -> Result<()> {
141 self.store_keyring(value)
142 }
143
144 fn store_keyring(&self, value: &str) -> Result<()> {
146 let entry = keyring::Entry::new(&self.service, &self.user)
147 .context("Failed to access OS keyring")?;
148
149 entry
150 .set_password(value)
151 .context("Failed to store credential in OS keyring")?;
152
153 tracing::debug!(
154 "Credential stored in OS keyring for {}/{}",
155 self.service,
156 self.user
157 );
158 Ok(())
159 }
160
161 pub fn load_with_mode(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
165 match mode.effective_mode() {
166 AuthCredentialsStoreMode::Keyring => self.load_keyring(),
167 AuthCredentialsStoreMode::File => Err(anyhow!(
168 "File storage requires the file_storage feature or custom implementation"
169 )),
170 _ => unreachable!(),
171 }
172 }
173
174 pub fn load(&self) -> Result<Option<String>> {
178 self.load_keyring()
179 }
180
181 fn load_keyring(&self) -> Result<Option<String>> {
183 let entry = match keyring::Entry::new(&self.service, &self.user) {
184 Ok(e) => e,
185 Err(_) => return Ok(None),
186 };
187
188 match entry.get_password() {
189 Ok(value) => Ok(Some(value)),
190 Err(keyring::Error::NoEntry) => Ok(None),
191 Err(e) => Err(anyhow!("Failed to read from keyring: {}", e)),
192 }
193 }
194
195 pub fn clear_with_mode(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
197 match mode.effective_mode() {
198 AuthCredentialsStoreMode::Keyring => self.clear_keyring(),
199 AuthCredentialsStoreMode::File => Ok(()), _ => unreachable!(),
201 }
202 }
203
204 pub fn clear(&self) -> Result<()> {
206 self.clear_keyring()
207 }
208
209 fn clear_keyring(&self) -> Result<()> {
211 let entry = match keyring::Entry::new(&self.service, &self.user) {
212 Ok(e) => e,
213 Err(_) => return Ok(()),
214 };
215
216 match entry.delete_credential() {
217 Ok(_) => {
218 tracing::debug!(
219 "Credential cleared from keyring for {}/{}",
220 self.service,
221 self.user
222 );
223 }
224 Err(keyring::Error::NoEntry) => {}
225 Err(e) => return Err(anyhow!("Failed to clear keyring entry: {}", e)),
226 }
227
228 Ok(())
229 }
230}
231
232#[cfg(test)]
233mod tests {
234 use super::*;
235
236 #[test]
237 fn test_storage_mode_default_is_keyring() {
238 assert_eq!(
239 AuthCredentialsStoreMode::default(),
240 AuthCredentialsStoreMode::Keyring
241 );
242 }
243
244 #[test]
245 fn test_storage_mode_effective_mode() {
246 assert_eq!(
247 AuthCredentialsStoreMode::Keyring.effective_mode(),
248 AuthCredentialsStoreMode::Keyring
249 );
250 assert_eq!(
251 AuthCredentialsStoreMode::File.effective_mode(),
252 AuthCredentialsStoreMode::File
253 );
254
255 let auto_mode = AuthCredentialsStoreMode::Auto.effective_mode();
257 assert!(
258 auto_mode == AuthCredentialsStoreMode::Keyring
259 || auto_mode == AuthCredentialsStoreMode::File
260 );
261 }
262
263 #[test]
264 fn test_storage_mode_serialization() {
265 let keyring_json = serde_json::to_string(&AuthCredentialsStoreMode::Keyring).unwrap();
266 assert_eq!(keyring_json, "\"keyring\"");
267
268 let file_json = serde_json::to_string(&AuthCredentialsStoreMode::File).unwrap();
269 assert_eq!(file_json, "\"file\"");
270
271 let auto_json = serde_json::to_string(&AuthCredentialsStoreMode::Auto).unwrap();
272 assert_eq!(auto_json, "\"auto\"");
273
274 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"keyring\"").unwrap();
276 assert_eq!(parsed, AuthCredentialsStoreMode::Keyring);
277
278 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"file\"").unwrap();
279 assert_eq!(parsed, AuthCredentialsStoreMode::File);
280
281 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"auto\"").unwrap();
282 assert_eq!(parsed, AuthCredentialsStoreMode::Auto);
283 }
284
285 #[test]
286 fn test_credential_storage_new() {
287 let storage = CredentialStorage::new("vtcode", "test_key");
288 assert_eq!(storage.service, "vtcode");
289 assert_eq!(storage.user, "test_key");
290 }
291
292 #[test]
293 fn test_is_keyring_functional_check() {
294 let _functional = is_keyring_functional();
297 }
298}
299
300pub struct CustomApiKeyStorage {
305 storage: CredentialStorage,
306}
307
308impl CustomApiKeyStorage {
309 pub fn new(provider: &str) -> Self {
314 Self {
315 storage: CredentialStorage::new("vtcode", format!("api_key_{}", provider.to_lowercase())),
316 }
317 }
318
319 pub fn store(&self, api_key: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
325 self.storage.store_with_mode(api_key, mode)
326 }
327
328 pub fn load(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
332 self.storage.load_with_mode(mode)
333 }
334
335 pub fn clear(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
337 self.storage.clear_with_mode(mode)
338 }
339}
340
341pub fn migrate_custom_api_keys_to_keyring(
354 custom_api_keys: &BTreeMap<String, String>,
355 mode: AuthCredentialsStoreMode,
356) -> Result<BTreeMap<String, bool>> {
357 let mut migration_results = BTreeMap::new();
358
359 for (provider, api_key) in custom_api_keys {
360 let storage = CustomApiKeyStorage::new(provider);
361 match storage.store(api_key, mode) {
362 Ok(()) => {
363 tracing::info!(
364 "Migrated API key for provider '{}' to secure storage",
365 provider
366 );
367 migration_results.insert(provider.clone(), true);
368 }
369 Err(e) => {
370 tracing::warn!(
371 "Failed to migrate API key for provider '{}': {}",
372 provider,
373 e
374 );
375 migration_results.insert(provider.clone(), false);
376 }
377 }
378 }
379
380 Ok(migration_results)
381}
382
383pub fn load_custom_api_keys(
394 providers: &[String],
395 mode: AuthCredentialsStoreMode,
396) -> Result<BTreeMap<String, String>> {
397 let mut api_keys = BTreeMap::new();
398
399 for provider in providers {
400 let storage = CustomApiKeyStorage::new(provider);
401 if let Some(key) = storage.load(mode)? {
402 api_keys.insert(provider.clone(), key);
403 }
404 }
405
406 Ok(api_keys)
407}
408
409pub fn clear_custom_api_keys(
415 providers: &[String],
416 mode: AuthCredentialsStoreMode,
417) -> Result<()> {
418 for provider in providers {
419 let storage = CustomApiKeyStorage::new(provider);
420 if let Err(e) = storage.clear(mode) {
421 tracing::warn!(
422 "Failed to clear API key for provider '{}': {}",
423 provider,
424 e
425 );
426 }
427 }
428 Ok(())
429}