eric / rustwarden

//! In-process [`GuardStore`] backend.
//!
//! Implemented with `DashMap` for O(1) hot-path reads. Events, bans,
//! and ACLs live only in RAM — restarting the process clears
//! everything. Use this in tests, or wrapped by a write-through
//! `HybridStore` in production.

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use dashmap::DashMap;
use rustwarden_core::error::AppError;
use uuid::Uuid;

use super::{AclEntry, BanEntry, GuardStore};
use crate::events::{GuardEvent, GuardEventKind};

/// Concurrent in-memory store.
#[derive(Default)]
pub struct MemoryStore {
    events: DashMap<String, Vec<GuardEvent>>,
    bans: DashMap<Uuid, BanEntry>,
    acl: DashMap<String, AclEntry>,
}

impl MemoryStore {
    /// Build an empty store.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl GuardStore for MemoryStore {
    async fn insert_event(&self, event: &GuardEvent) -> Result<(), AppError> {
        self.events
            .entry(event.domain.clone())
            .or_default()
            .push(event.clone());
        Ok(())
    }

    async fn count_events_since(
        &self,
        domain: &str,
        kind: GuardEventKind,
        since: DateTime<Utc>,
    ) -> Result<u64, AppError> {
        let count = self
            .events
            .get(domain)
            .map(|bucket| {
                bucket
                    .iter()
                    .filter(|e| e.kind == kind && e.created_at >= since)
                    .count() as u64
            })
            .unwrap_or(0);
        Ok(count)
    }

    async fn insert_ban(&self, ban: &BanEntry) -> Result<(), AppError> {
        self.bans.insert(ban.id, ban.clone());
        Ok(())
    }

    async fn active_bans_for_domain(
        &self,
        domain: &str,
        now: DateTime<Utc>,
    ) -> Result<Vec<BanEntry>, AppError> {
        let out = self
            .bans
            .iter()
            .filter(|kv| kv.value().domain == domain && kv.value().is_active_at(now))
            .map(|kv| kv.value().clone())
            .collect();
        Ok(out)
    }

    async fn all_active_bans(&self, now: DateTime<Utc>) -> Result<Vec<BanEntry>, AppError> {
        Ok(self
            .bans
            .iter()
            .filter(|kv| kv.value().is_active_at(now))
            .map(|kv| kv.value().clone())
            .collect())
    }

    async fn lift_ban(
        &self,
        ban_id: Uuid,
        by: Option<Uuid>,
        at: DateTime<Utc>,
    ) -> Result<(), AppError> {
        if let Some(mut row) = self.bans.get_mut(&ban_id) {
            row.lifted_at = Some(at);
            row.lifted_by = by;
        }
        Ok(())
    }

    async fn upsert_acl(&self, entry: &AclEntry) -> Result<(), AppError> {
        self.acl.insert(entry.domain.clone(), entry.clone());
        Ok(())
    }

    async fn acl_for_domain(&self, domain: &str) -> Result<Option<AclEntry>, AppError> {
        Ok(self.acl.get(domain).map(|kv| kv.value().clone()))
    }

    async fn delete_acl(&self, id: Uuid) -> Result<(), AppError> {
        self.acl.retain(|_, v| v.id != id);
        Ok(())
    }

    async fn list_acl(&self) -> Result<Vec<AclEntry>, AppError> {
        Ok(self.acl.iter().map(|kv| kv.value().clone()).collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::events::{GuardEvent, GuardEventKind, Severity};
    use crate::store::AclKind;

    fn evt(domain: &str, kind: GuardEventKind) -> GuardEvent {
        GuardEvent::new(kind, domain, Severity::Medium)
    }

    #[tokio::test]
    async fn insert_and_count_events_by_kind_and_window() {
        let store = MemoryStore::new();
        let e1 = evt("a.tld", GuardEventKind::SignatureFailure);
        let e2 = evt("a.tld", GuardEventKind::SignatureFailure);
        let e3 = evt("a.tld", GuardEventKind::ReplayAttempt);
        store.insert_event(&e1).await.unwrap();
        store.insert_event(&e2).await.unwrap();
        store.insert_event(&e3).await.unwrap();
        let cutoff = Utc::now() - chrono::Duration::minutes(1);
        let n = store
            .count_events_since("a.tld", GuardEventKind::SignatureFailure, cutoff)
            .await
            .unwrap();
        assert_eq!(n, 2);
    }

    #[tokio::test]
    async fn ban_lifecycle_insert_active_lift() {
        let store = MemoryStore::new();
        let now = Utc::now();
        let ban = BanEntry {
            id: Uuid::new_v4(),
            domain: "evil.tld".into(),
            actor_uri: None,
            reason: "auto-ban".into(),
            severity: "suspend".into(),
            triggered_by: "auto:test".into(),
            banned_at: now,
            expires_at: Some(now + chrono::Duration::hours(1)),
            lifted_at: None,
            lifted_by: None,
        };
        let ban_id = ban.id;
        store.insert_ban(&ban).await.unwrap();
        let active = store
            .active_bans_for_domain("evil.tld", now + chrono::Duration::minutes(10))
            .await
            .unwrap();
        assert_eq!(active.len(), 1);
        let admin = Uuid::new_v4();
        store
            .lift_ban(ban_id, Some(admin), Utc::now())
            .await
            .unwrap();
        let active_after = store
            .active_bans_for_domain("evil.tld", now + chrono::Duration::minutes(10))
            .await
            .unwrap();
        assert!(active_after.is_empty());
    }

    #[tokio::test]
    async fn acl_upsert_lookup_delete() {
        let store = MemoryStore::new();
        let entry = AclEntry {
            id: Uuid::new_v4(),
            kind: AclKind::Allow,
            domain: "trusted.tld".into(),
            reason: Some("peer".into()),
            created_by: Uuid::new_v4(),
            created_at: Utc::now(),
            updated_at: Utc::now(),
        };
        store.upsert_acl(&entry).await.unwrap();
        let found = store.acl_for_domain("trusted.tld").await.unwrap();
        assert!(found.is_some());
        store.delete_acl(entry.id).await.unwrap();
        let after = store.acl_for_domain("trusted.tld").await.unwrap();
        assert!(after.is_none());
    }
}