sz_orm_websocket/
handler.rs

1use crate::error::WsError;
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8#[derive(Debug, Clone)]
9pub struct WebSocketMessage {
10    pub msg_type: MessageType,
11    pub payload: Vec<u8>,
12    pub sender_id: Option<i64>,
13    pub room_id: Option<String>,
14    pub timestamp: i64,
15}
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
18pub enum MessageType {
19    #[default]
20    Text,
21    Binary,
22    Ping,
23    Pong,
24    Join,
25    Leave,
26    Subscribe,
27    Unsubscribe,
28    Notification,
29    System,
30}
31
32#[derive(Debug, Clone)]
33pub struct WebSocketConnection {
34    pub id: String,
35    pub user_id: Option<i64>,
36    pub remote_addr: Option<String>,
37    pub is_authenticated: bool,
38    pub subscriptions: Vec<String>,
39}
40
41impl WebSocketConnection {
42    pub fn new(id: impl Into<String>) -> Self {
43        Self {
44            id: id.into(),
45            user_id: None,
46            remote_addr: None,
47            is_authenticated: false,
48            subscriptions: Vec::new(),
49        }
50    }
51
52    pub fn with_user(mut self, user_id: i64) -> Self {
53        self.user_id = Some(user_id);
54        self.is_authenticated = true;
55        self
56    }
57
58    pub fn with_address(mut self, addr: impl Into<String>) -> Self {
59        self.remote_addr = Some(addr.into());
60        self
61    }
62
63    pub fn subscribe(&mut self, room: impl Into<String>) {
64        let room = room.into();
65        if !self.subscriptions.contains(&room) {
66            self.subscriptions.push(room);
67        }
68    }
69
70    pub fn unsubscribe(&mut self, room: &str) {
71        self.subscriptions.retain(|r| r != room);
72    }
73}
74
75#[async_trait]
76pub trait WebSocketHandler: Send + Sync {
77    async fn on_message(
78        &self,
79        conn: &WebSocketConnection,
80        msg: WebSocketMessage,
81    ) -> Result<Option<WebSocketMessage>, WsError>;
82
83    async fn on_connect(&self, conn: &WebSocketConnection) -> Result<(), WsError>;
84
85    async fn on_disconnect(&self, conn: &WebSocketConnection);
86
87    fn authenticate(&self, token: &str) -> Result<UserId, WsError>;
88}
89
90pub type UserId = i64;
91
92pub struct WsContext {
93    pub connection_id: String,
94    pub user_id: Option<i64>,
95    pub metadata: std::collections::HashMap<String, String>,
96}
97
98impl WsContext {
99    pub fn new(connection_id: impl Into<String>) -> Self {
100        Self {
101            connection_id: connection_id.into(),
102            user_id: None,
103            metadata: std::collections::HashMap::new(),
104        }
105    }
106
107    pub fn with_user(mut self, user_id: i64) -> Self {
108        self.user_id = Some(user_id);
109        self
110    }
111
112    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
113        self.metadata.insert(key.into(), value.into());
114        self
115    }
116}
117
118pub struct WsMessageBuilder {
119    msg_type: MessageType,
120    payload: Vec<u8>,
121    sender_id: Option<i64>,
122    room_id: Option<String>,
123}
124
125impl WsMessageBuilder {
126    pub fn new() -> Self {
127        Self {
128            msg_type: MessageType::Text,
129            payload: Vec::new(),
130            sender_id: None,
131            room_id: None,
132        }
133    }
134
135    pub fn text(mut self, text: impl Into<String>) -> Self {
136        self.msg_type = MessageType::Text;
137        self.payload = text.into().into_bytes();
138        self
139    }
140
141    pub fn binary(mut self, data: Vec<u8>) -> Self {
142        self.msg_type = MessageType::Binary;
143        self.payload = data;
144        self
145    }
146
147    pub fn json<T: serde::Serialize>(mut self, data: &T) -> Result<Self, WsError> {
148        self.msg_type = MessageType::Text;
149        self.payload = serde_json::to_vec(data)?;
150        Ok(self)
151    }
152
153    pub fn with_sender(mut self, user_id: i64) -> Self {
154        self.sender_id = Some(user_id);
155        self
156    }
157
158    pub fn with_room(mut self, room: impl Into<String>) -> Self {
159        self.room_id = Some(room.into());
160        self
161    }
162
163    pub fn notification(mut self) -> Self {
164        self.msg_type = MessageType::Notification;
165        self
166    }
167
168    pub fn system(mut self) -> Self {
169        self.msg_type = MessageType::System;
170        self
171    }
172
173    pub fn build(self) -> WebSocketMessage {
174        WebSocketMessage {
175            msg_type: self.msg_type,
176            payload: self.payload,
177            sender_id: self.sender_id,
178            room_id: self.room_id,
179            timestamp: current_timestamp(),
180        }
181    }
182}
183
184impl Default for WsMessageBuilder {
185    fn default() -> Self {
186        Self::new()
187    }
188}
189
190fn current_timestamp() -> i64 {
191    use std::time::{SystemTime, UNIX_EPOCH};
192    SystemTime::now()
193        .duration_since(UNIX_EPOCH)
194        .unwrap_or_default()
195        .as_millis() as i64
196}
197
198/// Default WebSocket handler that tracks connections and echoes messages.
199///
200/// - Text messages are echoed back to the sender with the sender's user_id.
201/// - Ping messages are answered with a Pong.
202/// - Subscribe/Unsubscribe/Join/Leave messages return a System acknowledgement.
203/// - All other messages are logged but produce no response.
204pub struct DefaultWebSocketHandler {
205    connections: Arc<RwLock<HashMap<String, WebSocketConnection>>>,
206    message_log: Arc<RwLock<Vec<WebSocketMessage>>>,
207}
208
209impl DefaultWebSocketHandler {
210    pub fn new() -> Self {
211        Self {
212            connections: Arc::new(RwLock::new(HashMap::new())),
213            message_log: Arc::new(RwLock::new(Vec::new())),
214        }
215    }
216
217    pub async fn connection_count(&self) -> usize {
218        self.connections.read().await.len()
219    }
220
221    pub async fn is_connected(&self, connection_id: &str) -> bool {
222        self.connections.read().await.contains_key(connection_id)
223    }
224
225    pub async fn message_count(&self) -> usize {
226        self.message_log.read().await.len()
227    }
228
229    pub async fn messages(&self) -> Vec<WebSocketMessage> {
230        self.message_log.read().await.clone()
231    }
232
233    pub async fn get_connection(&self, connection_id: &str) -> Option<WebSocketConnection> {
234        self.connections.read().await.get(connection_id).cloned()
235    }
236}
237
238impl Default for DefaultWebSocketHandler {
239    fn default() -> Self {
240        Self::new()
241    }
242}
243
244#[async_trait]
245impl WebSocketHandler for DefaultWebSocketHandler {
246    async fn on_message(
247        &self,
248        conn: &WebSocketConnection,
249        msg: WebSocketMessage,
250    ) -> Result<Option<WebSocketMessage>, WsError> {
251        self.message_log.write().await.push(msg.clone());
252
253        match msg.msg_type {
254            MessageType::Text => {
255                let response = WebSocketMessage {
256                    msg_type: MessageType::Text,
257                    payload: msg.payload.clone(),
258                    sender_id: conn.user_id,
259                    room_id: None,
260                    timestamp: current_timestamp(),
261                };
262                Ok(Some(response))
263            }
264            MessageType::Ping => {
265                let response = WebSocketMessage {
266                    msg_type: MessageType::Pong,
267                    payload: msg.payload.clone(),
268                    sender_id: None,
269                    room_id: None,
270                    timestamp: current_timestamp(),
271                };
272                Ok(Some(response))
273            }
274            MessageType::Subscribe => {
275                let room = String::from_utf8_lossy(&msg.payload).to_string();
276                let ack = format!("subscribed:{}", room);
277                let response = WebSocketMessage {
278                    msg_type: MessageType::System,
279                    payload: ack.into_bytes(),
280                    sender_id: None,
281                    room_id: Some(room),
282                    timestamp: current_timestamp(),
283                };
284                Ok(Some(response))
285            }
286            MessageType::Unsubscribe => {
287                let room = String::from_utf8_lossy(&msg.payload).to_string();
288                let ack = format!("unsubscribed:{}", room);
289                let response = WebSocketMessage {
290                    msg_type: MessageType::System,
291                    payload: ack.into_bytes(),
292                    sender_id: None,
293                    room_id: Some(room),
294                    timestamp: current_timestamp(),
295                };
296                Ok(Some(response))
297            }
298            MessageType::Join => {
299                let room = String::from_utf8_lossy(&msg.payload).to_string();
300                let ack = format!("joined:{}", room);
301                let response = WebSocketMessage {
302                    msg_type: MessageType::System,
303                    payload: ack.into_bytes(),
304                    sender_id: conn.user_id,
305                    room_id: Some(room),
306                    timestamp: current_timestamp(),
307                };
308                Ok(Some(response))
309            }
310            MessageType::Leave => {
311                let room = String::from_utf8_lossy(&msg.payload).to_string();
312                let ack = format!("left:{}", room);
313                let response = WebSocketMessage {
314                    msg_type: MessageType::System,
315                    payload: ack.into_bytes(),
316                    sender_id: conn.user_id,
317                    room_id: Some(room),
318                    timestamp: current_timestamp(),
319                };
320                Ok(Some(response))
321            }
322            _ => Ok(None),
323        }
324    }
325
326    async fn on_connect(&self, conn: &WebSocketConnection) -> Result<(), WsError> {
327        self.connections
328            .write()
329            .await
330            .insert(conn.id.clone(), conn.clone());
331        Ok(())
332    }
333
334    async fn on_disconnect(&self, conn: &WebSocketConnection) {
335        self.connections.write().await.remove(&conn.id);
336    }
337
338    fn authenticate(&self, token: &str) -> Result<UserId, WsError> {
339        if let Some(id_str) = token.strip_prefix("user_id:") {
340            id_str.parse::<i64>().map_err(|_| {
341                WsError::Authentication(format!("invalid user_id in token: {}", token))
342            })
343        } else {
344            token
345                .parse::<i64>()
346                .map_err(|_| WsError::Authentication(format!("invalid token: {}", token)))
347        }
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354
355    fn make_conn(id: &str, user_id: Option<i64>) -> WebSocketConnection {
356        let mut conn = WebSocketConnection::new(id);
357        if let Some(uid) = user_id {
358            conn = conn.with_user(uid);
359        }
360        conn
361    }
362
363    fn make_text(payload: &[u8]) -> WebSocketMessage {
364        WebSocketMessage {
365            msg_type: MessageType::Text,
366            payload: payload.to_vec(),
367            sender_id: None,
368            room_id: None,
369            timestamp: 1000,
370        }
371    }
372
373    #[tokio::test]
374    async fn test_on_connect_tracks_connection() {
375        let handler = DefaultWebSocketHandler::new();
376        assert_eq!(handler.connection_count().await, 0);
377
378        let conn = make_conn("c1", Some(123));
379        handler.on_connect(&conn).await.unwrap();
380
381        assert_eq!(handler.connection_count().await, 1);
382        assert!(handler.is_connected("c1").await);
383    }
384
385    #[tokio::test]
386    async fn test_on_disconnect_removes_connection() {
387        let handler = DefaultWebSocketHandler::new();
388        let conn = make_conn("c1", Some(123));
389        handler.on_connect(&conn).await.unwrap();
390        assert!(handler.is_connected("c1").await);
391
392        handler.on_disconnect(&conn).await;
393        assert!(!handler.is_connected("c1").await);
394        assert_eq!(handler.connection_count().await, 0);
395    }
396
397    #[tokio::test]
398    async fn test_on_message_text_echoes_back() {
399        let handler = DefaultWebSocketHandler::new();
400        let conn = make_conn("c1", Some(42));
401        let msg = make_text(b"hello");
402
403        let response = handler.on_message(&conn, msg).await.unwrap();
404        assert!(response.is_some());
405
406        let resp = response.unwrap();
407        assert_eq!(resp.msg_type, MessageType::Text);
408        assert_eq!(resp.payload, b"hello");
409        assert_eq!(resp.sender_id, Some(42));
410    }
411
412    #[tokio::test]
413    async fn test_on_message_ping_responds_pong() {
414        let handler = DefaultWebSocketHandler::new();
415        let conn = make_conn("c1", None);
416        let msg = WebSocketMessage {
417            msg_type: MessageType::Ping,
418            payload: b"ping".to_vec(),
419            sender_id: None,
420            room_id: None,
421            timestamp: 1,
422        };
423
424        let response = handler.on_message(&conn, msg).await.unwrap();
425        let resp = response.unwrap();
426        assert_eq!(resp.msg_type, MessageType::Pong);
427        assert_eq!(resp.payload, b"ping");
428    }
429
430    #[tokio::test]
431    async fn test_on_message_subscribe_returns_system_ack() {
432        let handler = DefaultWebSocketHandler::new();
433        let conn = make_conn("c1", None);
434        let msg = WebSocketMessage {
435            msg_type: MessageType::Subscribe,
436            payload: b"room1".to_vec(),
437            sender_id: None,
438            room_id: None,
439            timestamp: 1,
440        };
441
442        let response = handler.on_message(&conn, msg).await.unwrap();
443        let resp = response.unwrap();
444        assert_eq!(resp.msg_type, MessageType::System);
445        assert_eq!(resp.payload, b"subscribed:room1");
446        assert_eq!(resp.room_id, Some("room1".to_string()));
447    }
448
449    #[tokio::test]
450    async fn test_on_message_unsubscribe_returns_system_ack() {
451        let handler = DefaultWebSocketHandler::new();
452        let conn = make_conn("c1", None);
453        let msg = WebSocketMessage {
454            msg_type: MessageType::Unsubscribe,
455            payload: b"room1".to_vec(),
456            sender_id: None,
457            room_id: None,
458            timestamp: 1,
459        };
460
461        let response = handler.on_message(&conn, msg).await.unwrap();
462        let resp = response.unwrap();
463        assert_eq!(resp.msg_type, MessageType::System);
464        assert_eq!(resp.payload, b"unsubscribed:room1");
465    }
466
467    #[tokio::test]
468    async fn test_on_message_join_returns_system_ack() {
469        let handler = DefaultWebSocketHandler::new();
470        let conn = make_conn("c1", Some(7));
471        let msg = WebSocketMessage {
472            msg_type: MessageType::Join,
473            payload: b"lobby".to_vec(),
474            sender_id: None,
475            room_id: None,
476            timestamp: 1,
477        };
478
479        let response = handler.on_message(&conn, msg).await.unwrap();
480        let resp = response.unwrap();
481        assert_eq!(resp.msg_type, MessageType::System);
482        assert_eq!(resp.payload, b"joined:lobby");
483        assert_eq!(resp.sender_id, Some(7));
484        assert_eq!(resp.room_id, Some("lobby".to_string()));
485    }
486
487    #[tokio::test]
488    async fn test_on_message_leave_returns_system_ack() {
489        let handler = DefaultWebSocketHandler::new();
490        let conn = make_conn("c1", Some(7));
491        let msg = WebSocketMessage {
492            msg_type: MessageType::Leave,
493            payload: b"lobby".to_vec(),
494            sender_id: None,
495            room_id: None,
496            timestamp: 1,
497        };
498
499        let response = handler.on_message(&conn, msg).await.unwrap();
500        let resp = response.unwrap();
501        assert_eq!(resp.payload, b"left:lobby");
502    }
503
504    #[tokio::test]
505    async fn test_on_message_binary_returns_none() {
506        let handler = DefaultWebSocketHandler::new();
507        let conn = make_conn("c1", None);
508        let msg = WebSocketMessage {
509            msg_type: MessageType::Binary,
510            payload: vec![1, 2, 3],
511            sender_id: None,
512            room_id: None,
513            timestamp: 1,
514        };
515
516        let response = handler.on_message(&conn, msg).await.unwrap();
517        assert!(response.is_none());
518    }
519
520    #[tokio::test]
521    async fn test_message_log_records_all_messages() {
522        let handler = DefaultWebSocketHandler::new();
523        let conn = make_conn("c1", None);
524
525        handler.on_message(&conn, make_text(b"m1")).await.unwrap();
526        handler
527            .on_message(
528                &conn,
529                WebSocketMessage {
530                    msg_type: MessageType::Binary,
531                    payload: vec![1],
532                    sender_id: None,
533                    room_id: None,
534                    timestamp: 2,
535                },
536            )
537            .await
538            .unwrap();
539        handler.on_message(&conn, make_text(b"m3")).await.unwrap();
540
541        assert_eq!(handler.message_count().await, 3);
542        let msgs = handler.messages().await;
543        assert_eq!(msgs[0].payload, b"m1");
544        assert_eq!(msgs[1].msg_type, MessageType::Binary);
545        assert_eq!(msgs[2].payload, b"m3");
546    }
547
548    #[tokio::test]
549    async fn test_authenticate_valid_numeric_token() {
550        let handler = DefaultWebSocketHandler::new();
551        let user_id = handler.authenticate("12345").unwrap();
552        assert_eq!(user_id, 12345);
553    }
554
555    #[tokio::test]
556    async fn test_authenticate_valid_prefixed_token() {
557        let handler = DefaultWebSocketHandler::new();
558        let user_id = handler.authenticate("user_id:67890").unwrap();
559        assert_eq!(user_id, 67890);
560    }
561
562    #[test]
563    fn test_authenticate_invalid_token_returns_error() {
564        let handler = DefaultWebSocketHandler::new();
565        let result = handler.authenticate("not-a-number");
566        assert!(result.is_err());
567        assert!(matches!(result.unwrap_err(), WsError::Authentication(_)));
568    }
569
570    #[test]
571    fn test_authenticate_invalid_prefixed_token_returns_error() {
572        let handler = DefaultWebSocketHandler::new();
573        let result = handler.authenticate("user_id:abc");
574        assert!(result.is_err());
575        assert!(matches!(result.unwrap_err(), WsError::Authentication(_)));
576    }
577
578    #[tokio::test]
579    async fn test_get_connection_returns_stored_connection() {
580        let handler = DefaultWebSocketHandler::new();
581        let conn = make_conn("c1", Some(99));
582        handler.on_connect(&conn).await.unwrap();
583
584        let retrieved = handler.get_connection("c1").await.unwrap();
585        assert_eq!(retrieved.id, "c1");
586        assert_eq!(retrieved.user_id, Some(99));
587        assert!(retrieved.is_authenticated);
588    }
589
590    #[tokio::test]
591    async fn test_get_connection_not_found() {
592        let handler = DefaultWebSocketHandler::new();
593        assert!(handler.get_connection("missing").await.is_none());
594    }
595
596    #[tokio::test]
597    async fn test_multiple_connections_tracked_independently() {
598        let handler = DefaultWebSocketHandler::new();
599        let conn1 = make_conn("c1", Some(1));
600        let conn2 = make_conn("c2", Some(2));
601
602        handler.on_connect(&conn1).await.unwrap();
603        handler.on_connect(&conn2).await.unwrap();
604
605        assert_eq!(handler.connection_count().await, 2);
606
607        handler.on_disconnect(&conn1).await;
608        assert_eq!(handler.connection_count().await, 1);
609        assert!(!handler.is_connected("c1").await);
610        assert!(handler.is_connected("c2").await);
611    }
612
613    #[tokio::test]
614    async fn test_default_impl_creates_empty_handler() {
615        let handler = DefaultWebSocketHandler::default();
616        assert_eq!(handler.connection_count().await, 0);
617        assert_eq!(handler.message_count().await, 0);
618    }
619}
sz_orm_websocket/handler.rs

sz_orm_websocket/
handler.rs
