diff --git a/src/handlers/chat.rs b/src/handlers/chat.rs index 28b9e20..b91568a 100644 --- a/src/handlers/chat.rs +++ b/src/handlers/chat.rs @@ -1,28 +1,28 @@ use axum::{ - extract::{Path, State, Form}, + extract::{ConnectInfo, Form, Path, State}, response::Html, }; use chrono::Utc; -use std::sync::Arc; +use std::{net::SocketAddr, sync::Arc}; use uuid::Uuid; use crate::{ - models::{ChatSession, Message, MessageSender, SendMessageForm, NewSessionForm, SessionStatus}, - state::AppState, + models::{ChatSession, Message, MessageSender, NewSessionForm, SendMessageForm, SessionStatus}, + state::{AppState, MESSAGE_LIMIT_PER_WINDOW, SESSION_LIMIT_PER_WINDOW}, }; +const MAX_MESSAGE_CHARS: usize = 1_000; +const MAX_NAME_CHARS: usize = 40; + fn render_message(msg: &Message) -> String { let (bubble_class, label) = match msg.sender { - MessageSender::Customer => ("msg-customer", "You"), - MessageSender::Agent => ("msg-agent", "Support"), + MessageSender::User => ("msg-user", "You"), + MessageSender::Supporter => ("msg-supporter", "Support"), MessageSender::System => ("msg-system", ""), }; if msg.sender == MessageSender::System { - format!( - "
{}
", - html_escape(&msg.content) - ) + render_system_notice(&msg.content) } else { let time = msg.timestamp.format("%H:%M").to_string(); format!( @@ -39,11 +39,34 @@ fn render_message(msg: &Message) -> String { } } +fn render_system_notice(content: &str) -> String { + format!( + "
{}
", + html_escape(content) + ) +} + fn html_escape(s: &str) -> String { s.replace('&', "&") .replace('<', "<") .replace('>', ">") .replace('"', """) + .replace('\'', "'") +} + +fn limit_key(kind: &str, addr: SocketAddr) -> String { + format!("{}:{}", kind, addr.ip()) +} + +fn clean_text(input: &str, max_chars: usize) -> Result { + let trimmed = input.trim(); + if trimmed.is_empty() { + return Ok(String::new()); + } + if trimmed.chars().count() > max_chars { + return Err(format!("Please keep this under {max_chars} characters.")); + } + Ok(trimmed.to_string()) } /// GET /messages/:session_id @@ -52,76 +75,119 @@ pub async fn get_messages( State(state): State>, ) -> Html { let Some(session) = state.sessions.get(&session_id) else { - return Html("

Session not found.

".to_string()); + return Html(render_system_notice( + "Session not found. Chat history may have been cleared; please start a new chat.", + )); }; let html: String = session.messages.iter().map(render_message).collect(); Html(html) } -/// POST /send/:session_id — customer sends a message +/// POST /send/:session_id — user sends a message pub async fn send_message( + ConnectInfo(addr): ConnectInfo, Path(session_id): Path, State(state): State>, Form(form): Form, ) -> Html { - let content = form.content.trim().to_string(); + if state.is_rate_limited(limit_key("message", addr), MESSAGE_LIMIT_PER_WINDOW) { + return Html(render_system_notice( + "Too many messages from this network. Please wait until the 3-minute window resets.", + )); + } + + let content = match clean_text(&form.content, MAX_MESSAGE_CHARS) { + Ok(content) => content, + Err(message) => return Html(render_system_notice(&message)), + }; if content.is_empty() { return Html(String::new()); } + + let Some(mut session) = state.sessions.get_mut(&session_id) else { + return Html(render_system_notice( + "Session not found. Chat history may have been cleared; please start a new chat.", + )); + }; + let msg = Message { id: Uuid::new_v4().to_string(), session_id: session_id.clone(), - sender: MessageSender::Customer, + sender: MessageSender::User, content, timestamp: Utc::now(), }; let rendered = render_message(&msg); - if let Some(mut session) = state.sessions.get_mut(&session_id) { - session.messages.push(msg); - if let Some(tx) = state.notifiers.get(&session_id) { - let _ = tx.send(session_id.clone()); - } - let _ = state.agent_notifier.send(session_id.clone()); + session.messages.push(msg); + + if let Some(tx) = state.notifiers.get(&session_id) { + let _ = tx.send(session_id.clone()); } + let _ = state.supporter_notifier.send(session_id.clone()); + Html(rendered) } -/// POST /agent/send/:session_id — agent sends a message -pub async fn agent_send_message( +/// POST /supporter/send/:session_id — supporter sends a message +pub async fn supporter_send_message( + ConnectInfo(addr): ConnectInfo, Path(session_id): Path, State(state): State>, Form(form): Form, ) -> Html { - let content = form.content.trim().to_string(); + if state.is_rate_limited(limit_key("supporter-message", addr), MESSAGE_LIMIT_PER_WINDOW) { + return Html(render_system_notice( + "Too many messages from this network. Please wait until the 3-minute window resets.", + )); + } + + let content = match clean_text(&form.content, MAX_MESSAGE_CHARS) { + Ok(content) => content, + Err(message) => return Html(render_system_notice(&message)), + }; if content.is_empty() { return Html(String::new()); } + + let Some(mut session) = state.sessions.get_mut(&session_id) else { + return Html(render_system_notice( + "Session not found. Chat history may have been cleared.", + )); + }; + let msg = Message { id: Uuid::new_v4().to_string(), session_id: session_id.clone(), - sender: MessageSender::Agent, + sender: MessageSender::Supporter, content, timestamp: Utc::now(), }; let rendered = render_message(&msg); - if let Some(mut session) = state.sessions.get_mut(&session_id) { - session.messages.push(msg); - session.status = SessionStatus::Active; - if let Some(tx) = state.notifiers.get(&session_id) { - let _ = tx.send(session_id.clone()); - } + session.messages.push(msg); + session.status = SessionStatus::Active; + + if let Some(tx) = state.notifiers.get(&session_id) { + let _ = tx.send(session_id.clone()); } + let _ = state.supporter_notifier.send(session_id.clone()); + Html(rendered) } -/// GET /sessions — session list for agent dashboard -pub async fn get_sessions( - State(state): State>, -) -> Html { - let mut sessions: Vec<_> = state.sessions +/// GET /sessions — session list for supporter dashboard +pub async fn get_sessions(State(state): State>) -> Html { + let mut sessions: Vec<_> = state + .sessions .iter() .filter(|s| s.status != SessionStatus::Closed) - .map(|s| (s.id.clone(), s.customer_name.clone(), s.status.clone(), s.messages.len())) + .map(|s| { + ( + s.id.clone(), + s.user_name.clone(), + s.status.clone(), + s.messages.len(), + ) + }) .collect(); sessions.sort_by(|a, b| a.0.cmp(&b.0)); @@ -134,22 +200,24 @@ pub async fn get_sessions( for (id, name, status, count) in &sessions { let status_class = match status { SessionStatus::Waiting => "status-waiting", - SessionStatus::Active => "status-active", - SessionStatus::Closed => "status-closed", + SessionStatus::Active => "status-active", + SessionStatus::Closed => "status-closed", }; let status_label = match status { SessionStatus::Waiting => "Waiting", - SessionStatus::Active => "Active", - SessionStatus::Closed => "Closed", + SessionStatus::Active => "Active", + SessionStatus::Closed => "Closed", }; let short_id = &id[..8]; let escaped_name = html_escape(name); html.push_str(&format!( "\ \ \ diff --git a/src/handlers/pages.rs b/src/handlers/pages.rs index ddcb8d3..75e62dc 100644 --- a/src/handlers/pages.rs +++ b/src/handlers/pages.rs @@ -1,9 +1,21 @@ -use axum::response::Html; +use axum::{extract::State, response::Html}; +use std::sync::Arc; -pub async fn customer_page() -> Html<&'static str> { - Html(include_str!("../../templates/customer.html")) +use crate::state::AppState; + +pub async fn index_page() -> Html<&'static str> { + Html(include_str!("../../templates/index.html")) } -pub async fn agent_page() -> Html<&'static str> { - Html(include_str!("../../templates/agent.html")) +pub async fn user_page() -> Html<&'static str> { + Html(include_str!("../../templates/user.html")) +} + +pub async fn supporter_page() -> Html<&'static str> { + Html(include_str!("../../templates/supporter.html")) +} + +pub async fn reset_countdown(State(state): State>) -> Html { + let seconds = state.reset_seconds_remaining(); + Html(format!("Reset in {:02}:{:02}", seconds / 60, seconds % 60)) } diff --git a/src/handlers/ws.rs b/src/handlers/ws.rs index a028ac7..2d330fd 100644 --- a/src/handlers/ws.rs +++ b/src/handlers/ws.rs @@ -7,35 +7,44 @@ use axum::{ }; use std::sync::Arc; -use crate::state::AppState; +use crate::state::{AppState, RESET_EVENT}; -/// WebSocket for customer — pushes notifications when agent replies -pub async fn customer_ws( +const USER_TYPING_PREFIX: &str = "typing:user:"; +const SUPPORTER_TYPING_PREFIX: &str = "typing:supporter:"; + +/// WebSocket for user — pushes notifications when supporter replies +pub async fn user_ws( ws: WebSocketUpgrade, Path(session_id): Path, State(state): State>, ) -> impl IntoResponse { - ws.on_upgrade(move |socket| handle_customer_socket(socket, session_id, state)) + ws.on_upgrade(move |socket| handle_user_socket(socket, session_id, state)) } -async fn handle_customer_socket(mut socket: WebSocket, session_id: String, state: Arc) { +async fn handle_user_socket(mut socket: WebSocket, session_id: String, state: Arc) { let tx = state.get_or_create_notifier(&session_id); let mut rx = tx.subscribe(); loop { tokio::select! { - // Forward broadcast events to the customer browser + // Forward broadcast events to the user browser. Ok(sid) = rx.recv() => { - if sid == session_id { - // Signal HTMX to re-poll messages - if socket.send(Message::Text("refresh".into())).await.is_err() { + if sid == session_id || sid == RESET_EVENT || sid.starts_with(SUPPORTER_TYPING_PREFIX) || sid.starts_with(USER_TYPING_PREFIX) { + if socket.send(Message::Text(sid)).await.is_err() { break; } } } - // Handle ping/close from client + // Handle ping/close and typing events from client. msg = socket.recv() => { match msg { + Some(Ok(Message::Text(text))) => { + let expected = format!("{USER_TYPING_PREFIX}{session_id}"); + if text == expected { + let _ = tx.send(text.clone()); + let _ = state.supporter_notifier.send(text); + } + } Some(Ok(Message::Close(_))) | None => break, _ => {} } @@ -44,27 +53,35 @@ async fn handle_customer_socket(mut socket: WebSocket, session_id: String, state } } -/// WebSocket for agent dashboard — pushes notifications when any session updates -pub async fn agent_ws( +/// WebSocket for supporter dashboard — pushes notifications when any session updates +pub async fn supporter_ws( ws: WebSocketUpgrade, - Path(_agent_id): Path, + Path(_supporter_id): Path, State(state): State>, ) -> impl IntoResponse { - ws.on_upgrade(move |socket| handle_agent_socket(socket, state)) + ws.on_upgrade(move |socket| handle_supporter_socket(socket, state)) } -async fn handle_agent_socket(mut socket: WebSocket, state: Arc) { - let mut rx = state.agent_notifier.subscribe(); +async fn handle_supporter_socket(mut socket: WebSocket, state: Arc) { + let mut rx = state.supporter_notifier.subscribe(); loop { tokio::select! { Ok(session_id) = rx.recv() => { - if socket.send(Message::Text(session_id.into())).await.is_err() { + if socket.send(Message::Text(session_id)).await.is_err() { break; } } msg = socket.recv() => { match msg { + Some(Ok(Message::Text(text))) => { + if let Some(session_id) = text.strip_prefix(SUPPORTER_TYPING_PREFIX) { + if !session_id.is_empty() && state.sessions.contains_key(session_id) { + let tx = state.get_or_create_notifier(session_id); + let _ = tx.send(text); + } + } + } Some(Ok(Message::Close(_))) | None => break, _ => {} } diff --git a/src/main.rs b/src/main.rs index b06d033..83afdf6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,11 +6,11 @@ use axum::{ Router, routing::{get, post}, }; -use std::sync::Arc; +use std::{net::SocketAddr, sync::Arc}; use tower_http::services::ServeDir; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; -use state::AppState; +use state::{AppState, CHAT_RESET_INTERVAL}; #[tokio::main] async fn main() { @@ -23,16 +23,27 @@ async fn main() { let state = Arc::new(AppState::new()); + let reset_state = state.clone(); + tokio::spawn(async move { + loop { + tokio::time::sleep(CHAT_RESET_INTERVAL).await; + reset_state.clear_everything(); + tracing::info!("Cleared all in-memory chat sessions, messages, notifiers, and rate limits"); + } + }); + let app = Router::new() - .route("/", get(handlers::pages::customer_page)) - .route("/agent", get(handlers::pages::agent_page)) - .route("/ws/customer/:session_id", get(handlers::ws::customer_ws)) - .route("/ws/agent/:agent_id", get(handlers::ws::agent_ws)) + .route("/", get(handlers::pages::index_page)) + .route("/user", get(handlers::pages::user_page)) + .route("/supporter", get(handlers::pages::supporter_page)) + .route("/reset-countdown", get(handlers::pages::reset_countdown)) + .route("/ws/user/:session_id", get(handlers::ws::user_ws)) + .route("/ws/supporter/:supporter_id", get(handlers::ws::supporter_ws)) .route("/messages/:session_id", get(handlers::chat::get_messages)) .route("/send/:session_id", post(handlers::chat::send_message)) .route( - "/agent/send/:session_id", - post(handlers::chat::agent_send_message), + "/supporter/send/:session_id", + post(handlers::chat::supporter_send_message), ) .route("/sessions", get(handlers::chat::get_sessions)) .route("/sessions/new", post(handlers::chat::new_session)) @@ -43,7 +54,7 @@ async fn main() { tracing::info!("Chat service running — http://0.0.0.0:3111"); axum::Server::from_tcp(listener.into_std().unwrap()) .unwrap() - .serve(app.into_make_service()) + .serve(app.into_make_service_with_connect_info::()) .await .unwrap(); } diff --git a/src/models.rs b/src/models.rs index b2e51c2..5f23746 100644 --- a/src/models.rs +++ b/src/models.rs @@ -4,8 +4,8 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "lowercase")] pub enum MessageSender { - Customer, - Agent, + User, + Supporter, System, } @@ -29,25 +29,25 @@ pub enum SessionStatus { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChatSession { pub id: String, - pub customer_name: String, + pub user_name: String, pub status: SessionStatus, pub created_at: DateTime, pub messages: Vec, } impl ChatSession { - pub fn new(id: String, customer_name: String) -> Self { + pub fn new(id: String, user_name: String) -> Self { let now = Utc::now(); let welcome = Message { id: uuid::Uuid::new_v4().to_string(), session_id: id.clone(), sender: MessageSender::System, - content: "Welcome! A support agent will be with you shortly.".to_string(), + content: "Welcome! A supporter will be with you shortly.".to_string(), timestamp: now, }; Self { id, - customer_name, + user_name, status: SessionStatus::Waiting, created_at: now, messages: vec![welcome], diff --git a/src/state.rs b/src/state.rs index f50e172..ced64ec 100644 --- a/src/state.rs +++ b/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, // broadcast channel: session_id -> sender pub notifiers: DashMap>, - // agent broadcast for new sessions - pub agent_notifier: broadcast::Sender, + // supporter broadcast for new sessions and updates + pub supporter_notifier: broadcast::Sender, + // in-memory rate limits, wiped with all other state every three minutes + pub rate_limits: DashMap, + reset_started: Mutex, } 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()); + } } diff --git a/templates/index.html b/templates/index.html new file mode 100644 index 0000000..5ba4fcc --- /dev/null +++ b/templates/index.html @@ -0,0 +1,117 @@ + + + + + + Support Chat + + + + +
+ Reset in --:-- +

Support chat

+

Open the side you want to test. Each option opens in a new tab.

+ + +
+ + diff --git a/templates/supporter.html b/templates/supporter.html new file mode 100644 index 0000000..a95470c --- /dev/null +++ b/templates/supporter.html @@ -0,0 +1,487 @@ + + + + + + Support Dashboard + + + + + +
+
+

Support dashboard

+ Reset in --:-- +
+ + + +
+
+ Select a conversation + +
+ +
+
+

No conversation selected

+
+
+ +
+ + + +
+ +
+ + +
+
+
+ + + + + diff --git a/templates/user.html b/templates/user.html new file mode 100644 index 0000000..c55a869 --- /dev/null +++ b/templates/user.html @@ -0,0 +1,364 @@ + + + + + + Support Chat + + + + + Reset in --:-- + +
+

Talk to support

+

Start a conversation and a supporter can reply from the dashboard.

+
+ + +
+
+ +
+ + + +