dbx-tools-auth 0.6.167

Provider-neutral OAuth, credential storage, and locking
Documentation
use std::{
    collections::HashMap,
    sync::{Arc, Mutex as StdMutex},
    time::Duration,
};

use async_trait::async_trait;
use tokio::sync::{Mutex, OwnedMutexGuard, RwLock};

use super::{CredentialStore, StorageLock};
use crate::{Error, Result, Token};

pub struct MemoryStore {
    tokens: RwLock<HashMap<String, Token>>,
    locks: StdMutex<HashMap<String, Arc<Mutex<()>>>>,
}

impl MemoryStore {
    pub fn new() -> Self {
        Self {
            tokens: RwLock::new(HashMap::new()),
            locks: StdMutex::new(HashMap::new()),
        }
    }
}

impl Default for MemoryStore {
    fn default() -> Self {
        Self::new()
    }
}

struct MemoryLock {
    _guard: OwnedMutexGuard<()>,
}

impl StorageLock for MemoryLock {}

#[async_trait]
impl CredentialStore for MemoryStore {
    async fn load(&self, profile: &str) -> Result<Option<Token>> {
        Ok(self.tokens.read().await.get(profile).cloned())
    }

    async fn save(&self, profile: &str, token: &Token) -> Result<()> {
        self.tokens
            .write()
            .await
            .insert(profile.to_owned(), token.clone());
        Ok(())
    }

    async fn delete(&self, profile: &str) -> Result<()> {
        self.tokens.write().await.remove(profile);
        Ok(())
    }

    async fn lock(&self, profile: &str, timeout: Duration) -> Result<Box<dyn StorageLock>> {
        let gate = {
            let mut locks = self
                .locks
                .lock()
                .map_err(|_| Error::Storage("memory lock registry is poisoned".into()))?;
            Arc::clone(
                locks
                    .entry(profile.to_owned())
                    .or_insert_with(|| Arc::new(Mutex::new(()))),
            )
        };
        let guard = tokio::time::timeout(timeout, gate.lock_owned())
            .await
            .map_err(|_| Error::LockTimeout(profile.to_owned()))?;
        Ok(Box::new(MemoryLock { _guard: guard }))
    }

    fn name(&self) -> &'static str {
        "memory"
    }
}