elph_ai/auth/
credential_store.rs1use 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
10pub 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}