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

ri/ws/
server.rs
