1use anyhow::{Context, Result, anyhow};
26use serde::{Deserialize, Serialize};
27use std::collections::BTreeMap;
28
29#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
37#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
38#[serde(rename_all = "lowercase")]
39pub enum AuthCredentialsStoreMode {
40 Keyring,
44 File,
47 Auto,
49}
50
51impl Default for AuthCredentialsStoreMode {
52 fn default() -> Self {
55 Self::Keyring
56 }
57}
58
59impl AuthCredentialsStoreMode {
60 pub fn effective_mode(self) -> Self {
62 match self {
63 Self::Auto => {
64 if is_keyring_functional() {
66 Self::Keyring
67 } else {
68 tracing::debug!("Keyring not available, falling back to file storage");
69 Self::File
70 }
71 }
72 mode => mode,
73 }
74 }
75}
76
77pub(crate) fn is_keyring_functional() -> bool {
82 let test_user = format!("test_{}", std::process::id());
84 let entry = match keyring::Entry::new("vtcode", &test_user) {
85 Ok(e) => e,
86 Err(_) => return false,
87 };
88
89 if entry.set_password("test").is_err() {
91 return false;
92 }
93
94 let functional = entry.get_password().is_ok();
96
97 let _ = entry.delete_credential();
99
100 functional
101}
102
103pub struct CredentialStorage {
108 service: String,
109 user: String,
110}
111
112impl CredentialStorage {
113 pub fn new(service: impl Into<String>, user: impl Into<String>) -> Self {
119 Self {
120 service: service.into(),
121 user: user.into(),
122 }
123 }
124
125 pub fn store_with_mode(&self, value: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
131 match mode.effective_mode() {
132 AuthCredentialsStoreMode::Keyring => self.store_keyring(value),
133 AuthCredentialsStoreMode::File => Err(anyhow!(
134 "File storage requires the file_storage feature or custom implementation"
135 )),
136 _ => unreachable!(),
137 }
138 }
139
140 pub fn store(&self, value: &str) -> Result<()> {
142 self.store_keyring(value)
143 }
144
145 fn store_keyring(&self, value: &str) -> Result<()> {
147 let entry = keyring::Entry::new(&self.service, &self.user)
148 .context("Failed to access OS keyring")?;
149
150 entry
151 .set_password(value)
152 .context("Failed to store credential in OS keyring")?;
153
154 tracing::debug!(
155 "Credential stored in OS keyring for {}/{}",
156 self.service,
157 self.user
158 );
159 Ok(())
160 }
161
162 pub fn load_with_mode(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
166 match mode.effective_mode() {
167 AuthCredentialsStoreMode::Keyring => self.load_keyring(),
168 AuthCredentialsStoreMode::File => Err(anyhow!(
169 "File storage requires the file_storage feature or custom implementation"
170 )),
171 _ => unreachable!(),
172 }
173 }
174
175 pub fn load(&self) -> Result<Option<String>> {
179 self.load_keyring()
180 }
181
182 fn load_keyring(&self) -> Result<Option<String>> {
184 let entry = match keyring::Entry::new(&self.service, &self.user) {
185 Ok(e) => e,
186 Err(_) => return Ok(None),
187 };
188
189 match entry.get_password() {
190 Ok(value) => Ok(Some(value)),
191 Err(keyring::Error::NoEntry) => Ok(None),
192 Err(e) => Err(anyhow!("Failed to read from keyring: {}", e)),
193 }
194 }
195
196 pub fn clear_with_mode(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
198 match mode.effective_mode() {
199 AuthCredentialsStoreMode::Keyring => self.clear_keyring(),
200 AuthCredentialsStoreMode::File => Ok(()), _ => unreachable!(),
202 }
203 }
204
205 pub fn clear(&self) -> Result<()> {
207 self.clear_keyring()
208 }
209
210 fn clear_keyring(&self) -> Result<()> {
212 let entry = match keyring::Entry::new(&self.service, &self.user) {
213 Ok(e) => e,
214 Err(_) => return Ok(()),
215 };
216
217 match entry.delete_credential() {
218 Ok(_) => {
219 tracing::debug!(
220 "Credential cleared from keyring for {}/{}",
221 self.service,
222 self.user
223 );
224 }
225 Err(keyring::Error::NoEntry) => {}
226 Err(e) => return Err(anyhow!("Failed to clear keyring entry: {}", e)),
227 }
228
229 Ok(())
230 }
231}
232
233pub struct CustomApiKeyStorage {
238 storage: CredentialStorage,
239}
240
241impl CustomApiKeyStorage {
242 pub fn new(provider: &str) -> Self {
247 Self {
248 storage: CredentialStorage::new(
249 "vtcode",
250 format!("api_key_{}", provider.to_lowercase()),
251 ),
252 }
253 }
254
255 pub fn store(&self, api_key: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
261 self.storage.store_with_mode(api_key, mode)
262 }
263
264 pub fn load(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
268 self.storage.load_with_mode(mode)
269 }
270
271 pub fn clear(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
273 self.storage.clear_with_mode(mode)
274 }
275}
276
277pub fn migrate_custom_api_keys_to_keyring(
290 custom_api_keys: &BTreeMap<String, String>,
291 mode: AuthCredentialsStoreMode,
292) -> Result<BTreeMap<String, bool>> {
293 let mut migration_results = BTreeMap::new();
294
295 for (provider, api_key) in custom_api_keys {
296 let storage = CustomApiKeyStorage::new(provider);
297 match storage.store(api_key, mode) {
298 Ok(()) => {
299 tracing::info!(
300 "Migrated API key for provider '{}' to secure storage",
301 provider
302 );
303 migration_results.insert(provider.clone(), true);
304 }
305 Err(e) => {
306 tracing::warn!(
307 "Failed to migrate API key for provider '{}': {}",
308 provider,
309 e
310 );
311 migration_results.insert(provider.clone(), false);
312 }
313 }
314 }
315
316 Ok(migration_results)
317}
318
319pub fn load_custom_api_keys(
330 providers: &[String],
331 mode: AuthCredentialsStoreMode,
332) -> Result<BTreeMap<String, String>> {
333 let mut api_keys = BTreeMap::new();
334
335 for provider in providers {
336 let storage = CustomApiKeyStorage::new(provider);
337 if let Some(key) = storage.load(mode)? {
338 api_keys.insert(provider.clone(), key);
339 }
340 }
341
342 Ok(api_keys)
343}
344
345pub fn clear_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<()> {
351 for provider in providers {
352 let storage = CustomApiKeyStorage::new(provider);
353 if let Err(e) = storage.clear(mode) {
354 tracing::warn!("Failed to clear API key for provider '{}': {}", provider, e);
355 }
356 }
357 Ok(())
358}
359
360#[cfg(test)]
361mod tests {
362 use super::*;
363
364 #[test]
365 fn test_storage_mode_default_is_keyring() {
366 assert_eq!(
367 AuthCredentialsStoreMode::default(),
368 AuthCredentialsStoreMode::Keyring
369 );
370 }
371
372 #[test]
373 fn test_storage_mode_effective_mode() {
374 assert_eq!(
375 AuthCredentialsStoreMode::Keyring.effective_mode(),
376 AuthCredentialsStoreMode::Keyring
377 );
378 assert_eq!(
379 AuthCredentialsStoreMode::File.effective_mode(),
380 AuthCredentialsStoreMode::File
381 );
382
383 let auto_mode = AuthCredentialsStoreMode::Auto.effective_mode();
385 assert!(
386 auto_mode == AuthCredentialsStoreMode::Keyring
387 || auto_mode == AuthCredentialsStoreMode::File
388 );
389 }
390
391 #[test]
392 fn test_storage_mode_serialization() {
393 let keyring_json = serde_json::to_string(&AuthCredentialsStoreMode::Keyring).unwrap();
394 assert_eq!(keyring_json, "\"keyring\"");
395
396 let file_json = serde_json::to_string(&AuthCredentialsStoreMode::File).unwrap();
397 assert_eq!(file_json, "\"file\"");
398
399 let auto_json = serde_json::to_string(&AuthCredentialsStoreMode::Auto).unwrap();
400 assert_eq!(auto_json, "\"auto\"");
401
402 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"keyring\"").unwrap();
404 assert_eq!(parsed, AuthCredentialsStoreMode::Keyring);
405
406 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"file\"").unwrap();
407 assert_eq!(parsed, AuthCredentialsStoreMode::File);
408
409 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"auto\"").unwrap();
410 assert_eq!(parsed, AuthCredentialsStoreMode::Auto);
411 }
412
413 #[test]
414 fn test_credential_storage_new() {
415 let storage = CredentialStorage::new("vtcode", "test_key");
416 assert_eq!(storage.service, "vtcode");
417 assert_eq!(storage.user, "test_key");
418 }
419
420 #[test]
421 fn test_is_keyring_functional_check() {
422 let _functional = is_keyring_functional();
425 }
426}