1use anyhow::{Context, Result, anyhow};
29use base64::Engine;
30use base64::engine::general_purpose::STANDARD;
31use ring::aead::{self, Aad, LessSafeKey, NONCE_LEN, Nonce, UnboundKey};
32use ring::rand::{SecureRandom, SystemRandom};
33use serde::{Deserialize, Serialize};
34use std::collections::BTreeMap;
35use std::fs;
36
37use crate::storage_paths::auth_storage_dir;
38use crate::storage_paths::legacy_auth_storage_path;
39
40const ENCRYPTED_CREDENTIAL_VERSION: u8 = 1;
41
42#[derive(Debug, Serialize, Deserialize)]
43struct EncryptedCredential {
44 nonce: String,
45 ciphertext: String,
46 version: u8,
47}
48
49#[derive(Debug, Deserialize)]
50struct LegacyAuthFile {
51 mode: String,
52 provider: String,
53 api_key: String,
54}
55
56#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
64#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
65#[serde(rename_all = "lowercase")]
66pub enum AuthCredentialsStoreMode {
67 Keyring,
71 File,
74 Auto,
76}
77
78impl Default for AuthCredentialsStoreMode {
79 fn default() -> Self {
82 Self::Keyring
83 }
84}
85
86impl AuthCredentialsStoreMode {
87 pub fn effective_mode(self) -> Self {
89 match self {
90 Self::Auto => {
91 if is_keyring_functional() {
93 Self::Keyring
94 } else {
95 tracing::debug!("Keyring not available, falling back to file storage");
96 Self::File
97 }
98 }
99 mode => mode,
100 }
101 }
102}
103
104pub(crate) fn is_keyring_functional() -> bool {
109 let test_user = format!("test_{}", std::process::id());
111 let entry = match keyring::Entry::new("vtcode", &test_user) {
112 Ok(e) => e,
113 Err(_) => return false,
114 };
115
116 if entry.set_password("test").is_err() {
118 return false;
119 }
120
121 let functional = entry.get_password().is_ok();
123
124 let _ = entry.delete_credential();
126
127 functional
128}
129
130pub struct CredentialStorage {
135 service: String,
136 user: String,
137}
138
139impl CredentialStorage {
140 pub fn new(service: impl Into<String>, user: impl Into<String>) -> Self {
146 Self {
147 service: service.into(),
148 user: user.into(),
149 }
150 }
151
152 pub fn store_with_mode(&self, value: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
158 match mode.effective_mode() {
159 AuthCredentialsStoreMode::Keyring => match self.store_keyring(value) {
160 Ok(()) => {
161 let _ = self.clear_file();
162 Ok(())
163 }
164 Err(err) => {
165 tracing::warn!(
166 "Failed to store credential in OS keyring for {}/{}; falling back to encrypted file storage: {}",
167 self.service,
168 self.user,
169 err
170 );
171 self.store_file(value)
172 .context("failed to store credential in encrypted file")
173 }
174 },
175 AuthCredentialsStoreMode::File => self.store_file(value),
176 _ => unreachable!(),
177 }
178 }
179
180 pub fn store(&self, value: &str) -> Result<()> {
182 self.store_keyring(value)
183 }
184
185 fn store_keyring(&self, value: &str) -> Result<()> {
187 let entry = keyring::Entry::new(&self.service, &self.user)
188 .context("Failed to access OS keyring")?;
189
190 entry
191 .set_password(value)
192 .context("Failed to store credential in OS keyring")?;
193
194 tracing::debug!(
195 "Credential stored in OS keyring for {}/{}",
196 self.service,
197 self.user
198 );
199 Ok(())
200 }
201
202 pub fn load_with_mode(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
206 match mode.effective_mode() {
207 AuthCredentialsStoreMode::Keyring => match self.load_keyring() {
208 Ok(Some(value)) => Ok(Some(value)),
209 Ok(None) => self.load_file(),
210 Err(err) => {
211 tracing::warn!(
212 "Failed to read credential from OS keyring for {}/{}; falling back to encrypted file storage: {}",
213 self.service,
214 self.user,
215 err
216 );
217 self.load_file()
218 }
219 },
220 AuthCredentialsStoreMode::File => self.load_file(),
221 _ => unreachable!(),
222 }
223 }
224
225 pub fn load(&self) -> Result<Option<String>> {
229 self.load_keyring()
230 }
231
232 fn load_keyring(&self) -> Result<Option<String>> {
234 let entry = match keyring::Entry::new(&self.service, &self.user) {
235 Ok(e) => e,
236 Err(_) => return Ok(None),
237 };
238
239 match entry.get_password() {
240 Ok(value) => Ok(Some(value)),
241 Err(keyring::Error::NoEntry) => Ok(None),
242 Err(e) => Err(anyhow!("Failed to read from keyring: {}", e)),
243 }
244 }
245
246 pub fn clear_with_mode(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
248 match mode.effective_mode() {
249 AuthCredentialsStoreMode::Keyring => {
250 let mut errors = Vec::new();
251
252 if let Err(err) = self.clear_keyring() {
253 errors.push(err.to_string());
254 }
255 if let Err(err) = self.clear_file() {
256 errors.push(err.to_string());
257 }
258
259 if errors.is_empty() {
260 Ok(())
261 } else {
262 Err(anyhow!(
263 "Failed to clear credential from secure storage: {}",
264 errors.join("; ")
265 ))
266 }
267 }
268 AuthCredentialsStoreMode::File => self.clear_file(),
269 _ => unreachable!(),
270 }
271 }
272
273 pub fn clear(&self) -> Result<()> {
275 self.clear_keyring()
276 }
277
278 fn clear_keyring(&self) -> Result<()> {
280 let entry = match keyring::Entry::new(&self.service, &self.user) {
281 Ok(e) => e,
282 Err(_) => return Ok(()),
283 };
284
285 match entry.delete_credential() {
286 Ok(_) => {
287 tracing::debug!(
288 "Credential cleared from keyring for {}/{}",
289 self.service,
290 self.user
291 );
292 }
293 Err(keyring::Error::NoEntry) => {}
294 Err(e) => return Err(anyhow!("Failed to clear keyring entry: {}", e)),
295 }
296
297 Ok(())
298 }
299
300 fn store_file(&self, value: &str) -> Result<()> {
301 let path = self.file_path()?;
302 let encrypted = encrypt_credential(value)?;
303 let payload = serde_json::to_vec_pretty(&encrypted)
304 .context("failed to serialize encrypted credential")?;
305 fs::write(&path, payload).context("failed to write encrypted credential file")?;
306
307 #[cfg(unix)]
308 {
309 use std::os::unix::fs::PermissionsExt;
310
311 fs::set_permissions(&path, fs::Permissions::from_mode(0o600))
312 .context("failed to set credential file permissions")?;
313 }
314
315 Ok(())
316 }
317
318 fn load_file(&self) -> Result<Option<String>> {
319 let path = self.file_path()?;
320 let data = match fs::read(&path) {
321 Ok(data) => data,
322 Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
323 Err(err) => return Err(anyhow!("failed to read encrypted credential file: {err}")),
324 };
325
326 let encrypted: EncryptedCredential =
327 serde_json::from_slice(&data).context("failed to decode encrypted credential file")?;
328 decrypt_credential(&encrypted).map(Some)
329 }
330
331 fn clear_file(&self) -> Result<()> {
332 let path = self.file_path()?;
333 match fs::remove_file(path) {
334 Ok(()) => Ok(()),
335 Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
336 Err(err) => Err(anyhow!("failed to delete encrypted credential file: {err}")),
337 }
338 }
339
340 fn file_path(&self) -> Result<std::path::PathBuf> {
341 use sha2::Digest as _;
342
343 let mut hasher = sha2::Sha256::new();
344 hasher.update(self.service.as_bytes());
345 hasher.update([0]);
346 hasher.update(self.user.as_bytes());
347 let digest = hasher.finalize();
348 let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest);
349
350 Ok(auth_storage_dir()?.join(format!("credential_{encoded}.json")))
351 }
352}
353
354pub struct CustomApiKeyStorage {
359 provider: String,
360 storage: CredentialStorage,
361}
362
363impl CustomApiKeyStorage {
364 pub fn new(provider: &str) -> Self {
369 let normalized_provider = provider.to_lowercase();
370 Self {
371 provider: normalized_provider.clone(),
372 storage: CredentialStorage::new("vtcode", format!("api_key_{normalized_provider}")),
373 }
374 }
375
376 pub fn store(&self, api_key: &str, mode: AuthCredentialsStoreMode) -> Result<()> {
382 self.storage.store_with_mode(api_key, mode)?;
383 clear_legacy_auth_file_if_matches(&self.provider)?;
384 Ok(())
385 }
386
387 pub fn load(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
391 if let Some(key) = self.storage.load_with_mode(mode)? {
392 return Ok(Some(key));
393 }
394
395 self.load_legacy_auth_json(mode)
396 }
397
398 pub fn clear(&self, mode: AuthCredentialsStoreMode) -> Result<()> {
400 self.storage.clear_with_mode(mode)?;
401 clear_legacy_auth_file_if_matches(&self.provider)?;
402 Ok(())
403 }
404
405 fn load_legacy_auth_json(&self, mode: AuthCredentialsStoreMode) -> Result<Option<String>> {
406 let Some(legacy) = load_legacy_auth_file_for_provider(&self.provider)? else {
407 return Ok(None);
408 };
409
410 if let Err(err) = self.storage.store_with_mode(&legacy.api_key, mode) {
411 tracing::warn!(
412 "Failed to migrate legacy plaintext auth.json entry for provider '{}' into secure storage: {}",
413 self.provider,
414 err
415 );
416 return Ok(Some(legacy.api_key));
417 }
418
419 clear_legacy_auth_file_if_matches(&self.provider)?;
420 tracing::warn!(
421 "Migrated legacy plaintext auth.json entry for provider '{}' into secure storage",
422 self.provider
423 );
424 Ok(Some(legacy.api_key))
425 }
426}
427
428fn encrypt_credential(value: &str) -> Result<EncryptedCredential> {
429 let key = derive_file_encryption_key()?;
430 let rng = SystemRandom::new();
431 let mut nonce_bytes = [0_u8; NONCE_LEN];
432 rng.fill(&mut nonce_bytes)
433 .map_err(|_| anyhow!("failed to generate credential nonce"))?;
434
435 let mut ciphertext = value.as_bytes().to_vec();
436 key.seal_in_place_append_tag(
437 Nonce::assume_unique_for_key(nonce_bytes),
438 Aad::empty(),
439 &mut ciphertext,
440 )
441 .map_err(|_| anyhow!("failed to encrypt credential"))?;
442
443 Ok(EncryptedCredential {
444 nonce: STANDARD.encode(nonce_bytes),
445 ciphertext: STANDARD.encode(ciphertext),
446 version: ENCRYPTED_CREDENTIAL_VERSION,
447 })
448}
449
450fn decrypt_credential(encrypted: &EncryptedCredential) -> Result<String> {
451 if encrypted.version != ENCRYPTED_CREDENTIAL_VERSION {
452 return Err(anyhow!("unsupported encrypted credential format"));
453 }
454
455 let nonce_bytes = STANDARD
456 .decode(&encrypted.nonce)
457 .context("failed to decode credential nonce")?;
458 let nonce_array: [u8; NONCE_LEN] = nonce_bytes
459 .try_into()
460 .map_err(|_| anyhow!("invalid credential nonce length"))?;
461 let mut ciphertext = STANDARD
462 .decode(&encrypted.ciphertext)
463 .context("failed to decode credential ciphertext")?;
464
465 let key = derive_file_encryption_key()?;
466 let plaintext = key
467 .open_in_place(
468 Nonce::assume_unique_for_key(nonce_array),
469 Aad::empty(),
470 &mut ciphertext,
471 )
472 .map_err(|_| anyhow!("failed to decrypt credential"))?;
473
474 String::from_utf8(plaintext.to_vec()).context("failed to parse decrypted credential")
475}
476
477fn derive_file_encryption_key() -> Result<LessSafeKey> {
478 use ring::digest::SHA256;
479 use ring::digest::digest;
480
481 let mut key_material = Vec::new();
482 if let Ok(hostname) = hostname::get() {
483 key_material.extend_from_slice(hostname.as_encoded_bytes());
484 }
485
486 #[cfg(unix)]
487 {
488 key_material.extend_from_slice(&nix::unistd::getuid().as_raw().to_le_bytes());
489 }
490 #[cfg(not(unix))]
491 {
492 if let Ok(user) = std::env::var("USER").or_else(|_| std::env::var("USERNAME")) {
493 key_material.extend_from_slice(user.as_bytes());
494 }
495 }
496
497 key_material.extend_from_slice(b"vtcode-credentials-v1");
498
499 let hash = digest(&SHA256, &key_material);
500 let key_bytes: &[u8; 32] = hash.as_ref()[..32]
501 .try_into()
502 .context("credential encryption key was too short")?;
503 let unbound = UnboundKey::new(&aead::AES_256_GCM, key_bytes)
504 .map_err(|_| anyhow!("invalid credential encryption key"))?;
505 Ok(LessSafeKey::new(unbound))
506}
507
508fn load_legacy_auth_file_for_provider(provider: &str) -> Result<Option<LegacyAuthFile>> {
509 let path = legacy_auth_storage_path()?;
510 let data = match fs::read(&path) {
511 Ok(data) => data,
512 Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
513 Err(err) => return Err(anyhow!("failed to read legacy auth file: {err}")),
514 };
515
516 let legacy: LegacyAuthFile =
517 serde_json::from_slice(&data).context("failed to parse legacy auth file")?;
518 let matches_provider = legacy.provider.eq_ignore_ascii_case(provider);
519 let stores_api_key = legacy.mode.eq_ignore_ascii_case("api_key");
520 let has_key = !legacy.api_key.trim().is_empty();
521
522 if matches_provider && stores_api_key && has_key {
523 Ok(Some(legacy))
524 } else {
525 Ok(None)
526 }
527}
528
529fn clear_legacy_auth_file_if_matches(provider: &str) -> Result<()> {
530 let path = legacy_auth_storage_path()?;
531 let Some(_legacy) = load_legacy_auth_file_for_provider(provider)? else {
532 return Ok(());
533 };
534
535 match fs::remove_file(path) {
536 Ok(()) => Ok(()),
537 Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
538 Err(err) => Err(anyhow!("failed to delete legacy auth file: {err}")),
539 }
540}
541
542pub fn migrate_custom_api_keys_to_keyring(
555 custom_api_keys: &BTreeMap<String, String>,
556 mode: AuthCredentialsStoreMode,
557) -> Result<BTreeMap<String, bool>> {
558 let mut migration_results = BTreeMap::new();
559
560 for (provider, api_key) in custom_api_keys {
561 let storage = CustomApiKeyStorage::new(provider);
562 match storage.store(api_key, mode) {
563 Ok(()) => {
564 tracing::info!(
565 "Migrated API key for provider '{}' to secure storage",
566 provider
567 );
568 migration_results.insert(provider.clone(), true);
569 }
570 Err(e) => {
571 tracing::warn!(
572 "Failed to migrate API key for provider '{}': {}",
573 provider,
574 e
575 );
576 migration_results.insert(provider.clone(), false);
577 }
578 }
579 }
580
581 Ok(migration_results)
582}
583
584pub fn load_custom_api_keys(
595 providers: &[String],
596 mode: AuthCredentialsStoreMode,
597) -> Result<BTreeMap<String, String>> {
598 let mut api_keys = BTreeMap::new();
599
600 for provider in providers {
601 let storage = CustomApiKeyStorage::new(provider);
602 if let Some(key) = storage.load(mode)? {
603 api_keys.insert(provider.clone(), key);
604 }
605 }
606
607 Ok(api_keys)
608}
609
610pub fn clear_custom_api_keys(providers: &[String], mode: AuthCredentialsStoreMode) -> Result<()> {
616 for provider in providers {
617 let storage = CustomApiKeyStorage::new(provider);
618 if let Err(e) = storage.clear(mode) {
619 tracing::warn!("Failed to clear API key for provider '{}': {}", provider, e);
620 }
621 }
622 Ok(())
623}
624
625#[cfg(test)]
626mod tests {
627 use super::*;
628 use assert_fs::TempDir;
629 use serial_test::serial;
630
631 struct TestAuthDirGuard {
632 temp_dir: Option<TempDir>,
633 previous: Option<std::path::PathBuf>,
634 }
635
636 impl TestAuthDirGuard {
637 fn new() -> Self {
638 let temp_dir = TempDir::new().expect("create temp auth dir");
639 let previous = crate::storage_paths::auth_storage_dir_override_for_tests()
640 .expect("read auth dir override");
641 crate::storage_paths::set_auth_storage_dir_override_for_tests(Some(
642 temp_dir.path().to_path_buf(),
643 ))
644 .expect("set auth dir override");
645
646 Self {
647 temp_dir: Some(temp_dir),
648 previous,
649 }
650 }
651 }
652
653 impl Drop for TestAuthDirGuard {
654 fn drop(&mut self) {
655 crate::storage_paths::set_auth_storage_dir_override_for_tests(self.previous.clone())
656 .expect("restore auth dir override");
657 if let Some(temp_dir) = self.temp_dir.take() {
658 temp_dir.close().expect("remove temp auth dir");
659 }
660 }
661 }
662
663 #[test]
664 fn test_storage_mode_default_is_keyring() {
665 assert_eq!(
666 AuthCredentialsStoreMode::default(),
667 AuthCredentialsStoreMode::Keyring
668 );
669 }
670
671 #[test]
672 fn test_storage_mode_effective_mode() {
673 assert_eq!(
674 AuthCredentialsStoreMode::Keyring.effective_mode(),
675 AuthCredentialsStoreMode::Keyring
676 );
677 assert_eq!(
678 AuthCredentialsStoreMode::File.effective_mode(),
679 AuthCredentialsStoreMode::File
680 );
681
682 let auto_mode = AuthCredentialsStoreMode::Auto.effective_mode();
684 assert!(
685 auto_mode == AuthCredentialsStoreMode::Keyring
686 || auto_mode == AuthCredentialsStoreMode::File
687 );
688 }
689
690 #[test]
691 fn test_storage_mode_serialization() {
692 let keyring_json = serde_json::to_string(&AuthCredentialsStoreMode::Keyring).unwrap();
693 assert_eq!(keyring_json, "\"keyring\"");
694
695 let file_json = serde_json::to_string(&AuthCredentialsStoreMode::File).unwrap();
696 assert_eq!(file_json, "\"file\"");
697
698 let auto_json = serde_json::to_string(&AuthCredentialsStoreMode::Auto).unwrap();
699 assert_eq!(auto_json, "\"auto\"");
700
701 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"keyring\"").unwrap();
703 assert_eq!(parsed, AuthCredentialsStoreMode::Keyring);
704
705 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"file\"").unwrap();
706 assert_eq!(parsed, AuthCredentialsStoreMode::File);
707
708 let parsed: AuthCredentialsStoreMode = serde_json::from_str("\"auto\"").unwrap();
709 assert_eq!(parsed, AuthCredentialsStoreMode::Auto);
710 }
711
712 #[test]
713 fn test_credential_storage_new() {
714 let storage = CredentialStorage::new("vtcode", "test_key");
715 assert_eq!(storage.service, "vtcode");
716 assert_eq!(storage.user, "test_key");
717 }
718
719 #[test]
720 fn test_is_keyring_functional_check() {
721 let _functional = is_keyring_functional();
724 }
725
726 #[test]
727 #[serial]
728 fn credential_storage_file_mode_round_trips_without_plaintext() {
729 let _guard = TestAuthDirGuard::new();
730 let storage = CredentialStorage::new("vtcode", "test_key");
731
732 storage
733 .store_with_mode("secret_api_key", AuthCredentialsStoreMode::File)
734 .expect("store encrypted credential");
735
736 let loaded = storage
737 .load_with_mode(AuthCredentialsStoreMode::File)
738 .expect("load encrypted credential");
739 assert_eq!(loaded.as_deref(), Some("secret_api_key"));
740
741 let stored = fs::read_to_string(storage.file_path().expect("credential path"))
742 .expect("read encrypted credential file");
743 assert!(!stored.contains("secret_api_key"));
744 }
745
746 #[test]
747 #[serial]
748 fn keyring_mode_load_falls_back_to_encrypted_file() {
749 let _guard = TestAuthDirGuard::new();
750 let storage = CredentialStorage::new("vtcode", "test_key");
751
752 storage
753 .store_with_mode("secret_api_key", AuthCredentialsStoreMode::File)
754 .expect("store encrypted credential");
755
756 let loaded = storage
757 .load_with_mode(AuthCredentialsStoreMode::Keyring)
758 .expect("load credential");
759 assert_eq!(loaded.as_deref(), Some("secret_api_key"));
760 }
761
762 #[test]
763 #[serial]
764 fn custom_api_key_load_migrates_legacy_auth_json() {
765 let _guard = TestAuthDirGuard::new();
766 let legacy_path = legacy_auth_storage_path().expect("legacy auth path");
767 fs::write(
768 &legacy_path,
769 r#"{
770 "version": 1,
771 "mode": "api_key",
772 "provider": "openai",
773 "api_key": "legacy-secret",
774 "authenticated_at": 1768406185
775}"#,
776 )
777 .expect("write legacy auth file");
778
779 let storage = CustomApiKeyStorage::new("openai");
780 let loaded = storage
781 .load(AuthCredentialsStoreMode::File)
782 .expect("load migrated api key");
783 assert_eq!(loaded.as_deref(), Some("legacy-secret"));
784 assert!(!legacy_path.exists());
785
786 let encrypted = fs::read_to_string(storage.storage.file_path().expect("credential path"))
787 .expect("read migrated credential file");
788 assert!(!encrypted.contains("legacy-secret"));
789 }
790}