Skip to main content

aether_auth/
memory.rs

1use std::collections::HashMap;
2use std::sync::Mutex;
3
4use async_trait::async_trait;
5
6use crate::credential::{OAuthCredential, OAuthCredentialStorage};
7use crate::error::OAuthError;
8
9#[derive(Default)]
10pub struct MemoryCredentialStore {
11    values: Mutex<HashMap<String, serde_json::Value>>,
12}
13
14impl MemoryCredentialStore {
15    pub fn new() -> Self {
16        Self::default()
17    }
18
19    pub fn with_credential(self, key: &str, credential: OAuthCredential) -> Self {
20        self.values.lock().unwrap().insert(key.to_string(), serde_json::to_value(credential).unwrap());
21        self
22    }
23
24    pub fn with_value(self, key: &str, value: serde_json::Value) -> Self {
25        self.values.lock().unwrap().insert(key.to_string(), value);
26        self
27    }
28}
29
30#[async_trait]
31impl OAuthCredentialStorage for MemoryCredentialStore {
32    async fn load(&self, key: &str) -> Result<Option<serde_json::Value>, OAuthError> {
33        Ok(self.values.lock().unwrap().get(key).cloned())
34    }
35
36    async fn save(&self, key: &str, value: serde_json::Value) -> Result<(), OAuthError> {
37        self.values.lock().unwrap().insert(key.to_string(), value);
38        Ok(())
39    }
40
41    async fn delete(&self, key: &str) -> Result<(), OAuthError> {
42        self.values.lock().unwrap().remove(key);
43        Ok(())
44    }
45
46    fn contains(&self, key: &str) -> bool {
47        self.values.lock().unwrap().contains_key(key)
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54    use crate::credential::assert_storage_contract;
55    use serde_json::json;
56
57    #[tokio::test]
58    async fn satisfies_the_storage_contract() {
59        assert_storage_contract(&MemoryCredentialStore::new()).await;
60    }
61
62    #[tokio::test]
63    async fn builders_seed_the_store() {
64        let credential = OAuthCredential::new("client", "token");
65        let store =
66            MemoryCredentialStore::new().with_credential("codex", credential.clone()).with_value("mcp:slack", json!(1));
67
68        assert_eq!(store.load_credential("codex").await.unwrap(), Some(credential));
69        assert_eq!(store.load("mcp:slack").await.unwrap(), Some(json!(1)));
70    }
71}