Skip to main content

elph_ai/auth/
credential_store.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use tokio::sync::Mutex;
7
8use super::types::{Credential, CredentialStore};
9
10/// In-memory credential store with per-provider serialized writes.
11pub struct InMemoryCredentialStore {
12    credentials: Mutex<HashMap<String, Credential>>,
13    chains: Mutex<HashMap<String, Arc<Mutex<()>>>>,
14}
15
16impl Default for InMemoryCredentialStore {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl InMemoryCredentialStore {
23    pub fn new() -> Self {
24        Self {
25            credentials: Mutex::new(HashMap::new()),
26            chains: Mutex::new(HashMap::new()),
27        }
28    }
29
30    async fn lock_chain(&self, provider_id: &str) -> Arc<Mutex<()>> {
31        let mut chains = self.chains.lock().await;
32        chains
33            .entry(provider_id.to_string())
34            .or_insert_with(|| Arc::new(Mutex::new(())))
35            .clone()
36    }
37}
38
39#[async_trait::async_trait]
40impl CredentialStore for InMemoryCredentialStore {
41    async fn read(&self, provider_id: &str) -> Option<Credential> {
42        self.credentials.lock().await.get(provider_id).cloned()
43    }
44
45    async fn modify(
46        &self,
47        provider_id: &str,
48        f: Box<dyn FnOnce(Option<Credential>) -> Pin<Box<dyn Future<Output = Option<Credential>> + Send>> + Send>,
49    ) -> Option<Credential> {
50        let chain = self.lock_chain(provider_id).await;
51        let _guard = chain.lock().await;
52        let current = self.credentials.lock().await.get(provider_id).cloned();
53        let next = f(current).await;
54        if let Some(ref cred) = next {
55            self.credentials
56                .lock()
57                .await
58                .insert(provider_id.to_string(), cred.clone());
59        }
60        if next.is_some() {
61            next
62        } else {
63            self.credentials.lock().await.get(provider_id).cloned()
64        }
65    }
66
67    async fn delete(&self, provider_id: &str) {
68        let chain = self.lock_chain(provider_id).await;
69        let _guard = chain.lock().await;
70        self.credentials.lock().await.remove(provider_id);
71    }
72}