ultrafast_gateway/dashboard/websocket/
message_handler.rs

1use crate::dashboard::websocket::rate_limiter::WebSocketRateLimiter;
2use crate::dashboard::websocket::subscription_manager::SubscriptionManager;
3use crate::dashboard::websocket::{ClientMessage, DashboardMessage, MessageType};
4use crate::gateway_error::GatewayError;
5use std::sync::Arc;
6use tokio::sync::mpsc;
7
8pub struct MessageHandler {
9    subscription_manager: Arc<SubscriptionManager>,
10    rate_limiter: Arc<WebSocketRateLimiter>,
11}
12
13impl MessageHandler {
14    pub fn new(
15        subscription_manager: Arc<SubscriptionManager>,
16        rate_limiter: Arc<WebSocketRateLimiter>,
17    ) -> Self {
18        Self {
19            subscription_manager,
20            rate_limiter,
21        }
22    }
23
24    pub async fn handle_message(
25        &self,
26        message: ClientMessage,
27        user_id: &str,
28        connection_id: &str,
29        response_tx: &mpsc::Sender<DashboardMessage>,
30    ) -> Result<(), GatewayError> {
31        // Check rate limit
32        if !self.rate_limiter.check_rate_limit(user_id).await? {
33            return Err(GatewayError::RateLimit {
34                message: "WebSocket rate limit exceeded".to_string(),
35            });
36        }
37
38        match message.message_type {
39            crate::dashboard::websocket::ClientMessageType::Subscribe => {
40                self.handle_subscribe(message.data, user_id, connection_id, response_tx)
41                    .await
42            }
43            crate::dashboard::websocket::ClientMessageType::Unsubscribe => {
44                self.handle_unsubscribe(message.data, user_id, connection_id, response_tx)
45                    .await
46            }
47            crate::dashboard::websocket::ClientMessageType::Ping => {
48                self.handle_ping(response_tx).await
49            }
50            crate::dashboard::websocket::ClientMessageType::RequestUpdate => {
51                self.handle_request_update(message.data, user_id, response_tx)
52                    .await
53            }
54        }
55    }
56
57    async fn handle_subscribe(
58        &self,
59        data: serde_json::Value,
60        user_id: &str,
61        connection_id: &str,
62        response_tx: &mpsc::Sender<DashboardMessage>,
63    ) -> Result<(), GatewayError> {
64        if let Some(topic) = data.get("topic").and_then(|t| t.as_str()) {
65            self.subscription_manager
66                .subscribe(user_id, connection_id, topic.to_string())
67                .await?;
68
69            let response = DashboardMessage {
70                message_type: MessageType::Subscribe,
71                data: serde_json::json!({
72                    "success": true,
73                    "topic": topic
74                }),
75                timestamp: chrono::Utc::now().timestamp(),
76                user_id: Some(user_id.to_string()),
77                topic: Some(topic.to_string()),
78            };
79
80            response_tx
81                .send(response)
82                .await
83                .map_err(|_| GatewayError::Internal {
84                    message: "Failed to send subscription response".to_string(),
85                })?;
86        }
87
88        Ok(())
89    }
90
91    async fn handle_unsubscribe(
92        &self,
93        data: serde_json::Value,
94        user_id: &str,
95        connection_id: &str,
96        response_tx: &mpsc::Sender<DashboardMessage>,
97    ) -> Result<(), GatewayError> {
98        if let Some(topic) = data.get("topic").and_then(|t| t.as_str()) {
99            self.subscription_manager
100                .unsubscribe(user_id, connection_id, topic)
101                .await?;
102
103            let response = DashboardMessage {
104                message_type: MessageType::Unsubscribe,
105                data: serde_json::json!({
106                    "success": true,
107                    "topic": topic
108                }),
109                timestamp: chrono::Utc::now().timestamp(),
110                user_id: Some(user_id.to_string()),
111                topic: Some(topic.to_string()),
112            };
113
114            response_tx
115                .send(response)
116                .await
117                .map_err(|_| GatewayError::Internal {
118                    message: "Failed to send unsubscription response".to_string(),
119                })?;
120        }
121
122        Ok(())
123    }
124
125    async fn handle_ping(
126        &self,
127        response_tx: &mpsc::Sender<DashboardMessage>,
128    ) -> Result<(), GatewayError> {
129        let response = DashboardMessage {
130            message_type: MessageType::Pong,
131            data: serde_json::json!({
132                "timestamp": chrono::Utc::now().timestamp()
133            }),
134            timestamp: chrono::Utc::now().timestamp(),
135            user_id: None,
136            topic: None,
137        };
138
139        response_tx
140            .send(response)
141            .await
142            .map_err(|_| GatewayError::Internal {
143                message: "Failed to send pong response".to_string(),
144            })?;
145
146        Ok(())
147    }
148
149    async fn handle_request_update(
150        &self,
151        _data: serde_json::Value,
152        user_id: &str,
153        response_tx: &mpsc::Sender<DashboardMessage>,
154    ) -> Result<(), GatewayError> {
155        // Fetch current metrics and send update
156        let metrics = crate::metrics::get_aggregated_metrics().await;
157
158        let response = DashboardMessage {
159            message_type: MessageType::Update,
160            data: serde_json::to_value(metrics)?,
161            timestamp: chrono::Utc::now().timestamp(),
162            user_id: Some(user_id.to_string()),
163            topic: None,
164        };
165
166        response_tx
167            .send(response)
168            .await
169            .map_err(|_| GatewayError::Internal {
170                message: "Failed to send update response".to_string(),
171            })?;
172
173        Ok(())
174    }
175}
ultrafast_gateway/dashboard/websocket/message_handler.rs

ultrafast_gateway/dashboard/websocket/
message_handler.rs
