yapi_ai/auth/
store.rs

1//! `auth.json`: stored credentials keyed by provider id.
2//!
3//! Port of `core/auth-storage.ts` in pi `v1.0.0`. Reads use a cached
4//! document that reloads when the file's revision changes. Changes re-read
5//! the file under the lock and write it back as pi does, so other entries keep
6//! their bytes.
7
8use std::future::Future;
9use std::path::{Path, PathBuf};
10use std::sync::{Mutex, MutexGuard};
11use std::time::SystemTime;
12
13use serde_json::{Map, Value};
14use tokio_util::sync::CancellationToken;
15use yapi_types::auth::Credential;
16use yapi_types::config::ConfigFile;
17
18use super::AuthError;
19use super::lock::{FileLock, write_atomic};
20
21type Document = Map<String, Value>;
22type Revision = (Option<SystemTime>, u64);
23
24#[derive(Default)]
25struct State {
26    document: Document,
27    revision: Option<Revision>,
28    /// Why the file at this revision could not be read, if it could not.
29    error: Option<String>,
30}
31
32/// The credential store: `auth.json` in the agent directory, or memory.
33#[derive(Default)]
34pub struct CredentialStore {
35    path: Option<PathBuf>,
36    state: Mutex<State>,
37}
38
39impl std::fmt::Debug for CredentialStore {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.debug_struct("CredentialStore")
42            .field("path", &self.path)
43            .finish_non_exhaustive()
44    }
45}
46
47/// Which kind of credential is stored.
48#[derive(Clone, Copy, Debug, PartialEq, Eq)]
49pub enum CredentialKind {
50    /// `"api_key"`.
51    ApiKey,
52    /// `"oauth"`.
53    OAuth,
54}
55
56fn revision(path: &Path) -> Option<Revision> {
57    let metadata = std::fs::metadata(path).ok()?;
58    Some((metadata.modified().ok(), metadata.len()))
59}
60
61fn parse(text: &str) -> Result<Document, String> {
62    let text = text.strip_prefix('\u{feff}').unwrap_or(text);
63    if text.trim().is_empty() {
64        return Ok(Document::new());
65    }
66    match serde_json::from_str::<Value>(text) {
67        Ok(Value::Object(document)) => Ok(document),
68        Ok(_) => Err("Invalid auth.json: expected an object".into()),
69        Err(err) => Err(err.to_string()),
70    }
71}
72
73fn failed(err: impl std::fmt::Display) -> AuthError {
74    AuthError::Failed(err.to_string())
75}
76
77impl CredentialStore {
78    /// The store backed by `path`, loaded now. An unreadable file reads as empty.
79    pub fn open(path: impl Into<PathBuf>) -> CredentialStore {
80        let store = CredentialStore {
81            path: Some(path.into()),
82            state: Mutex::default(),
83        };
84        drop(store.state());
85        store
86    }
87
88    /// The file behind the store, if any.
89    pub fn path(&self) -> Option<&Path> {
90        self.path.as_deref()
91    }
92
93    /// The cached state, reloaded first when the file changed. A file that no
94    /// longer parses keeps the last good document, as pi does.
95    fn state(&self) -> MutexGuard<'_, State> {
96        let mut state = self
97            .state
98            .lock()
99            .unwrap_or_else(std::sync::PoisonError::into_inner);
100        if let Some(path) = &self.path {
101            let current = revision(path);
102            if current.is_none() || current != state.revision {
103                state.error = None;
104                match std::fs::read_to_string(path) {
105                    Ok(text) => match parse(&text) {
106                        Ok(document) => state.document = document,
107                        Err(error) => state.error = Some(error),
108                    },
109                    Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
110                        state.document = Document::new();
111                    }
112                    Err(error) => state.error = Some(error.to_string()),
113                }
114                state.revision = current;
115            }
116        }
117        state
118    }
119
120    /// Why the file cannot be read now, as pi's store reports it. Reads
121    /// meanwhile see the last document that could be.
122    pub fn read_error(&self) -> Option<String> {
123        self.state().error.clone()
124    }
125
126    /// The stored credential for `provider`, if it is a valid one.
127    pub fn get(&self, provider: &str) -> Option<Credential> {
128        let value = self.state().document.get(provider).cloned()?;
129        serde_json::from_value(value).ok()
130    }
131
132    /// Stored providers and their credential kinds, in file order.
133    pub fn list(&self) -> Vec<(String, CredentialKind)> {
134        self.state()
135            .document
136            .iter()
137            .filter_map(|(provider, value)| {
138                let kind = match value.get("type").and_then(Value::as_str)? {
139                    "api_key" => CredentialKind::ApiKey,
140                    "oauth" => CredentialKind::OAuth,
141                    _ => return None,
142                };
143                Some((provider.clone(), kind))
144            })
145            .collect()
146    }
147
148    /// Replaces `provider`'s credential with what `change` returns for the
149    /// current one, under the lock. `None` leaves the file untouched. Returns the
150    /// credential now stored.
151    pub async fn modify<F, Fut>(
152        &self,
153        provider: &str,
154        change: F,
155        cancel: &CancellationToken,
156    ) -> Result<Option<Credential>, AuthError>
157    where
158        F: FnOnce(Option<Credential>) -> Fut,
159        Fut: Future<Output = Result<Option<Credential>, AuthError>>,
160    {
161        let Some(path) = self.path.clone() else {
162            let current = self.get(provider);
163            let next = change(current.clone()).await?;
164            if let Some(next) = &next {
165                let value = serde_json::to_value(next).map_err(failed)?;
166                self.state().document.insert(provider.to_owned(), value);
167            }
168            return Ok(next.or(current));
169        };
170        prepare(&path)?;
171        let _lock = FileLock::acquire(&path, cancel).await.map_err(|err| {
172            if cancel.is_cancelled() {
173                AuthError::Cancelled
174            } else {
175                failed(err)
176            }
177        })?;
178        let mut document = read(&path)?;
179        let current = document
180            .get(provider)
181            .cloned()
182            .and_then(|value| serde_json::from_value::<Credential>(value).ok());
183        let Some(next) = change(current.clone()).await? else {
184            self.remember(&path, document);
185            return Ok(current);
186        };
187        if cancel.is_cancelled() {
188            return Err(AuthError::Cancelled);
189        }
190        let value = serde_json::to_value(&next).map_err(failed)?;
191        document.insert(provider.to_owned(), value);
192        write(&path, &document)?;
193        self.remember(&path, document);
194        Ok(Some(next))
195    }
196
197    /// Stores `credential` for `provider`.
198    pub async fn set(
199        &self,
200        provider: &str,
201        credential: Credential,
202        cancel: &CancellationToken,
203    ) -> Result<(), AuthError> {
204        self.modify(provider, |_| async { Ok(Some(credential)) }, cancel)
205            .await
206            .map(drop)
207    }
208
209    /// Removes `provider`'s credential.
210    pub async fn delete(
211        &self,
212        provider: &str,
213        cancel: &CancellationToken,
214    ) -> Result<(), AuthError> {
215        let Some(path) = self.path.clone() else {
216            self.state().document.shift_remove(provider);
217            return Ok(());
218        };
219        prepare(&path)?;
220        let _lock = FileLock::acquire(&path, cancel).await.map_err(failed)?;
221        let mut document = read(&path)?;
222        document.shift_remove(provider);
223        write(&path, &document)?;
224        self.remember(&path, document);
225        Ok(())
226    }
227
228    fn remember(&self, path: &Path, document: Document) {
229        let mut state = self
230            .state
231            .lock()
232            .unwrap_or_else(std::sync::PoisonError::into_inner);
233        state.document = document;
234        state.revision = revision(path);
235    }
236}
237
238/// pi's `FileAuthStorageBackend.withLock` for another credential file, such
239/// as `mcp-auth.json`: creates the directory (mode 0700) and the file (mode
240/// 0600, `{}`) when missing, locks the file, passes its text to `change`, and
241/// writes the text `change` returns, if any.
242pub async fn with_file_lock<T>(
243    path: &Path,
244    change: impl FnOnce(&str) -> Result<(T, Option<String>), AuthError>,
245) -> Result<T, AuthError> {
246    prepare(path)?;
247    let _lock = FileLock::acquire(path, &CancellationToken::new())
248        .await
249        .map_err(failed)?;
250    let current = match std::fs::read_to_string(path) {
251        Ok(text) => text,
252        Err(err) if err.kind() == std::io::ErrorKind::NotFound => String::new(),
253        Err(err) => return Err(failed(err)),
254    };
255    let (result, next) = change(&current)?;
256    if let Some(next) = next {
257        write_text(path, &next)?;
258    }
259    Ok(result)
260}
261
262/// Creates the directory (mode 0700) and an empty file (mode 0600) when missing.
263fn prepare(path: &Path) -> Result<(), AuthError> {
264    if let Some(dir) = path.parent()
265        && !dir.exists()
266    {
267        let mut builder = std::fs::DirBuilder::new();
268        builder.recursive(true);
269        #[cfg(unix)]
270        std::os::unix::fs::DirBuilderExt::mode(&mut builder, 0o700);
271        builder.create(dir).map_err(failed)?;
272    }
273    if !path.exists() {
274        write_text(path, "{}")?;
275    }
276    Ok(())
277}
278
279fn read(path: &Path) -> Result<Document, AuthError> {
280    match std::fs::read_to_string(path) {
281        Ok(text) => parse(&text).map_err(AuthError::Failed),
282        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(Document::new()),
283        Err(err) => Err(failed(err)),
284    }
285}
286
287fn write(path: &Path, document: &Document) -> Result<(), AuthError> {
288    let text = ConfigFile::Auth.render(document).map_err(failed)?;
289    write_text(path, &text)
290}
291
292/// Writes with mode 0600 when creating. An existing file keeps its mode.
293fn write_text(path: &Path, text: &str) -> Result<(), AuthError> {
294    write_atomic(path, text.as_bytes(), 0o600).map_err(failed)
295}
296
297#[cfg(test)]
298mod tests {
299    use yapi_types::auth::ApiKeyCredential;
300
301    use super::*;
302
303    fn temp(name: &str) -> PathBuf {
304        let dir = std::env::temp_dir().join(format!("yapi-store-{name}-{}", std::process::id()));
305        let _ = std::fs::remove_dir_all(&dir);
306        dir.join("agent").join("auth.json")
307    }
308
309    fn key(value: &str) -> Credential {
310        Credential::ApiKey(ApiKeyCredential {
311            key: Some(value.into()),
312            env: None,
313        })
314    }
315
316    #[test]
317    fn reports_unparsable_file() {
318        let path = temp("corrupt");
319        std::fs::create_dir_all(path.parent().unwrap()).unwrap();
320        std::fs::write(&path, "{bad").unwrap();
321        let store = CredentialStore::open(&path);
322        assert!(store.read_error().is_some());
323        assert_eq!(store.get("anthropic"), None);
324        std::fs::write(
325            &path,
326            "{\"anthropic\": {\"type\": \"api_key\", \"key\": \"k\"}}",
327        )
328        .unwrap();
329        // A rewrite of the same length may share the old modification time.
330        let store = CredentialStore::open(&path);
331        assert_eq!(store.read_error(), None);
332        assert_eq!(store.get("anthropic"), Some(key("k")));
333        let _ = std::fs::remove_dir_all(path.parent().unwrap().parent().unwrap());
334    }
335
336    #[tokio::test]
337    async fn other_entries_byte_identical() {
338        let path = temp("bytes");
339        std::fs::create_dir_all(path.parent().unwrap()).unwrap();
340        let original = "{\n  \"zeta\": {\n    \"type\": \"oauth\",\n    \"refresh\": \"r\",\n    \"access\": \"a\",\n    \"expires\": 1700000000000.5,\n    \"accountId\": \"x\"\n  },\n  \"alpha\": {\n    \"type\": \"api_key\",\n    \"key\": \"$KEY\"\n  }\n}";
341        std::fs::write(&path, original).unwrap();
342        let store = CredentialStore::open(&path);
343        assert_eq!(
344            store.list(),
345            [
346                ("zeta".to_owned(), CredentialKind::OAuth),
347                ("alpha".to_owned(), CredentialKind::ApiKey)
348            ]
349        );
350        let cancel = CancellationToken::new();
351        store.set("beta", key("k"), &cancel).await.unwrap();
352        let text = std::fs::read_to_string(&path).unwrap();
353        assert_eq!(
354            text,
355            original.replace(
356                "\n}",
357                ",\n  \"beta\": {\n    \"type\": \"api_key\",\n    \"key\": \"k\"\n  }\n}"
358            )
359        );
360        store.delete("beta", &cancel).await.unwrap();
361        assert_eq!(std::fs::read_to_string(&path).unwrap(), original);
362        // pi accepts any finite expiry, and a rewrite keeps its bytes.
363        let Some(Credential::OAuth(zeta)) = store.get("zeta") else {
364            panic!("zeta is not an OAuth credential");
365        };
366        assert_eq!(zeta.expires, 1_700_000_000_000.5);
367        store
368            .set("zeta", Credential::OAuth(zeta), &cancel)
369            .await
370            .unwrap();
371        let text = std::fs::read_to_string(&path).unwrap();
372        assert!(text.contains("\"expires\": 1700000000000.5,"), "{text}");
373        assert!(!path.with_file_name("auth.json.lock").exists());
374        std::fs::remove_dir_all(path.parent().unwrap().parent().unwrap()).unwrap();
375    }
376
377    #[tokio::test]
378    async fn external_edits_and_creation() {
379        let path = temp("reload");
380        let store = CredentialStore::open(&path);
381        assert!(store.get("openai").is_none());
382        let cancel = CancellationToken::new();
383        store.set("openai", key("one"), &cancel).await.unwrap();
384        #[cfg(unix)]
385        {
386            use std::os::unix::fs::PermissionsExt;
387            let mode = std::fs::metadata(&path).unwrap().permissions().mode();
388            assert_eq!(mode & 0o777, 0o600);
389        }
390        std::fs::write(
391            &path,
392            "{\"openai\": {\"type\": \"api_key\", \"key\": \"two-longer\"}}",
393        )
394        .unwrap();
395        assert_eq!(store.get("openai"), Some(key("two-longer")));
396        let unchanged = store
397            .modify("openai", |_| async { Ok(None) }, &cancel)
398            .await
399            .unwrap();
400        assert_eq!(unchanged, Some(key("two-longer")));
401        std::fs::remove_dir_all(path.parent().unwrap().parent().unwrap()).unwrap();
402    }
403}
yapi_ai/auth/store.rs

yapi_ai/auth/
store.rs
