ri/ws/
mod.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 Support
19
20use crate::core::{RiResult, RiError};
21use async_trait::async_trait;
22use std::sync::Arc;
23use tokio::sync::RwLock;
24use std::net::SocketAddr;
25use futures::stream::SplitStream;
26use tokio::net::TcpListener;
27use tokio::sync::broadcast;
28use std::collections::HashMap as FxHashMap;
29use tungstenite::Message;
30
31#[cfg(feature = "pyo3")]
32use pyo3::prelude::*;
33
34#[cfg(feature = "websocket")]
35mod server;
36
37#[cfg(feature = "websocket")]
38mod client;
39
40#[cfg(feature = "websocket")]
41pub use server::RiWSServer;
42
43#[cfg(feature = "websocket")]
44pub use client::RiWSClient;
45
46#[cfg(feature = "websocket")]
47pub use client::RiWSClientConfig;
48
49#[cfg(feature = "websocket")]
50pub use client::RiWSClientStats;
51
52#[cfg(all(feature = "websocket", feature = "pyo3"))]
53pub use server::RiWSServerPy;
54
55#[cfg(all(feature = "websocket", feature = "pyo3"))]
56pub use client::RiWSClientPy;
57
58#[derive(Debug, Clone, PartialEq, Eq)]
59#[cfg_attr(feature = "pyo3", pyclass)]
60pub struct RiWSServerConfig {
61    pub addr: String,
62    pub port: u16,
63    pub max_connections: usize,
64    pub heartbeat_interval: u64,
65    pub heartbeat_timeout: u64,
66    pub max_message_size: usize,
67    pub ping_interval: u64,
68    pub allowed_origins: Vec<String>,
69}
70
71#[cfg(feature = "pyo3")]
72#[pymethods]
73impl RiWSServerConfig {
74    #[new]
75    fn new() -> Self {
76        Self::default()
77    }
78    
79    #[getter]
80    fn get_addr(&self) -> String {
81        self.addr.clone()
82    }
83    
84    #[setter]
85    fn set_addr(&mut self, addr: String) {
86        self.addr = addr;
87    }
88    
89    #[getter]
90    fn get_port(&self) -> u16 {
91        self.port
92    }
93    
94    #[setter]
95    fn set_port(&mut self, port: u16) {
96        self.port = port;
97    }
98}
99
100impl Default for RiWSServerConfig {
101    fn default() -> Self {
102        Self {
103            addr: "127.0.0.1".to_string(),
104            port: 8080,
105            max_connections: 1000,
106            heartbeat_interval: 30,
107            heartbeat_timeout: 60,
108            max_message_size: 65536,
109            ping_interval: 25,
110            allowed_origins: vec![],
111        }
112    }
113}
114
115#[derive(Debug, Clone, PartialEq, Eq)]
116#[cfg_attr(feature = "pyo3", pyclass)]
117pub enum RiWSEvent {
118    Connected { session_id: String },
119    Disconnected { session_id: String },
120    Message { session_id: String, data: Vec<u8> },
121    Error { session_id: String, message: String },
122}
123
124#[derive(Debug, Clone)]
125#[cfg_attr(feature = "pyo3", pyclass)]
126pub struct RiWSSessionInfo {
127    pub session_id: String,
128    pub remote_addr: String,
129    pub connected_at: u64,
130    pub messages_sent: u64,
131    pub messages_received: u64,
132    pub bytes_sent: u64,
133    pub bytes_received: u64,
134    pub is_active: bool,
135    pub last_heartbeat: u64,
136}
137
138#[cfg(feature = "pyo3")]
139#[pymethods]
140impl RiWSSessionInfo {
141    #[getter]
142    fn get_session_id(&self) -> String {
143        self.session_id.clone()
144    }
145    
146    #[getter]
147    fn get_remote_addr(&self) -> String {
148        self.remote_addr.clone()
149    }
150    
151    #[getter]
152    fn get_connected_at(&self) -> u64 {
153        self.connected_at
154    }
155    
156    #[getter]
157    fn get_messages_sent(&self) -> u64 {
158        self.messages_sent
159    }
160    
161    #[getter]
162    fn get_messages_received(&self) -> u64 {
163        self.messages_received
164    }
165    
166    #[getter]
167    fn get_bytes_sent(&self) -> u64 {
168        self.bytes_sent
169    }
170    
171    #[getter]
172    fn get_bytes_received(&self) -> u64 {
173        self.bytes_received
174    }
175    
176    #[getter]
177    fn get_is_active(&self) -> bool {
178        self.is_active
179    }
180    
181    #[getter]
182    fn get_last_heartbeat(&self) -> u64 {
183        self.last_heartbeat
184    }
185}
186
187impl Default for RiWSSessionInfo {
188    fn default() -> Self {
189        Self {
190            session_id: String::new(),
191            remote_addr: String::new(),
192            connected_at: 0,
193            messages_sent: 0,
194            messages_received: 0,
195            bytes_sent: 0,
196            bytes_received: 0,
197            is_active: false,
198            last_heartbeat: 0,
199        }
200    }
201}
202
203#[async_trait]
204pub trait RiWSSessionHandler: Send + Sync {
205    async fn on_connect(&self, session_id: &str, remote_addr: &str) -> RiResult<()>;
206    async fn on_disconnect(&self, session_id: &str) -> RiResult<()>;
207    async fn on_message(&self, session_id: &str, data: &[u8]) -> RiResult<Vec<u8>>;
208    async fn on_error(&self, session_id: &str, error: &str) -> RiResult<()>;
209}
210
211#[derive(Debug, thiserror::Error)]
212pub enum WSError {
213    #[error("Server error: {message}")]
214    Server { message: String },
215    #[error("Session error: {message}")]
216    Session { message: String },
217    #[error("Connection error: {message}")]
218    Connection { message: String },
219    #[error("Message too large: {size} bytes (max: {max_size})")]
220    MessageTooLarge { size: usize, max_size: usize },
221    #[error("Session not found: {session_id}")]
222    SessionNotFound { session_id: String },
223    #[error("Invalid message format")]
224    InvalidFormat,
225}
226
227impl From<WSError> for RiError {
228    fn from(error: WSError) -> Self {
229        RiError::Other(format!("WebSocket error: {}", error))
230    }
231}
232
233pub struct RiWSSession {
234    pub id: String,
235    pub sender: tokio::sync::mpsc::Sender<std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>,
236    pub receiver: SplitStream<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>>,
237    pub info: Arc<RwLock<RiWSSessionInfo>>,
238}
239
240impl RiWSSession {
241    pub fn new(
242        id: String,
243        sender: tokio::sync::mpsc::Sender<std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>,
244        receiver: SplitStream<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>>,
245        remote_addr: String,
246    ) -> Self {
247            let now = chrono::Utc::now().timestamp() as u64;
248        let session_id = id.clone();
249        Self {
250            id,
251            sender,
252            receiver,
253            info: Arc::new(RwLock::new(RiWSSessionInfo {
254                session_id,
255                remote_addr,
256                connected_at: now,
257                messages_sent: 0,
258                messages_received: 0,
259                bytes_sent: 0,
260                bytes_received: 0,
261                is_active: true,
262                last_heartbeat: now,
263            })),
264        }
265    }
266
267    pub async fn send(&self, data: &[u8]) -> RiResult<()> {
268        let message = Message::Binary(data.to_vec());
269        
270        self.sender.send(Ok(message))
271            .await
272            .map_err(|e| WSError::Session {
273                message: format!("Failed to send message: {}", e)
274            })?;
275
276        let mut info = self.info.write().await;
277        info.messages_sent += 1;
278        info.bytes_sent += data.len() as u64;
279
280        Ok(())
281    }
282
283    pub async fn send_text(&self, text: &str) -> RiResult<()> {
284        let message = Message::Text(text.to_string());
285        
286        self.sender.send(Ok(message))
287            .await
288            .map_err(|e| WSError::Session {
289                message: format!("Failed to send message: {}", e)
290            })?;
291
292        let mut info = self.info.write().await;
293        info.messages_sent += 1;
294        info.bytes_sent += text.len() as u64;
295
296        Ok(())
297    }
298
299    pub async fn close(&self) -> RiResult<()> {
300        self.sender.send(Ok(Message::Close(None)))
301            .await
302            .map_err(|e| WSError::Session {
303                message: format!("Failed to close session: {}", e)
304            })?;
305
306        let mut info = self.info.write().await;
307        info.is_active = false;
308
309        Ok(())
310    }
311
312    pub fn get_info(&self) -> RiWSSessionInfo {
313        self.info.try_read()
314            .map(|guard| guard.clone())
315            .unwrap_or_else(|_| RiWSSessionInfo::default())
316    }
317}
318
319pub struct RiWSSessionManager {
320    sessions: Arc<RwLock<FxHashMap<String, Arc<RiWSSession>>>>,
321    max_connections: usize,
322}
323
324impl Clone for RiWSSessionManager {
325    fn clone(&self) -> Self {
326        Self {
327            sessions: self.sessions.clone(),
328            max_connections: self.max_connections,
329        }
330    }
331}
332
333impl RiWSSessionManager {
334    pub fn new(max_connections: usize) -> Self {
335        Self {
336            sessions: Arc::new(RwLock::new(FxHashMap::new())),
337            max_connections,
338        }
339    }
340
341    pub async fn add_session(&self, session: Arc<RiWSSession>) -> RiResult<()> {
342        let mut sessions = self.sessions.write().await;
343        
344        if sessions.len() >= self.max_connections {
345            return Err(WSError::Session {
346                message: format!("Max connections reached: {}", self.max_connections)
347            }.into());
348        }
349
350        sessions.insert(session.id.clone(), session);
351        Ok(())
352    }
353
354    pub async fn remove_session(&self, session_id: &str) {
355        let mut sessions = self.sessions.write().await;
356        sessions.remove(session_id);
357    }
358
359    pub async fn get_session(&self, session_id: &str) -> Option<Arc<RiWSSession>> {
360        let sessions = self.sessions.read().await;
361        sessions.get(session_id).cloned()
362    }
363
364    pub async fn broadcast(&self, data: &[u8]) -> RiResult<usize> {
365        let sessions = self.sessions.read().await;
366        let mut count = 0;
367
368        for session in sessions.values() {
369            if session.send(data).await.is_ok() {
370                count += 1;
371            }
372        }
373
374        Ok(count)
375    }
376
377    pub async fn get_session_count(&self) -> usize {
378        self.sessions.read().await.len()
379    }
380
381    pub async fn get_all_sessions(&self) -> Vec<RiWSSessionInfo> {
382        let sessions = self.sessions.read().await;
383        sessions.values().map(|s| s.get_info()).collect()
384    }
385}
386
387#[cfg(feature = "pyo3")]
388#[pyclass]
389pub struct RiWSPythonHandler {
390    on_connect: Arc<Py<PyAny>>,
391    on_disconnect: Arc<Py<PyAny>>,
392    on_message: Arc<Py<PyAny>>,
393    on_error: Arc<Py<PyAny>>,
394}
395
396#[cfg(feature = "pyo3")]
397#[pymethods]
398impl RiWSPythonHandler {
399    #[new]
400    fn new(
401        on_connect: Py<PyAny>,
402        on_disconnect: Py<PyAny>,
403        on_message: Py<PyAny>,
404        on_error: Py<PyAny>,
405    ) -> Self {
406        Self {
407            on_connect: Arc::new(on_connect),
408            on_disconnect: Arc::new(on_disconnect),
409            on_message: Arc::new(on_message),
410            on_error: Arc::new(on_error),
411        }
412    }
413}
414
415#[cfg(feature = "pyo3")]
416#[async_trait]
417impl RiWSSessionHandler for RiWSPythonHandler {
418    async fn on_connect(&self, session_id: &str, remote_addr: &str) -> RiResult<()> {
419        let on_connect = Arc::clone(&self.on_connect);
420        let session_id = session_id.to_string();
421        let remote_addr = remote_addr.to_string();
422        
423        tokio::task::spawn_blocking(move || {
424            Python::attach(|py| {
425                let handler = on_connect.clone_ref(py);
426                let _ = handler.call(py, (session_id, remote_addr), None);
427            });
428        }).await.ok();
429        
430        Ok(())
431    }
432    
433    async fn on_disconnect(&self, session_id: &str) -> RiResult<()> {
434        let on_disconnect = Arc::clone(&self.on_disconnect);
435        let session_id = session_id.to_string();
436        
437        tokio::task::spawn_blocking(move || {
438            Python::attach(|py| {
439                let handler = on_disconnect.clone_ref(py);
440                let _ = handler.call(py, (session_id,), None);
441            });
442        }).await.ok();
443        
444        Ok(())
445    }
446    
447    async fn on_message(&self, session_id: &str, data: &[u8]) -> RiResult<Vec<u8>> {
448        let on_message = Arc::clone(&self.on_message);
449        let session_id = session_id.to_string();
450        let data_vec = data.to_vec();
451        
452        let result = tokio::task::spawn_blocking(move || {
453            Python::attach(|py| {
454                let handler = on_message.clone_ref(py);
455                match handler.call(py, (session_id, data_vec), None) {
456                    Ok(obj) => obj.extract::<Vec<u8>>(py).ok(),
457                    Err(_) => None,
458                }
459            })
460        }).await.ok().flatten();
461        
462        Ok(result.unwrap_or_default())
463    }
464    
465    async fn on_error(&self, session_id: &str, error: &str) -> RiResult<()> {
466        let on_error = Arc::clone(&self.on_error);
467        let session_id = session_id.to_string();
468        let error = error.to_string();
469        
470        tokio::task::spawn_blocking(move || {
471            Python::attach(|py| {
472                let handler = on_error.clone_ref(py);
473                let _ = handler.call(py, (session_id, error), None);
474            });
475        }).await.ok();
476        
477        Ok(())
478    }
479}
480
481#[cfg(feature = "pyo3")]
482#[pyclass]
483pub struct RiWSSessionManagerPy {
484    manager: RiWSSessionManager,
485}
486
487#[cfg(feature = "pyo3")]
488#[pymethods]
489impl RiWSSessionManagerPy {
490    #[new]
491    fn new(max_connections: usize) -> Self {
492        Self {
493            manager: RiWSSessionManager::new(max_connections),
494        }
495    }
496    
497    fn get_session_count(&self) -> usize {
498        tokio::runtime::Handle::try_current()
499            .map(|handle| handle.block_on(async { self.manager.get_session_count().await }))
500            .unwrap_or(0)
501    }
502    
503    fn get_all_sessions(&self) -> Vec<RiWSSessionInfo> {
504        tokio::runtime::Handle::try_current()
505            .map(|handle| handle.block_on(async { self.manager.get_all_sessions().await }))
506            .unwrap_or_default()
507    }
508    
509    fn broadcast(&self, data: Vec<u8>) -> usize {
510        tokio::runtime::Handle::try_current()
511            .map(|handle| handle.block_on(async { self.manager.broadcast(&data).await.unwrap_or(0) }))
512            .unwrap_or(0)
513    }
514}
515
516#[derive(Debug, Clone)]
517#[cfg_attr(feature = "pyo3", pyclass)]
518pub struct RiWSServerStats {
519    pub total_connections: u64,
520    pub active_connections: u64,
521    pub total_messages_sent: u64,
522    pub total_messages_received: u64,
523    pub total_bytes_sent: u64,
524    pub total_bytes_received: u64,
525    pub connection_errors: u64,
526    pub message_errors: u64,
527}
528
529#[cfg(feature = "pyo3")]
530#[pymethods]
531impl RiWSServerStats {
532    #[getter]
533    fn get_total_connections(&self) -> u64 {
534        self.total_connections
535    }
536    
537    #[getter]
538    fn get_active_connections(&self) -> u64 {
539        self.active_connections
540    }
541    
542    #[getter]
543    fn get_total_messages_sent(&self) -> u64 {
544        self.total_messages_sent
545    }
546    
547    #[getter]
548    fn get_total_messages_received(&self) -> u64 {
549        self.total_messages_received
550    }
551    
552    #[getter]
553    fn get_total_bytes_sent(&self) -> u64 {
554        self.total_bytes_sent
555    }
556    
557    #[getter]
558    fn get_total_bytes_received(&self) -> u64 {
559        self.total_bytes_received
560    }
561    
562    #[getter]
563    fn get_connection_errors(&self) -> u64 {
564        self.connection_errors
565    }
566    
567    #[getter]
568    fn get_message_errors(&self) -> u64 {
569        self.message_errors
570    }
571}
572
573impl RiWSServerStats {
574    pub fn new() -> Self {
575        Self {
576            total_connections: 0,
577            active_connections: 0,
578            total_messages_sent: 0,
579            total_messages_received: 0,
580            total_bytes_sent: 0,
581            total_bytes_received: 0,
582            connection_errors: 0,
583            message_errors: 0,
584        }
585    }
586
587    pub fn record_connection(&mut self) {
588        self.total_connections += 1;
589        self.active_connections += 1;
590    }
591
592    pub fn record_disconnection(&mut self) {
593        if self.active_connections > 0 {
594            self.active_connections -= 1;
595        }
596    }
597
598    pub fn record_message_sent(&mut self, size: usize) {
599        self.total_messages_sent += 1;
600        self.total_bytes_sent += size as u64;
601    }
602
603    pub fn record_message_received(&mut self, size: usize) {
604        self.total_messages_received += 1;
605        self.total_bytes_received += size as u64;
606    }
607
608    pub fn record_connection_error(&mut self) {
609        self.connection_errors += 1;
610        if self.active_connections > 0 {
611            self.active_connections -= 1;
612        }
613    }
614
615    pub fn record_message_error(&mut self) {
616        self.message_errors += 1;
617    }
618}
619
620impl Default for RiWSServerStats {
621    fn default() -> Self {
622        Self::new()
623    }
624}
ri/ws/mod.rs

ri/ws/
mod.rs
