ri/ws/
server.rs

1//! Copyright 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18//! # WebSocket Server Implementation
19
20use super::*;
21use uuid::Uuid;
22use tokio::sync::mpsc;
23use tokio::time::Duration;
24use futures::StreamExt;
25use tungstenite::Message;
26
27pub struct RiWSServer {
28    config: RiWSServerConfig,
29    stats: Arc<RwLock<RiWSServerStats>>,
30    session_manager: Arc<RiWSSessionManager>,
31    event_tx: Arc<RwLock<Option<broadcast::Sender<RiWSEvent>>>>,
32    shutdown_tx: Option<mpsc::Sender<()>>,
33    running: Arc<RwLock<bool>>,
34    handler: Arc<RwLock<Option<Arc<dyn RiWSSessionHandler>>>>,
35}
36
37#[cfg(feature = "pyo3")]
38#[pyclass]
39pub struct RiWSServerPy {
40    inner: RiWSServer,
41}
42
43#[cfg(feature = "pyo3")]
44#[pymethods]
45impl RiWSServerPy {
46    #[new]
47    fn new(config: RiWSServerConfig) -> Self {
48        Self {
49            inner: RiWSServer::new(config),
50        }
51    }
52
53    fn get_stats(&self) -> RiWSServerStats {
54        self.inner.get_stats()
55    }
56
57    fn is_running(&self) -> bool {
58        tokio::runtime::Handle::try_current()
59            .map(|handle| handle.block_on(async { self.inner.is_running().await }))
60            .unwrap_or(false)
61    }
62
63    fn start(&mut self) -> PyResult<()> {
64        let rt = tokio::runtime::Runtime::new()
65            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
66        
67        rt.block_on(async {
68            self.inner.start().await
69        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
70    }
71
72    fn stop(&mut self) -> PyResult<()> {
73        let rt = tokio::runtime::Runtime::new()
74            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
75        
76        rt.block_on(async {
77            self.inner.stop().await
78        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
79    }
80
81    fn broadcast(&self, data: Vec<u8>) -> usize {
82        tokio::runtime::Handle::try_current()
83            .map(|handle| handle.block_on(async { self.inner.broadcast(&data).await.unwrap_or(0) }))
84            .unwrap_or(0)
85    }
86
87    fn get_active_session_count(&self) -> usize {
88        tokio::runtime::Handle::try_current()
89            .map(|handle| handle.block_on(async { self.inner.get_active_session_count().await }))
90            .unwrap_or(0)
91    }
92}
93
94impl RiWSServer {
95    pub fn new(config: RiWSServerConfig) -> Self {
96        Self {
97            config,
98            stats: Arc::new(RwLock::new(RiWSServerStats::new())),
99            session_manager: Arc::new(RiWSSessionManager::new(1000)),
100            event_tx: Arc::new(RwLock::new(None)),
101            shutdown_tx: None,
102            running: Arc::new(RwLock::new(false)),
103            handler: Arc::new(RwLock::new(None)),
104        }
105    }
106
107    pub async fn set_handler<H: RiWSSessionHandler + 'static>(&self, handler: H) {
108        *self.handler.write().await = Some(Arc::new(handler));
109    }
110
111    pub async fn start(&mut self) -> RiResult<()> {
112        let addr: SocketAddr = format!("{}:{}", self.config.addr, self.config.port)
113            .parse()
114            .map_err(|e| WSError::Server {
115                message: format!("Invalid address: {}", e)
116            })?;
117
118        let listener = TcpListener::bind(&addr)
119            .await
120            .map_err(|e| WSError::Server {
121                message: format!("Failed to bind: {}", e)
122            })?;
123
124        let (event_tx, _) = broadcast::channel(100);
125        *self.event_tx.write().await = Some(event_tx);
126
127        let (shutdown_tx, shutdown_rx) = mpsc::channel(1);
128        self.shutdown_tx = Some(shutdown_tx);
129
130        let running = self.running.clone();
131        let stats = self.stats.clone();
132        let session_manager = self.session_manager.clone();
133        let handler = self.handler.clone();
134        let config = self.config.clone();
135
136        *running.write().await = true;
137
138        tokio::spawn(async move {
139            Self::accept_connections(
140                listener,
141                session_manager,
142                stats,
143                handler,
144                config,
145                shutdown_rx,
146                running,
147            ).await;
148        });
149
150        tracing::info!("WebSocket server started on {}", addr);
151        Ok(())
152    }
153
154    async fn accept_connections(
155        listener: TcpListener,
156        session_manager: Arc<RiWSSessionManager>,
157        stats: Arc<RwLock<RiWSServerStats>>,
158        handler: Arc<RwLock<Option<Arc<dyn RiWSSessionHandler>>>>,
159        config: RiWSServerConfig,
160        mut shutdown_rx: mpsc::Receiver<()>,
161        running: Arc<RwLock<bool>>,
162    ) {
163        let mut shutdown = false;
164
165        while !shutdown {
166            let result = listener.accept().await;
167            
168            if shutdown {
169                break;
170            }
171
172            match result {
173                Ok((stream, remote_addr)) => {
174                    let session_id = Uuid::new_v4().to_string();
175                    let remote_addr_str = remote_addr.to_string();
176                    
177                    tracing::info!("New WebSocket connection attempt: {} (session: {})", remote_addr_str, session_id);
178
179                    let allowed_origins = config.allowed_origins.clone();
180                    
181                    match Self::accept_with_origin_check(stream, &allowed_origins).await {
182                        Ok(ws_stream) => {
183                            let (_sender, receiver) = ws_stream.split();
184                            let (tx, rx) = mpsc::channel(100);
185
186                            let session = Arc::new(RiWSSession::new(
187                                session_id.clone(),
188                                tx,
189                                receiver,
190                                remote_addr_str.clone(),
191                            ));
192
193                            if session_manager.add_session(session.clone()).await.is_ok() {
194                                stats.write().await.record_connection();
195
196                                let handler_clone = handler.clone();
197                                let session_manager_clone = session_manager.clone();
198                                let stats_clone = stats.clone();
199                                let max_msg_size = config.max_message_size;
200
201                                tokio::spawn(async move {
202                                    Self::handle_session(
203                                        session.clone(),
204                                        rx,
205                                        handler_clone,
206                                        session_manager_clone,
207                                        stats_clone,
208                                        max_msg_size,
209                                    ).await;
210                                });
211                            } else {
212                                tracing::trace!("Failed to add session: {}", session_id);
213                            }
214                        }
215                        Err(e) => {
216                            tracing::error!("WebSocket upgrade failed: {}", e);
217                            stats.write().await.record_connection_error();
218                        }
219                    }
220                }
221                Err(e) => {
222                    tracing::error!("Failed to accept connection: {}", e);
223                    stats.write().await.record_connection_error();
224                }
225            }
226
227            tokio::time::sleep(Duration::from_secs(config.heartbeat_interval)).await;
228            
229            if !*running.read().await {
230                break;
231            }
232            
233            let _timeout = Duration::from_secs(config.heartbeat_timeout);
234            let sessions = session_manager.get_all_sessions().await;
235            for session_info in sessions {
236                let last_heartbeat_time = chrono::DateTime::from_timestamp(session_info.last_heartbeat as i64, 0)
237                    .unwrap_or_else(|| chrono::Utc::now());
238                let elapsed = chrono::Utc::now().signed_duration_since(last_heartbeat_time);
239                let elapsed_secs = elapsed.num_seconds() as u64;
240
241                if elapsed_secs > config.heartbeat_timeout {
242                    if let Some(session) = session_manager.get_session(&session_info.session_id).await {
243                        let _ = session.close().await;
244                    }
245                }
246            }
247            
248            if shutdown_rx.try_recv().is_ok() {
249                shutdown = true;
250            }
251        }
252
253        tracing::info!("WebSocket server stopped");
254    }
255
256    async fn handle_session(
257        session: Arc<RiWSSession>,
258        mut rx: mpsc::Receiver<std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>,
259        handler: Arc<RwLock<Option<Arc<dyn RiWSSessionHandler>>>>,
260        session_manager: Arc<RiWSSessionManager>,
261        stats: Arc<RwLock<RiWSServerStats>>,
262        max_message_size: usize,
263    ) {
264        let session_id = session.id.clone();
265
266        while let Some(message_result) = rx.recv().await {
267            match message_result {
268                Ok(message) => {
269                    match message {
270                        Message::Binary(data) => {
271                            // Security: Check message size to prevent memory exhaustion
272                            if data.len() > max_message_size {
273                                tracing::warn!(
274                                    "WebSocket message too large for session {}: {} bytes (max {} bytes)",
275                                    session_id, data.len(), max_message_size
276                                );
277                                stats.write().await.record_message_error();
278                                break;
279                            }
280                            
281                            stats.write().await.record_message_received(data.len());
282
283                            let handler_read = handler.read().await;
284                            if let Some(handler) = &*handler_read {
285                                if let Ok(response) = handler.on_message(&session_id, &data).await {
286                                    // Security: Check response size too
287                                    if response.len() > max_message_size {
288                                        tracing::warn!(
289                                            "WebSocket response too large for session {}: {} bytes",
290                                            session_id, response.len()
291                                        );
292                                        stats.write().await.record_message_error();
293                                    } else if session.send(&response).await.is_err() {
294                                        break;
295                                    }
296                                    stats.write().await.record_message_sent(response.len());
297                                }
298                            } else {
299                                if session.send(&data).await.is_err() {
300                                    break;
301                                }
302                                stats.write().await.record_message_sent(data.len());
303                            }
304                        }
305                        Message::Text(text) => {
306                            let data = text.into_bytes();
307                            
308                            // Security: Check message size to prevent memory exhaustion
309                            if data.len() > max_message_size {
310                                tracing::warn!(
311                                    "WebSocket message too large for session {}: {} bytes (max {} bytes)",
312                                    session_id, data.len(), max_message_size
313                                );
314                                stats.write().await.record_message_error();
315                                break;
316                            }
317                            
318                            stats.write().await.record_message_received(data.len());
319                            
320                            let handler_read = handler.read().await;
321                            if let Some(handler) = &*handler_read {
322                                if let Ok(response) = handler.on_message(&session_id, &data).await {
323                                    // Security: Check response size too
324                                    if response.len() > max_message_size {
325                                        tracing::warn!(
326                                            "WebSocket response too large for session {}: {} bytes",
327                                            session_id, response.len()
328                                        );
329                                        stats.write().await.record_message_error();
330                                    } else if session.send(&response).await.is_err() {
331                                        break;
332                                    }
333                                }
334                            }
335                        }
336                        Message::Ping(ping_data) => {
337                            // Security: Check ping data size
338                            if ping_data.len() > 125 {  // WebSocket ping limit
339                                tracing::warn!("WebSocket ping data too large for session {}", session_id);
340                                break;
341                            }
342                            if session.send(&ping_data).await.is_err() {
343                                break;
344                            }
345                        }
346                        Message::Pong(_) => {}
347                        Message::Close(_) => {
348                            break;
349                        }
350                        _ => {}
351                    }
352                }
353                Err(e) => {
354                    tracing::error!("WebSocket error for session {}: {}", session_id, e);
355                    stats.write().await.record_message_error();
356                    break;
357                }
358            }
359        }
360
361        session_manager.remove_session(&session_id).await;
362        stats.write().await.record_disconnection();
363
364        let handler_read = handler.read().await;
365        if let Some(handler) = &*handler_read {
366            let _ = handler.on_disconnect(&session_id).await;
367        }
368    }
369
370    pub async fn stop(&mut self) -> RiResult<()> {
371        *self.running.write().await = false;
372
373        let sessions = self.session_manager.get_all_sessions().await;
374        for session_info in sessions {
375            if let Some(session) = self.session_manager.get_session(&session_info.session_id).await {
376                let _ = session.close().await;
377            }
378        }
379
380        if let Some(tx) = self.shutdown_tx.take() {
381            tx.send(()).await.map_err(|e| WSError::Server {
382                message: format!("Shutdown error: {}", e)
383            })?;
384        }
385
386        tracing::info!("WebSocket server stopped");
387        Ok(())
388    }
389
390    pub fn get_stats(&self) -> RiWSServerStats {
391        self.stats.try_read()
392            .map(|guard| guard.clone())
393            .unwrap_or_else(|_| RiWSServerStats::new())
394    }
395
396    pub async fn get_session_info(&self, session_id: &str) -> Option<RiWSSessionInfo> {
397        self.session_manager.get_session(session_id).await
398            .map(|s| s.get_info())
399    }
400
401    pub async fn broadcast(&self, data: &[u8]) -> RiResult<usize> {
402        let count = self.session_manager.broadcast(data).await?;
403        self.stats.write().await.record_message_sent(data.len() * count);
404        Ok(count)
405    }
406
407    pub async fn is_running(&self) -> bool {
408        *self.running.read().await
409    }
410
411    pub async fn get_active_session_count(&self) -> usize {
412        self.session_manager.get_session_count().await
413    }
414
415    async fn accept_with_origin_check(
416        stream: tokio::net::TcpStream,
417        allowed_origins: &[String],
418    ) -> std::result::Result<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>, tungstenite::Error> {
419        use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
420        
421        let mut origin_valid = true;
422        let origins_empty = allowed_origins.is_empty();
423        let origins = allowed_origins.to_vec();
424        
425        let callback = |request: &Request, response: Response| {
426            if !origins_empty {
427                if let Some(origin_header) = request.headers().get("origin") {
428                    if let Ok(origin_str) = origin_header.to_str() {
429                        if !origins.contains(&origin_str.to_string()) {
430                            tracing::warn!("WebSocket connection rejected: Origin '{}' not in whitelist", origin_str);
431                            origin_valid = false;
432                            let bad_response = Response::builder()
433                                .status(403)
434                                .body(Some("Origin not allowed".to_string()))
435                                .unwrap();
436                            return Err(bad_response);
437                        }
438                    }
439                } else {
440                    tracing::warn!("WebSocket connection rejected: No Origin header present");
441                    origin_valid = false;
442                    let bad_response = Response::builder()
443                        .status(403)
444                        .body(Some("Origin header required".to_string()))
445                        .unwrap();
446                    return Err(bad_response);
447                }
448            }
449            Ok(response)
450        };
451        
452        let ws_stream = tokio_tungstenite::accept_hdr_async(stream, callback).await?;
453        Ok(ws_stream)
454    }
455}
456
457impl Clone for RiWSServer {
458    fn clone(&self) -> Self {
459        Self {
460            config: self.config.clone(),
461            stats: self.stats.clone(),
462            session_manager: self.session_manager.clone(),
463            event_tx: self.event_tx.clone(),
464            shutdown_tx: None,
465            running: self.running.clone(),
466            handler: self.handler.clone(),
467        }
468    }
469}
470
471impl Default for RiWSServer {
472    fn default() -> Self {
473        Self::new(RiWSServerConfig::default())
474    }
475}
ri/ws/server.rs

ri/ws/
server.rs
