security

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<String>,
     State(state): State<Arc<AppState>>,
 ) -> 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<AppState>) {
+async fn handle_user_socket(mut socket: WebSocket, session_id: String, state: Arc<AppState>) {
     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<String>,
+    Path(_supporter_id): Path<String>,
     State(state): State<Arc<AppState>>,
 ) -> 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<AppState>) {
-    let mut rx = state.agent_notifier.subscribe();
+async fn handle_supporter_socket(mut socket: WebSocket, state: Arc<AppState>) {
+    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,
                     _ => {}
                 }