Skip to main content

grid_rs/
ws.rs

1//! Real-time WebSocket.
2//!
3//! This module provides WebSocket functionality for both client and server connections.
4//! It currently supports TCP with plans to support TLS connections (including support for RA certs), with automatic resource cleanup.
5//!
6//! # Examples
7//!
8//! ## Server
9//! ```rust
10//! use grid_rs::ws::WebSocketServer;
11//!
12//! let ws = WebSocketServer::create("127.0.0.1:9002")?;
13//! if let Some(conn) = ws.accept()? {
14//!     if let Some(msg) = conn.receive()? {
15//!         conn.send(conn.id, &msg)?; // Echo back
16//!     }
17//! }
18//! ```
19//!
20//! ## Client
21//! ```rust
22//! use grid_rs::ws::WebSocketClient;
23//!
24//! let ws = WebSocket::connect("ws://localhost:9002")?;
25//! ws.send(b"Hello WebSocket!")?;
26//! if let Some(response) = ws.receive()? {
27//!     return Ok(response.to_vec());
28//! }
29//! ```
30//! See chat room [example](https://github.com/s3ndotxyz/grid-rs/tree/main/examples/websocket) for
31//! more complex usage.
32//!
33
34use crate::region::Region;
35use std::io::Result;
36
37unsafe extern "C" {
38    fn ws_server_create(addr_ptr: usize) -> usize;
39    fn ws_server_accept(server_id: usize) -> usize;
40    fn ws_server_send(server_id: usize, conn_id: usize, data_ptr: usize);
41    fn ws_server_receive(server_id: usize, conn_id: usize) -> usize;
42    fn ws_server_close(server_id: usize);
43
44    fn ws_client_connect(url_ptr: usize) -> usize;
45    fn ws_client_send(client_id: usize, data_ptr: usize);
46    fn ws_client_receive(client_id: usize) -> usize;
47    fn ws_client_close(client_id: usize);
48}
49
50#[derive(Debug)]
51pub struct WebSocketServer {
52    pub id: u32,
53}
54
55#[derive(Debug)]
56pub struct WebSocketClient {
57    pub id: u32,
58}
59
60#[derive(Debug)]
61pub struct WebSocketConnection {
62    pub id: u32,
63    pub server_id: u32,
64}
65
66impl WebSocketServer {
67    pub fn create(addr: &str) -> Result<Self> {
68        let addr = Region::build(addr.as_bytes());
69        let addr_ptr = &*addr as *const Region;
70
71        let id = unsafe { ws_server_create(addr_ptr as usize) } as u32;
72        Ok(Self { id })
73    }
74
75    pub fn accept(&self) -> Result<Option<WebSocketConnection>> {
76        let conn_id = unsafe { ws_server_accept(self.id as usize) } as u32;
77        if conn_id == 0 {
78            return Ok(None);
79        }
80        Ok(Some(WebSocketConnection { id: conn_id, server_id: self.id }))
81    }
82
83    pub fn close(&self) -> Result<()> {
84        unsafe { ws_server_close(self.id as usize) }
85        Ok(())
86    }
87}
88
89impl WebSocketClient {
90    pub fn connect(url: &str) -> Result<Self> {
91        let url = Region::build(url.as_bytes());
92        let url_ptr = &*url as *const Region;
93        let id = unsafe { ws_client_connect(url_ptr as usize) } as u32;
94        Ok(Self { id })
95    }
96
97    pub fn send(&self, data: &[u8]) -> Result<()> {
98        let data = Region::build(data);
99        let data_ptr = &*data as *const Region;
100
101        unsafe { ws_client_send(self.id as usize, data_ptr as usize) }
102        Ok(())
103    }
104
105    pub fn receive(&self) -> Result<Option<Vec<u8>>> {
106        let data_ptr = unsafe { ws_client_receive(self.id as usize) };
107        if data_ptr == 0 {
108            return Ok(None);
109        }
110        unsafe { Ok(Some(Region::consume(data_ptr as *mut Region))) }
111    }
112
113    pub fn close(&self) -> Result<()> {
114        unsafe { ws_client_close(self.id as usize) }
115        Ok(())
116    }
117}
118
119impl WebSocketConnection {
120    pub fn send(&self, client_id: u32, data: &[u8]) -> Result<()> {
121        let data = Region::build(data);
122        let data_ptr = &*data as *const Region;
123        unsafe { ws_server_send(self.server_id as usize, client_id as usize, data_ptr as usize) }
124        Ok(())
125    }
126
127    pub fn receive(&self) -> Result<Option<Vec<u8>>> {
128        let data_ptr = unsafe { ws_server_receive(self.server_id as usize, self.id as usize) };
129        if data_ptr == 0 {
130            return Ok(None);
131        } else if data_ptr == 1 {
132            return Err(std::io::Error::other("Connection closed"));
133        }
134        unsafe { Ok(Some(Region::consume(data_ptr as *mut Region))) }
135    }
136}