94 lines
2.7 KiB
Rust
94 lines
2.7 KiB
Rust
use crate::models::ChatSession;
|
|
use dashmap::DashMap;
|
|
use std::{
|
|
sync::Mutex,
|
|
time::{Duration, Instant},
|
|
};
|
|
use tokio::sync::broadcast;
|
|
|
|
pub const RESET_EVENT: &str = "__reset__";
|
|
// Changed from 3 minutes to 20 minutes
|
|
pub const CHAT_RESET_INTERVAL: Duration = Duration::from_secs(20 * 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>,
|
|
pub notifiers: DashMap<String, broadcast::Sender<String>>,
|
|
pub supporter_notifier: broadcast::Sender<String>,
|
|
pub rate_limits: DashMap<String, RateLimitEntry>,
|
|
reset_started: Mutex<Instant>,
|
|
}
|
|
|
|
impl AppState {
|
|
pub fn new() -> Self {
|
|
let (supporter_tx, _) = broadcast::channel(64);
|
|
Self {
|
|
sessions: DashMap::new(),
|
|
notifiers: DashMap::new(),
|
|
supporter_notifier: supporter_tx,
|
|
rate_limits: DashMap::new(),
|
|
reset_started: Mutex::new(Instant::now()),
|
|
}
|
|
}
|
|
|
|
pub fn get_or_create_notifier(&self, session_id: &str) -> broadcast::Sender<String> {
|
|
if let Some(tx) = self.notifiers.get(session_id) {
|
|
tx.clone()
|
|
} else {
|
|
let (tx, _) = broadcast::channel(64);
|
|
self.notifiers.insert(session_id.to_string(), tx.clone());
|
|
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());
|
|
}
|
|
}
|