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, } #[derive(Clone)] pub struct QueryCache { pub map: Arc>, pub hits: Arc, pub misses: Arc, } #[derive(Clone, Debug)] pub struct CacheEntry { pub sql: String, pub last_accessed_secs: Arc, } 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 { 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>>, } #[derive(Clone)] pub struct CompiledEntry { pub entry: BlacklistEntry, pub pattern: Pattern, pub bypass_mask: Option, } impl BlacklistCache { pub fn new() -> Self { Self { inner: Arc::new(RwLock::new(Vec::new())), } } pub async fn load(&self, entries: Vec) { let compiled: Vec = entries .into_iter() .filter_map(|e| { let bypass_mask = e .bypass_mask .as_deref() .and_then(|s| s.parse::().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); } }