armature_websocket/handler.rs
1//! WebSocket handler trait for implementing custom message handling.
2
3use crate::message::Message;
4use async_trait::async_trait;
5
6/// Trait for handling WebSocket events.
7///
8/// Implement this trait to define custom behavior for WebSocket connections.
9#[async_trait]
10pub trait WebSocketHandler: Send + Sync + 'static {
11 /// Called when a new client connects.
12 ///
13 /// # Arguments
14 /// * `connection_id` - The unique identifier for the connection
15 async fn on_connect(&self, connection_id: &str) {
16 let _ = connection_id;
17 }
18
19 /// Called when a message is received from a client.
20 ///
21 /// # Arguments
22 /// * `connection_id` - The unique identifier for the connection
23 /// * `message` - The received message
24 async fn on_message(&self, connection_id: &str, message: Message);
25
26 /// Called when a client disconnects.
27 ///
28 /// # Arguments
29 /// * `connection_id` - The unique identifier for the connection
30 async fn on_disconnect(&self, connection_id: &str) {
31 let _ = connection_id;
32 }
33
34 /// Called when an error occurs on a connection.
35 ///
36 /// # Arguments
37 /// * `connection_id` - The unique identifier for the connection
38 /// * `error` - The error that occurred
39 async fn on_error(&self, connection_id: &str, error: &crate::error::WebSocketError) {
40 tracing::error!(connection_id = %connection_id, error = %error, "WebSocket error");
41 }
42
43 /// Called when a ping is received. Return the pong payload.
44 ///
45 /// # Arguments
46 /// * `connection_id` - The unique identifier for the connection
47 /// * `payload` - The ping payload
48 ///
49 /// # Returns
50 /// The payload to send in the pong response
51 async fn on_ping(&self, connection_id: &str, payload: &[u8]) -> Vec<u8> {
52 let _ = connection_id;
53 payload.to_vec()
54 }
55
56 /// Called when a pong is received.
57 ///
58 /// # Arguments
59 /// * `connection_id` - The unique identifier for the connection
60 /// * `payload` - The pong payload
61 async fn on_pong(&self, connection_id: &str, payload: &[u8]) {
62 let _ = (connection_id, payload);
63 }
64}
65
66/// A no-op handler that logs messages.
67#[derive(Debug, Default, Clone)]
68pub struct LoggingHandler;
69
70#[async_trait]
71impl WebSocketHandler for LoggingHandler {
72 async fn on_connect(&self, connection_id: &str) {
73 tracing::info!(connection_id = %connection_id, "Client connected");
74 }
75
76 async fn on_message(&self, connection_id: &str, message: Message) {
77 tracing::debug!(
78 connection_id = %connection_id,
79 message_type = ?message.message_type,
80 payload_len = message.payload.len(),
81 "Received message"
82 );
83 }
84
85 async fn on_disconnect(&self, connection_id: &str) {
86 tracing::info!(connection_id = %connection_id, "Client disconnected");
87 }
88}
89