Mercury/src/state.rs

use std::sync::{
    atomic::{AtomicU64, Ordering},
    Arc,
};
use std::time::{SystemTime, UNIX_EPOCH};

use dashmap::DashMap;
use glob::Pattern;
use sqlx::PgPool;
use tokio::sync::RwLock;

use crate::config::Config;
use crate::models::blacklist::BlacklistEntry;

pub fn unix_now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

#[derive(Clone)]
pub struct AppState {
    pub pool: PgPool,
    pub query_cache: QueryCache,
    pub blacklist_cache: BlacklistCache,
    pub config: Arc<Config>,
}

#[derive(Clone)]
pub struct QueryCache {
    pub map: Arc<DashMap<String, CacheEntry>>,
    pub hits: Arc<AtomicU64>,
    pub misses: Arc<AtomicU64>,
}

#[derive(Clone, Debug)]
pub struct CacheEntry {
    pub sql: String,
    pub last_accessed_secs: Arc<AtomicU64>,
}

impl CacheEntry {
    pub fn new(sql: String) -> Self {
        Self {
            sql,
            last_accessed_secs: Arc::new(AtomicU64::new(unix_now())),
        }
    }

    pub fn touch(&self) {
        self.last_accessed_secs.store(unix_now(), Ordering::Relaxed);
    }

    pub fn last_accessed(&self) -> u64 {
        self.last_accessed_secs.load(Ordering::Relaxed)
    }
}

impl QueryCache {
    pub fn new() -> Self {
        Self {
            map: Arc::new(DashMap::new()),
            hits: Arc::new(AtomicU64::new(0)),
            misses: Arc::new(AtomicU64::new(0)),
        }
    }

    pub fn get(&self, key: &str) -> Option<CacheEntry> {
        if let Some(entry) = self.map.get(key) {
            entry.touch();
            self.hits.fetch_add(1, Ordering::Relaxed);
            Some(entry.clone())
        } else {
            self.misses.fetch_add(1, Ordering::Relaxed);
            None
        }
    }

    pub fn insert(&self, key: String, entry: CacheEntry, max_capacity: usize) {
        if self.map.len() >= max_capacity {
            let oldest_key = self
                .map
                .iter()
                .min_by_key(|e| e.last_accessed())
                .map(|e| e.key().clone());
            if let Some(k) = oldest_key {
                self.map.remove(&k);
            }
        }
        self.map.insert(key, entry);
    }

    #[allow(dead_code)]
    pub fn remove(&self, key: &str) {
        self.map.remove(key);
    }

    pub fn hits(&self) -> u64 {
        self.hits.load(Ordering::Relaxed)
    }

    pub fn misses(&self) -> u64 {
        self.misses.load(Ordering::Relaxed)
    }

    pub fn len(&self) -> usize {
        self.map.len()
    }

    pub fn flush(&self) {
        self.map.clear();
    }
}

#[derive(Clone)]
pub struct BlacklistCache {
    pub inner: Arc<RwLock<Vec<CompiledEntry>>>,
}

#[derive(Clone)]
pub struct CompiledEntry {
    pub entry: BlacklistEntry,
    pub pattern: Pattern,
    pub bypass_mask: Option<u128>,
}

impl BlacklistCache {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(RwLock::new(Vec::new())),
        }
    }

    pub async fn load(&self, entries: Vec<BlacklistEntry>) {
        let compiled: Vec<CompiledEntry> = entries
            .into_iter()
            .filter_map(|e| {
                let bypass_mask = e
                    .bypass_mask
                    .as_deref()
                    .and_then(|s| s.parse::<u128>().ok());
                Pattern::new(&e.pattern).ok().map(|pattern| CompiledEntry {
                    entry: e,
                    pattern,
                    bypass_mask,
                })
            })
            .collect();
        let mut guard = self.inner.write().await;
        *guard = compiled;
    }

    /// Returns true if the request should be blocked.
    /// `caller_mask` is 0 for unauthenticated requests; bypass only applies
    /// when the caller holds the permission bit stored in bypass_mask.
    pub async fn is_blocked(&self, method: &str, path: &str, caller_mask: u128) -> bool {
        let path = path.trim_end_matches('/');
        let path = if path.is_empty() { "/" } else { path };
        let guard = self.inner.read().await;
        guard.iter().any(|compiled| {
            if !compiled.entry.active {
                return false;
            }
            let method_matches = compiled
                .entry
                .method
                .as_deref()
                .map(|m| {
                    m.split(',')
                        .any(|part| part.trim().eq_ignore_ascii_case(method))
                })
                .unwrap_or(true);
            if !method_matches || !compiled.pattern.matches(path) {
                return false;
            }
            // If caller holds the bypass permission, they are not blocked.
            !matches!(compiled.bypass_mask, Some(mask) if caller_mask & mask != 0)
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cache_insert_and_get() {
        let cache = QueryCache::new();
        let entry = CacheEntry::new("SELECT 1".into());
        cache.insert("key1".into(), entry, 100);
        let got = cache.get("key1");
        assert!(got.is_some());
        assert_eq!(got.unwrap().sql, "SELECT 1");
        assert_eq!(cache.hits(), 1);
        assert_eq!(cache.misses(), 0);
    }

    #[test]
    fn test_cache_miss() {
        let cache = QueryCache::new();
        let got = cache.get("missing");
        assert!(got.is_none());
        assert_eq!(cache.misses(), 1);
    }

    #[test]
    fn test_cache_capacity_evicts_oldest() {
        let cache = QueryCache::new();
        let e1 = CacheEntry::new("SELECT 1".into());
        // force e1 to be older
        e1.last_accessed_secs.store(1, Ordering::Relaxed);
        cache.map.insert("old".into(), e1);
        let e2 = CacheEntry::new("SELECT 2".into());
        cache.insert("new".into(), e2, 1); // capacity=1, should evict "old"
        assert!(cache.map.get("old").is_none());
        assert!(cache.map.get("new").is_some());
    }

    #[tokio::test]
    async fn test_blacklist_blocks_pattern() {
        use chrono::Utc;
        let cache = BlacklistCache::new();
        let entry = BlacklistEntry {
            id: 1,
            pattern: "/api/users/**".into(),
            method: None,
            reason: None,
            active: true,
            bypass_mask: None,
            created_at: Utc::now(),
        };
        cache.load(vec![entry]).await;
        assert!(cache.is_blocked("GET", "/api/users/42", 0).await);
        assert!(!cache.is_blocked("GET", "/api/orders/1", 0).await);
    }

    #[tokio::test]
    async fn test_blacklist_method_specific() {
        use chrono::Utc;
        let cache = BlacklistCache::new();
        let entry = BlacklistEntry {
            id: 2,
            pattern: "/api/secrets".into(),
            method: Some("GET".into()),
            reason: None,
            active: true,
            bypass_mask: None,
            created_at: Utc::now(),
        };
        cache.load(vec![entry]).await;
        assert!(cache.is_blocked("GET", "/api/secrets", 0).await);
        assert!(!cache.is_blocked("POST", "/api/secrets", 0).await);
    }
}