Skip to main content

dbx_tools_auth/storage/
memory.rs

1use std::{
2    collections::HashMap,
3    sync::{Arc, Mutex as StdMutex},
4    time::Duration,
5};
6
7use async_trait::async_trait;
8use tokio::sync::{Mutex, OwnedMutexGuard, RwLock};
9
10use super::{CredentialStore, StorageLock};
11use crate::{Error, Result, Token};
12
13pub struct MemoryStore {
14    tokens: RwLock<HashMap<String, Token>>,
15    locks: StdMutex<HashMap<String, Arc<Mutex<()>>>>,
16}
17
18impl MemoryStore {
19    pub fn new() -> Self {
20        Self {
21            tokens: RwLock::new(HashMap::new()),
22            locks: StdMutex::new(HashMap::new()),
23        }
24    }
25}
26
27impl Default for MemoryStore {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33struct MemoryLock {
34    _guard: OwnedMutexGuard<()>,
35}
36
37impl StorageLock for MemoryLock {}
38
39#[async_trait]
40impl CredentialStore for MemoryStore {
41    async fn load(&self, profile: &str) -> Result<Option<Token>> {
42        Ok(self.tokens.read().await.get(profile).cloned())
43    }
44
45    async fn save(&self, profile: &str, token: &Token) -> Result<()> {
46        self.tokens
47            .write()
48            .await
49            .insert(profile.to_owned(), token.clone());
50        Ok(())
51    }
52
53    async fn delete(&self, profile: &str) -> Result<()> {
54        self.tokens.write().await.remove(profile);
55        Ok(())
56    }
57
58    async fn lock(&self, profile: &str, timeout: Duration) -> Result<Box<dyn StorageLock>> {
59        let gate = {
60            let mut locks = self
61                .locks
62                .lock()
63                .map_err(|_| Error::Storage("memory lock registry is poisoned".into()))?;
64            Arc::clone(
65                locks
66                    .entry(profile.to_owned())
67                    .or_insert_with(|| Arc::new(Mutex::new(()))),
68            )
69        };
70        let guard = tokio::time::timeout(timeout, gate.lock_owned())
71            .await
72            .map_err(|_| Error::LockTimeout(profile.to_owned()))?;
73        Ok(Box::new(MemoryLock { _guard: guard }))
74    }
75
76    fn name(&self) -> &'static str {
77        "memory"
78    }
79}