security

src/state.rs

@@ -1,22 +1,40 @@
-use dashmap::DashMap;
-use tokio::sync::broadcast;
 use crate::models::ChatSession;
+use dashmap::DashMap;
+use std::{sync::Mutex, time::{Duration, Instant}};
+use tokio::sync::broadcast;
+
+pub const RESET_EVENT: &str = "__reset__";
+pub const CHAT_RESET_INTERVAL: Duration = Duration::from_secs(3 * 60);
+pub const RATE_LIMIT_WINDOW: Duration = CHAT_RESET_INTERVAL;
+pub const MESSAGE_LIMIT_PER_WINDOW: u32 = 60;
+pub const SESSION_LIMIT_PER_WINDOW: u32 = 10;
+
+#[derive(Debug, Clone, Copy)]
+pub struct RateLimitEntry {
+    window_started: Instant,
+    count: u32,
+}
 
 pub struct AppState {
     pub sessions: DashMap<String, ChatSession>,
     // broadcast channel: session_id -> sender
     pub notifiers: DashMap<String, broadcast::Sender<String>>,
-    // agent broadcast for new sessions
-    pub agent_notifier: broadcast::Sender<String>,
+    // supporter broadcast for new sessions and updates
+    pub supporter_notifier: broadcast::Sender<String>,
+    // in-memory rate limits, wiped with all other state every three minutes
+    pub rate_limits: DashMap<String, RateLimitEntry>,
+    reset_started: Mutex<Instant>,
 }
 
 impl AppState {
     pub fn new() -> Self {
-        let (agent_tx, _) = broadcast::channel(64);
+        let (supporter_tx, _) = broadcast::channel(64);
         Self {
             sessions: DashMap::new(),
             notifiers: DashMap::new(),
-            agent_notifier: agent_tx,
+            supporter_notifier: supporter_tx,
+            rate_limits: DashMap::new(),
+            reset_started: Mutex::new(Instant::now()),
         }
     }
 
@@ -29,4 +47,46 @@ impl AppState {
             tx
         }
     }
+
+    pub fn is_rate_limited(&self, key: String, limit: u32) -> bool {
+        let now = Instant::now();
+        let mut entry = self.rate_limits.entry(key).or_insert(RateLimitEntry {
+            window_started: now,
+            count: 0,
+        });
+
+        if now.duration_since(entry.window_started) >= RATE_LIMIT_WINDOW {
+            entry.window_started = now;
+            entry.count = 0;
+        }
+
+        entry.count += 1;
+        entry.count > limit
+    }
+
+    pub fn reset_seconds_remaining(&self) -> u64 {
+        let elapsed = self
+            .reset_started
+            .lock()
+            .map(|started| started.elapsed())
+            .unwrap_or(Duration::ZERO);
+
+        CHAT_RESET_INTERVAL.saturating_sub(elapsed).as_secs()
+    }
+
+    pub fn clear_everything(&self) {
+        for tx in self.notifiers.iter() {
+            let _ = tx.value().send(RESET_EVENT.to_string());
+        }
+
+        self.sessions.clear();
+        self.notifiers.clear();
+        self.rate_limits.clear();
+
+        if let Ok(mut reset_started) = self.reset_started.lock() {
+            *reset_started = Instant::now();
+        }
+
+        let _ = self.supporter_notifier.send(RESET_EVENT.to_string());
+    }
 }