1use std::fs::{File, OpenOptions, TryLockError};
13use std::future::Future;
14use std::io::Write;
15use std::path::{Path, PathBuf};
16use std::sync::{Arc, Mutex, PoisonError, RwLock};
17use std::time::{Duration, Instant, SystemTime};
18
19use indexmap::IndexMap;
20use oino_ai::{JsonObject, JsonValue};
21
22use super::credential::{AuthCredential, CredentialInfo};
23use crate::config_value::resolve_config_value;
24use crate::env::{EnvSource, ProcessEnv};
25use crate::json_text::strip_bom;
26
27const LOCK_TIMEOUT: Duration = Duration::from_secs(30);
29
30#[derive(Debug, thiserror::Error)]
32pub enum AuthStorageError {
33 #[error("Failed to read auth.json: {0}")]
35 Read(String),
36 #[error("Invalid auth.json: expected an object")]
38 NotAnObject,
39 #[error("Invalid auth.json credential for provider \"{0}\"")]
41 InvalidCredential(String),
42 #[error("Read-only credential storage cannot modify auth.json")]
44 ReadOnly,
45 #[error("Failed to write auth.json: {0}")]
47 Write(#[source] std::io::Error),
48 #[error("Timed out waiting for auth.json lock: {}", .0.display())]
50 LockTimeout(PathBuf),
51}
52
53#[derive(Clone, Copy, Debug, PartialEq, Eq)]
54struct Revision {
55 modified: SystemTime,
56 len: u64,
57}
58
59fn revision(path: &Path) -> Option<Revision> {
60 let metadata = std::fs::metadata(path).ok()?;
61 Some(Revision {
62 modified: metadata.modified().ok()?,
63 len: metadata.len(),
64 })
65}
66
67#[derive(Debug)]
68enum Backend {
69 File { path: PathBuf, read_only: bool },
70 Memory(Mutex<JsonObject>),
71}
72
73#[derive(Debug, Default)]
74struct Snapshot {
75 data: JsonObject,
76 revision: Option<Revision>,
77 loaded: bool,
78}
79
80pub struct AuthStorage {
83 backend: Backend,
84 runtime: RwLock<IndexMap<String, String>>,
85 snapshot: Mutex<Snapshot>,
86 async_gate: futures::lock::Mutex<()>,
87 env: Arc<dyn EnvSource>,
88}
89
90impl std::fmt::Debug for AuthStorage {
91 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92 f.debug_struct("AuthStorage")
93 .field("backend", &self.backend)
94 .finish_non_exhaustive()
95 }
96}
97
98struct LockGuard {
100 file: File,
101}
102
103impl Drop for LockGuard {
104 fn drop(&mut self) {
105 let _ = self.file.unlock();
107 }
108}
109
110fn parse_data(content: &str) -> Result<JsonObject, AuthStorageError> {
111 let content = strip_bom(content);
112 if content.is_empty() {
113 return Ok(JsonObject::new());
114 }
115 match serde_json::from_str::<JsonValue>(content) {
116 Ok(JsonValue::Object(map)) => Ok(map),
117 Ok(_) => Err(AuthStorageError::NotAnObject),
118 Err(error) => Err(AuthStorageError::Read(error.to_string())),
119 }
120}
121
122fn validate_entries(data: &JsonObject) -> Result<(), AuthStorageError> {
123 for (provider, value) in data {
124 if AuthCredential::from_value(value).is_none() {
125 return Err(AuthStorageError::InvalidCredential(provider.clone()));
126 }
127 }
128 Ok(())
129}
130
131fn serialize_data(data: &JsonObject) -> String {
132 serde_json::to_string_pretty(data).unwrap_or_else(|_| "{}".to_owned())
133}
134
135impl AuthStorage {
136 fn with_backend(backend: Backend) -> Self {
137 Self {
138 backend,
139 runtime: RwLock::new(IndexMap::new()),
140 snapshot: Mutex::new(Snapshot::default()),
141 async_gate: futures::lock::Mutex::new(()),
142 env: Arc::new(ProcessEnv),
143 }
144 }
145
146 #[must_use]
148 pub fn open(path: impl Into<PathBuf>) -> Self {
149 Self::with_backend(Backend::File {
150 path: path.into(),
151 read_only: false,
152 })
153 }
154
155 #[must_use]
157 pub fn open_read_only(path: impl Into<PathBuf>) -> Self {
158 Self::with_backend(Backend::File {
159 path: path.into(),
160 read_only: true,
161 })
162 }
163
164 #[must_use]
166 pub fn in_memory(data: IndexMap<String, AuthCredential>) -> Self {
167 let object = data
168 .into_iter()
169 .map(|(provider, credential)| (provider, credential.to_value()))
170 .collect();
171 Self::with_backend(Backend::Memory(Mutex::new(object)))
172 }
173
174 #[must_use]
176 pub fn with_env(mut self, env: Arc<dyn EnvSource>) -> Self {
177 self.env = env;
178 self
179 }
180
181 #[must_use]
183 pub fn path(&self) -> Option<&Path> {
184 match &self.backend {
185 Backend::File { path, .. } => Some(path),
186 Backend::Memory(_) => None,
187 }
188 }
189
190 #[must_use]
192 pub fn is_read_only(&self) -> bool {
193 matches!(
194 self.backend,
195 Backend::File {
196 read_only: true,
197 ..
198 }
199 )
200 }
201
202 fn read_file(path: &Path, read_only: bool) -> Result<JsonObject, AuthStorageError> {
203 let content = match std::fs::read_to_string(path) {
204 Ok(content) => content,
205 Err(error) if error.kind() == std::io::ErrorKind::NotFound => String::new(),
206 Err(error) => return Err(AuthStorageError::Read(error.to_string())),
207 };
208 let data = parse_data(&content)?;
209 if read_only {
210 validate_entries(&data)?;
211 }
212 Ok(data)
213 }
214
215 fn data(&self) -> Result<JsonObject, AuthStorageError> {
220 match &self.backend {
221 Backend::Memory(data) => {
222 Ok(data.lock().unwrap_or_else(PoisonError::into_inner).clone())
223 }
224 Backend::File { path, read_only } => {
225 let current = revision(path);
226 let mut snapshot = self.snapshot.lock().unwrap_or_else(PoisonError::into_inner);
227 if snapshot.loaded && current.is_some() && snapshot.revision == current {
228 return Ok(snapshot.data.clone());
229 }
230 match Self::read_file(path, *read_only) {
231 Ok(data) => {
232 snapshot.data = data;
233 snapshot.revision = current;
234 snapshot.loaded = true;
235 Ok(snapshot.data.clone())
236 }
237 Err(error) if !*read_only => {
238 tracing::warn!(%error, path = %path.display(), "keeping last auth.json snapshot");
239 Ok(snapshot.data.clone())
240 }
241 Err(error) => Err(error),
242 }
243 }
244 }
245 }
246
247 pub fn reload(&self) {
249 let mut snapshot = self.snapshot.lock().unwrap_or_else(PoisonError::into_inner);
250 snapshot.revision = None;
251 snapshot.loaded = false;
252 }
253
254 pub fn stored(&self, provider: &str) -> Result<Option<AuthCredential>, AuthStorageError> {
261 let data = self.data()?;
262 Ok(data.get(provider).and_then(AuthCredential::from_value))
263 }
264
265 pub fn get(&self, provider: &str) -> Result<Option<AuthCredential>, AuthStorageError> {
273 if let Some(key) = self.runtime_api_key(provider) {
274 return Ok(Some(AuthCredential::api_key(key)));
275 }
276 Ok(self.stored(provider)?.map(|credential| match credential {
277 AuthCredential::ApiKey(mut api_key) => {
278 if let Some(key) = &api_key.key {
279 api_key.key =
280 resolve_config_value(key, api_key.env.as_ref(), self.env.as_ref());
281 }
282 AuthCredential::ApiKey(api_key)
283 }
284 oauth @ AuthCredential::Oauth(_) => oauth,
285 }))
286 }
287
288 pub fn list(&self) -> Result<Vec<CredentialInfo>, AuthStorageError> {
294 let mut entries: IndexMap<String, String> = self
295 .data()?
296 .iter()
297 .map(|(provider, value)| {
298 let kind = value
299 .get("type")
300 .and_then(JsonValue::as_str)
301 .unwrap_or_default()
302 .to_owned();
303 (provider.clone(), kind)
304 })
305 .collect();
306 for provider in self.read_runtime().keys() {
307 entries.insert(provider.clone(), "api_key".to_owned());
308 }
309 Ok(entries
310 .into_iter()
311 .map(|(provider_id, kind)| CredentialInfo { provider_id, kind })
312 .collect())
313 }
314
315 #[must_use]
317 pub fn has(&self, provider: &str) -> bool {
318 self.has_runtime_api_key(provider)
319 || self.data().is_ok_and(|data| data.contains_key(provider))
320 }
321
322 #[must_use]
324 pub fn has_stored(&self, provider: &str) -> bool {
325 self.data().is_ok_and(|data| data.contains_key(provider))
326 }
327
328 pub fn set(&self, provider: &str, credential: &AuthCredential) -> Result<(), AuthStorageError> {
334 self.write_with(|data| {
335 data.insert(provider.to_owned(), credential.to_value());
336 true
337 })
338 }
339
340 pub fn remove(&self, provider: &str) -> Result<(), AuthStorageError> {
346 self.write_with(|data| {
347 data.shift_remove(provider);
348 true
349 })?;
350 self.write_runtime().shift_remove(provider);
351 Ok(())
352 }
353
354 pub fn modify(
361 &self,
362 provider: &str,
363 f: impl FnOnce(Option<AuthCredential>) -> Option<AuthCredential>,
364 ) -> Result<Option<AuthCredential>, AuthStorageError> {
365 let mut result = None;
366 self.write_with(|data| {
367 let current = data.get(provider).and_then(AuthCredential::from_value);
368 if let Some(next) = f(current.clone()) {
369 data.insert(provider.to_owned(), next.to_value());
370 result = Some(next);
371 true
372 } else {
373 result = current;
374 false
375 }
376 })?;
377 Ok(result)
378 }
379
380 pub async fn modify_async<F, Fut, E>(
389 &self,
390 provider: &str,
391 f: F,
392 ) -> Result<Option<AuthCredential>, E>
393 where
394 F: FnOnce(Option<AuthCredential>) -> Fut,
395 Fut: Future<Output = Result<Option<AuthCredential>, E>>,
396 E: From<AuthStorageError>,
397 {
398 let _gate = self.async_gate.lock().await;
399 if self.is_read_only() {
400 return Err(AuthStorageError::ReadOnly.into());
401 }
402 let _lock = self.acquire_lock()?;
403 let mut data = self.read_locked()?;
404 let current = data.get(provider).and_then(AuthCredential::from_value);
405 match f(current.clone()).await? {
406 Some(next) => {
407 data.insert(provider.to_owned(), next.to_value());
408 self.write_locked(&data)?;
409 Ok(Some(next))
410 }
411 None => Ok(current),
412 }
413 }
414
415 pub fn set_runtime_api_key(&self, provider: &str, key: impl Into<String>) {
417 self.write_runtime().insert(provider.to_owned(), key.into());
418 }
419
420 pub fn remove_runtime_api_key(&self, provider: &str) {
422 self.write_runtime().shift_remove(provider);
423 }
424
425 #[must_use]
427 pub fn has_runtime_api_key(&self, provider: &str) -> bool {
428 self.read_runtime().contains_key(provider)
429 }
430
431 fn runtime_api_key(&self, provider: &str) -> Option<String> {
432 self.read_runtime().get(provider).cloned()
433 }
434
435 fn read_runtime(&self) -> std::sync::RwLockReadGuard<'_, IndexMap<String, String>> {
436 self.runtime.read().unwrap_or_else(PoisonError::into_inner)
437 }
438
439 fn write_runtime(&self) -> std::sync::RwLockWriteGuard<'_, IndexMap<String, String>> {
440 self.runtime.write().unwrap_or_else(PoisonError::into_inner)
441 }
442
443 fn write_with(
445 &self,
446 mutate: impl FnOnce(&mut JsonObject) -> bool,
447 ) -> Result<(), AuthStorageError> {
448 if self.is_read_only() {
449 return Err(AuthStorageError::ReadOnly);
450 }
451 let _lock = self.acquire_lock()?;
452 let mut data = self.read_locked()?;
453 if mutate(&mut data) {
454 self.write_locked(&data)?;
455 }
456 Ok(())
457 }
458
459 fn read_locked(&self) -> Result<JsonObject, AuthStorageError> {
460 match &self.backend {
461 Backend::Memory(data) => {
462 Ok(data.lock().unwrap_or_else(PoisonError::into_inner).clone())
463 }
464 Backend::File { path, .. } => Self::read_file(path, false),
465 }
466 }
467
468 fn write_locked(&self, data: &JsonObject) -> Result<(), AuthStorageError> {
469 match &self.backend {
470 Backend::Memory(stored) => {
471 stored
472 .lock()
473 .unwrap_or_else(PoisonError::into_inner)
474 .clone_from(data);
475 Ok(())
476 }
477 Backend::File { path, .. } => {
478 write_atomic(path, &serialize_data(data)).map_err(AuthStorageError::Write)?;
479 let mut snapshot = self.snapshot.lock().unwrap_or_else(PoisonError::into_inner);
480 snapshot.data.clone_from(data);
481 snapshot.revision = revision(path);
482 snapshot.loaded = true;
483 Ok(())
484 }
485 }
486 }
487
488 fn acquire_lock(&self) -> Result<Option<LockGuard>, AuthStorageError> {
489 let Backend::File { path, .. } = &self.backend else {
490 return Ok(None);
491 };
492 let lock_path = lock_path(path);
493 if let Some(parent) = lock_path.parent() {
494 create_private_dir(parent).map_err(AuthStorageError::Write)?;
495 }
496 let file = open_private(&lock_path, false).map_err(AuthStorageError::Write)?;
497 let deadline = Instant::now() + LOCK_TIMEOUT;
498 let mut delay = Duration::from_millis(10);
499 loop {
500 match file.try_lock() {
501 Ok(()) => return Ok(Some(LockGuard { file })),
502 Err(TryLockError::WouldBlock) if Instant::now() < deadline => {
503 std::thread::sleep(delay);
504 delay = (delay * 2).min(Duration::from_secs(1));
505 }
506 Err(TryLockError::WouldBlock) => {
507 return Err(AuthStorageError::LockTimeout(lock_path));
508 }
509 Err(TryLockError::Error(error)) => return Err(AuthStorageError::Write(error)),
510 }
511 }
512 }
513}
514
515fn lock_path(path: &Path) -> PathBuf {
516 let mut name = path
517 .file_name()
518 .map(std::ffi::OsStr::to_os_string)
519 .unwrap_or_default();
520 name.push(".lock");
521 path.with_file_name(name)
522}
523
524fn create_private_dir(dir: &Path) -> std::io::Result<()> {
525 if dir.as_os_str().is_empty() || dir.exists() {
526 return Ok(());
527 }
528 let mut builder = std::fs::DirBuilder::new();
529 builder.recursive(true);
530 #[cfg(unix)]
531 {
532 use std::os::unix::fs::DirBuilderExt;
533 builder.mode(0o700);
534 }
535 builder.create(dir)
536}
537
538fn open_private(path: &Path, truncate: bool) -> std::io::Result<File> {
539 let mut options = OpenOptions::new();
540 options
541 .read(true)
542 .write(true)
543 .create(true)
544 .truncate(truncate);
545 #[cfg(unix)]
546 {
547 use std::os::unix::fs::OpenOptionsExt;
548 options.mode(0o600);
549 }
550 options.open(path)
551}
552
553pub(crate) fn write_atomic(path: &Path, content: &str) -> std::io::Result<()> {
556 if let Some(parent) = path.parent() {
557 create_private_dir(parent)?;
558 }
559 let mut temp_name = path
560 .file_name()
561 .map(std::ffi::OsStr::to_os_string)
562 .unwrap_or_default();
563 temp_name.push(format!(".tmp-{}", std::process::id()));
564 let temp = path.with_file_name(temp_name);
565 {
566 let mut file = open_private(&temp, true)?;
567 if let Ok(metadata) = std::fs::metadata(path) {
568 file.set_permissions(metadata.permissions())?;
569 }
570 file.write_all(content.as_bytes())?;
571 file.sync_all()?;
572 }
573 std::fs::rename(&temp, path).inspect_err(|_| {
574 let _ = std::fs::remove_file(&temp);
575 })
576}
577
578#[cfg(test)]
579mod tests {
580 use super::*;
581 use crate::env::MapEnv;
582
583 fn temp_dir() -> tempfile::TempDir {
584 let root = std::env::var_os("OINO_TEST_TMP").map_or_else(
585 || PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../target/tmp"),
586 PathBuf::from,
587 );
588 std::fs::create_dir_all(&root).unwrap();
589 tempfile::tempdir_in(root).unwrap()
590 }
591
592 #[test]
593 fn file_round_trip_and_permissions() {
594 let dir = temp_dir();
595 let path = dir.path().join("agent/auth.json");
596 let storage = AuthStorage::open(&path);
597 assert_eq!(storage.get("a").unwrap(), None);
598 storage.set("a", &AuthCredential::api_key("k1")).unwrap();
599 let text = std::fs::read_to_string(&path).unwrap();
600 assert_eq!(
601 text,
602 "{\n \"a\": {\n \"type\": \"api_key\",\n \"key\": \"k1\"\n }\n}"
603 );
604 #[cfg(unix)]
605 {
606 use std::os::unix::fs::PermissionsExt;
607 let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777;
608 assert_eq!(mode, 0o600);
609 }
610 let other = AuthStorage::open(&path);
611 assert_eq!(other.get("a").unwrap(), Some(AuthCredential::api_key("k1")));
612 other.remove("a").unwrap();
613 assert!(!storage.has("a"));
614 assert_eq!(storage.list().unwrap(), vec![]);
615 }
616
617 #[test]
618 fn concurrent_writers_do_not_lose_updates() {
619 let dir = temp_dir();
620 let path = dir.path().join("auth.json");
621 let handles: Vec<_> = (0..8)
622 .map(|i| {
623 let path = path.clone();
624 std::thread::spawn(move || {
625 let storage = AuthStorage::open(&path);
626 storage
627 .set(&format!("p{i}"), &AuthCredential::api_key(format!("k{i}")))
628 .unwrap();
629 })
630 })
631 .collect();
632 for handle in handles {
633 handle.join().unwrap();
634 }
635 let storage = AuthStorage::open(&path);
636 assert_eq!(storage.list().unwrap().len(), 8);
637 }
638
639 #[test]
640 fn runtime_keys_overlay() {
641 let storage =
642 AuthStorage::in_memory([("a".to_owned(), AuthCredential::api_key("stored"))].into());
643 storage.set_runtime_api_key("a", "runtime");
644 storage.set_runtime_api_key("b", "rb");
645 assert_eq!(
646 storage.get("a").unwrap(),
647 Some(AuthCredential::api_key("runtime"))
648 );
649 let kinds: Vec<_> = storage
650 .list()
651 .unwrap()
652 .into_iter()
653 .map(|i| (i.provider_id, i.kind))
654 .collect();
655 assert_eq!(
656 kinds,
657 vec![
658 ("a".into(), "api_key".into()),
659 ("b".into(), "api_key".into())
660 ]
661 );
662 storage.remove("a").unwrap();
663 assert!(!storage.has("a"));
664 }
665
666 #[test]
667 fn resolves_env_keys() {
668 let storage =
669 AuthStorage::in_memory([("a".to_owned(), AuthCredential::api_key("$KEY"))].into())
670 .with_env(Arc::new(MapEnv::new().with("KEY", "secret")));
671 assert_eq!(
672 storage.get("a").unwrap(),
673 Some(AuthCredential::api_key("secret"))
674 );
675 assert_eq!(
676 storage.stored("a").unwrap(),
677 Some(AuthCredential::api_key("$KEY"))
678 );
679 }
680
681 #[test]
682 fn read_only_validates() {
683 let dir = temp_dir();
684 let path = dir.path().join("auth.json");
685 std::fs::write(&path, r#"{"a": {"type": "api_key", "key": 5}}"#).unwrap();
686 let storage = AuthStorage::open_read_only(&path);
687 assert_eq!(
688 storage.get("a").unwrap_err().to_string(),
689 "Invalid auth.json credential for provider \"a\""
690 );
691 assert!(matches!(
692 storage.set("b", &AuthCredential::api_key("x")),
693 Err(AuthStorageError::ReadOnly)
694 ));
695 std::fs::write(&path, "[]").unwrap();
696 storage.reload();
697 assert_eq!(
698 storage.list().unwrap_err().to_string(),
699 "Invalid auth.json: expected an object"
700 );
701 }
702
703 #[test]
704 fn modify_async_serializes() {
705 let storage = AuthStorage::in_memory(IndexMap::new());
706 let result: Result<_, AuthStorageError> =
707 futures::executor::block_on(storage.modify_async("a", |current| async move {
708 assert!(current.is_none());
709 Ok(Some(AuthCredential::api_key("x")))
710 }));
711 assert_eq!(result.unwrap(), Some(AuthCredential::api_key("x")));
712 assert!(storage.has("a"));
713 }
714}