clawft_kernel/
mesh_ws.rs

1//! WebSocket transport backend for mesh networking (K6).
2//!
3//! Uses tokio-tungstenite for node-to-node communication.  This enables
4//! browser-based nodes (via wasm-bindgen) and traversal of HTTP proxies
5//! that block raw TCP.
6
7use std::net::SocketAddr;
8
9use async_trait::async_trait;
10use futures_util::{SinkExt, StreamExt};
11use tokio::net::TcpListener;
12use tokio_tungstenite::MaybeTlsStream;
13use tokio_tungstenite::tungstenite::Message;
14
15use crate::mesh::{MAX_MESSAGE_SIZE, MeshError, MeshStream, MeshTransport, TransportListener};
16
17// ── Wire helpers ──────────────────────────────────────────────
18
19/// Encode a payload into a length-prefixed binary WebSocket message.
20fn encode_ws_message(data: &[u8]) -> Message {
21    let len = (data.len() as u32).to_be_bytes();
22    let mut payload = Vec::with_capacity(4 + data.len());
23    payload.extend_from_slice(&len);
24    payload.extend_from_slice(data);
25    Message::Binary(payload)
26}
27
28/// Decode a length-prefixed binary WebSocket message payload.
29fn decode_ws_payload(payload: &[u8]) -> Result<Vec<u8>, MeshError> {
30    if payload.len() < 4 {
31        return Err(MeshError::Io(
32            "binary message too short for length prefix".into(),
33        ));
34    }
35    let len = u32::from_be_bytes([payload[0], payload[1], payload[2], payload[3]]) as usize;
36    if len > MAX_MESSAGE_SIZE {
37        return Err(MeshError::MessageTooLarge {
38            size: len,
39            max: MAX_MESSAGE_SIZE,
40        });
41    }
42    if payload.len() < 4 + len {
43        return Err(MeshError::Io("truncated binary message".into()));
44    }
45    Ok(payload[4..4 + len].to_vec())
46}
47
48// ── WsMeshStream (server-side, plain TcpStream) ──────────────
49
50/// A [`MeshStream`] backed by a server-side WebSocket connection
51/// (accepted from a [`TcpListener`]).
52pub struct WsMeshStream {
53    sink: futures_util::stream::SplitSink<
54        tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>,
55        Message,
56    >,
57    stream: futures_util::stream::SplitStream<
58        tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>,
59    >,
60    remote: SocketAddr,
61}
62
63#[async_trait]
64impl MeshStream for WsMeshStream {
65    async fn send(&mut self, data: &[u8]) -> Result<(), MeshError> {
66        self.sink
67            .send(encode_ws_message(data))
68            .await
69            .map_err(|e| MeshError::Io(e.to_string()))
70    }
71
72    async fn recv(&mut self) -> Result<Vec<u8>, MeshError> {
73        loop {
74            let msg = self
75                .stream
76                .next()
77                .await
78                .ok_or(MeshError::ConnectionClosed)?
79                .map_err(|e| MeshError::Io(e.to_string()))?;
80            match msg {
81                Message::Binary(payload) => return decode_ws_payload(&payload),
82                Message::Close(_) => return Err(MeshError::ConnectionClosed),
83                _ => continue,
84            }
85        }
86    }
87
88    async fn close(&mut self) -> Result<(), MeshError> {
89        self.sink
90            .send(Message::Close(None))
91            .await
92            .map_err(|e| MeshError::Io(e.to_string()))
93    }
94
95    fn remote_addr(&self) -> Option<SocketAddr> {
96        Some(self.remote)
97    }
98}
99
100// ── WsClientStream (client-side, MaybeTlsStream) ─────────────
101
102/// A [`MeshStream`] backed by a client-side WebSocket connection
103/// (created via [`tokio_tungstenite::connect_async`]).
104pub struct WsClientStream {
105    sink: futures_util::stream::SplitSink<
106        tokio_tungstenite::WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>,
107        Message,
108    >,
109    stream: futures_util::stream::SplitStream<
110        tokio_tungstenite::WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>,
111    >,
112    remote: SocketAddr,
113}
114
115#[async_trait]
116impl MeshStream for WsClientStream {
117    async fn send(&mut self, data: &[u8]) -> Result<(), MeshError> {
118        self.sink
119            .send(encode_ws_message(data))
120            .await
121            .map_err(|e| MeshError::Io(e.to_string()))
122    }
123
124    async fn recv(&mut self) -> Result<Vec<u8>, MeshError> {
125        loop {
126            let msg = self
127                .stream
128                .next()
129                .await
130                .ok_or(MeshError::ConnectionClosed)?
131                .map_err(|e| MeshError::Io(e.to_string()))?;
132            match msg {
133                Message::Binary(payload) => return decode_ws_payload(&payload),
134                Message::Close(_) => return Err(MeshError::ConnectionClosed),
135                _ => continue,
136            }
137        }
138    }
139
140    async fn close(&mut self) -> Result<(), MeshError> {
141        self.sink
142            .send(Message::Close(None))
143            .await
144            .map_err(|e| MeshError::Io(e.to_string()))
145    }
146
147    fn remote_addr(&self) -> Option<SocketAddr> {
148        Some(self.remote)
149    }
150}
151
152// ── WsTransportListener ──────────────────────────────────────
153
154/// A [`TransportListener`] that accepts TCP connections and upgrades
155/// them to WebSocket.
156pub struct WsTransportListener {
157    listener: TcpListener,
158}
159
160#[async_trait]
161impl TransportListener for WsTransportListener {
162    async fn accept(&mut self) -> Result<(Box<dyn MeshStream>, SocketAddr), MeshError> {
163        let (tcp_stream, addr) = self
164            .listener
165            .accept()
166            .await
167            .map_err(|e| MeshError::Io(e.to_string()))?;
168
169        let ws_stream = tokio_tungstenite::accept_async(tcp_stream)
170            .await
171            .map_err(|e| MeshError::Handshake(e.to_string()))?;
172
173        let (sink, stream) = ws_stream.split();
174        Ok((
175            Box::new(WsMeshStream {
176                sink,
177                stream,
178                remote: addr,
179            }),
180            addr,
181        ))
182    }
183
184    fn local_addr(&self) -> Result<SocketAddr, MeshError> {
185        self.listener
186            .local_addr()
187            .map_err(|e| MeshError::Io(e.to_string()))
188    }
189}
190
191// ── WsTransport ──────────────────────────────────────────────
192
193/// WebSocket mesh transport.
194///
195/// Implements [`MeshTransport`] using tokio-tungstenite.  Addresses
196/// use the `ws://` or `wss://` scheme.
197pub struct WsTransport;
198
199#[async_trait]
200impl MeshTransport for WsTransport {
201    fn name(&self) -> &str {
202        "websocket"
203    }
204
205    async fn listen(&self, addr: &str) -> Result<Box<dyn TransportListener>, MeshError> {
206        let bind_addr = addr
207            .strip_prefix("ws://")
208            .or_else(|| addr.strip_prefix("wss://"))
209            .unwrap_or(addr);
210        let listener = TcpListener::bind(bind_addr)
211            .await
212            .map_err(|e| MeshError::Io(e.to_string()))?;
213        Ok(Box::new(WsTransportListener { listener }))
214    }
215
216    async fn connect(&self, addr: &str) -> Result<Box<dyn MeshStream>, MeshError> {
217        // Ensure the address has a ws:// scheme for tungstenite.
218        let url = if addr.starts_with("ws://") || addr.starts_with("wss://") {
219            addr.to_string()
220        } else {
221            format!("ws://{addr}")
222        };
223
224        let (ws_stream, _response) = tokio_tungstenite::connect_async(&url)
225            .await
226            .map_err(|e| MeshError::Io(e.to_string()))?;
227
228        // Extract remote address from the URL for bookkeeping.
229        let remote: SocketAddr = url
230            .strip_prefix("ws://")
231            .or_else(|| url.strip_prefix("wss://"))
232            .unwrap_or(&url)
233            .parse()
234            .map_err(|e: std::net::AddrParseError| MeshError::Io(e.to_string()))?;
235
236        let (sink, stream) = ws_stream.split();
237        Ok(Box::new(WsClientStream {
238            sink,
239            stream,
240            remote,
241        }))
242    }
243
244    fn supports(&self, addr: &str) -> bool {
245        addr.starts_with("ws://") || addr.starts_with("wss://")
246    }
247}
248
249// ── Tests ────────────────────────────────────────────────────
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254
255    #[tokio::test]
256    async fn ws_transport_connect_send_recv() {
257        let transport = WsTransport;
258
259        // Bind to an OS-assigned port.
260        let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
261        let addr = listener.local_addr().unwrap();
262
263        // "Client" node sends, then reads response.
264        let ws_url = format!("ws://{addr}");
265        let send_task = tokio::spawn(async move {
266            let mut stream = WsTransport.connect(&ws_url).await.unwrap();
267            stream.send(b"hello from node A").await.unwrap();
268            let response = stream.recv().await.unwrap();
269            assert_eq!(response, b"hello from node B");
270            stream.close().await.unwrap();
271        });
272
273        // "Server" node accepts, reads, responds.
274        let (mut server_stream, _client_addr) = listener.accept().await.unwrap();
275        let received = server_stream.recv().await.unwrap();
276        assert_eq!(received, b"hello from node A");
277        server_stream.send(b"hello from node B").await.unwrap();
278
279        send_task.await.unwrap();
280    }
281
282    #[test]
283    fn ws_transport_supports() {
284        let t = WsTransport;
285        assert!(t.supports("ws://127.0.0.1:9489"));
286        assert!(t.supports("wss://example.com:443"));
287        assert!(!t.supports("tcp://127.0.0.1:9489"));
288        assert!(!t.supports("127.0.0.1:9489"));
289    }
290
291    #[tokio::test]
292    async fn browser_node_connects_via_websocket() {
293        // Simulates a browser node joining the mesh over WebSocket.
294        let transport = WsTransport;
295        let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
296        let addr = listener.local_addr().unwrap();
297
298        let ws_url = format!("ws://{addr}");
299        let browser_task = tokio::spawn(async move {
300            let mut stream = WsTransport.connect(&ws_url).await.unwrap();
301
302            // Browser sends a join request payload.
303            let join_payload = br#"{"node_id":"browser-1","type":"join"}"#;
304            stream.send(join_payload).await.unwrap();
305
306            // Receive acknowledgement.
307            let ack = stream.recv().await.unwrap();
308            assert_eq!(ack, b"welcome");
309
310            stream.close().await.unwrap();
311        });
312
313        // Server side accepts the browser node.
314        let (mut peer, _) = listener.accept().await.unwrap();
315        let data = peer.recv().await.unwrap();
316        assert!(data.starts_with(b"{"));
317        let parsed: serde_json::Value = serde_json::from_slice(&data).unwrap();
318        assert_eq!(parsed["node_id"], "browser-1");
319
320        peer.send(b"welcome").await.unwrap();
321
322        browser_task.await.unwrap();
323    }
324}
clawft_kernel/mesh_ws.rs

clawft_kernel/
mesh_ws.rs
