ferric_tensor/
ws.rs

1//! A tiny dependency-free WebSocket bridge so the scheduler can reach a **browser tab** as a device.
2//! Browsers can't accept raw TCP, but a tab connects *out* via WebSocket — so we run a small HTTP
3//! server that serves the worker page + wasm and upgrades `/ws` to a WebSocket. Op frames (the same
4//! host-buffer format as Device::Remote) ride the socket to the tab, which computes on WebGPU and
5//! sends the result back. RFC 6455 handshake (SHA-1 + base64) and frame codec implemented inline.
6
7use std::io::{Read, Write};
8use std::net::{TcpListener, TcpStream};
9use std::sync::{Arc, Mutex};
10
11// ---- SHA-1 (for the WebSocket accept key) ----
12fn sha1(data: &[u8]) -> [u8; 20] {
13    let mut h: [u32; 5] = [0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0];
14    let ml = (data.len() as u64) * 8;
15    let mut msg = data.to_vec();
16    msg.push(0x80);
17    while msg.len() % 64 != 56 { msg.push(0); }
18    msg.extend_from_slice(&ml.to_be_bytes());
19    for chunk in msg.chunks(64) {
20        let mut w = [0u32; 80];
21        for i in 0..16 { w[i] = u32::from_be_bytes([chunk[i * 4], chunk[i * 4 + 1], chunk[i * 4 + 2], chunk[i * 4 + 3]]); }
22        for i in 16..80 { w[i] = (w[i - 3] ^ w[i - 8] ^ w[i - 14] ^ w[i - 16]).rotate_left(1); }
23        let (mut a, mut b, mut c, mut d, mut e) = (h[0], h[1], h[2], h[3], h[4]);
24        for (i, &wi) in w.iter().enumerate() {
25            let (f, k) = match i {
26                0..=19 => ((b & c) | ((!b) & d), 0x5A827999u32),
27                20..=39 => (b ^ c ^ d, 0x6ED9EBA1),
28                40..=59 => ((b & c) | (b & d) | (c & d), 0x8F1BBCDC),
29                _ => (b ^ c ^ d, 0xCA62C1D6),
30            };
31            let t = a.rotate_left(5).wrapping_add(f).wrapping_add(e).wrapping_add(k).wrapping_add(wi);
32            e = d; d = c; c = b.rotate_left(30); b = a; a = t;
33        }
34        h[0] = h[0].wrapping_add(a); h[1] = h[1].wrapping_add(b); h[2] = h[2].wrapping_add(c);
35        h[3] = h[3].wrapping_add(d); h[4] = h[4].wrapping_add(e);
36    }
37    let mut out = [0u8; 20];
38    for i in 0..5 { out[i * 4..i * 4 + 4].copy_from_slice(&h[i].to_be_bytes()); }
39    out
40}
41fn base64(data: &[u8]) -> String {
42    const T: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
43    let mut s = String::new();
44    for c in data.chunks(3) {
45        let b = [c[0], *c.get(1).unwrap_or(&0), *c.get(2).unwrap_or(&0)];
46        let n = (b[0] as u32) << 16 | (b[1] as u32) << 8 | b[2] as u32;
47        s.push(T[(n >> 18 & 63) as usize] as char);
48        s.push(T[(n >> 12 & 63) as usize] as char);
49        s.push(if c.len() > 1 { T[(n >> 6 & 63) as usize] as char } else { '=' });
50        s.push(if c.len() > 2 { T[(n & 63) as usize] as char } else { '=' });
51    }
52    s
53}
54
55/// One WebSocket connection to a browser tab. Binary messages only.
56pub struct WsConn(Mutex<TcpStream>);
57impl WsConn {
58    pub fn send(&self, payload: &[u8]) -> std::io::Result<()> {
59        let mut s = self.0.lock().unwrap();
60        let mut f = vec![0x82u8]; // FIN + binary
61        let n = payload.len();
62        if n < 126 { f.push(n as u8); }
63        else if n < 65536 { f.push(126); f.extend_from_slice(&(n as u16).to_be_bytes()); }
64        else { f.push(127); f.extend_from_slice(&(n as u64).to_be_bytes()); }
65        f.extend_from_slice(payload);
66        s.write_all(&f)
67    }
68    /// Blocking receive of one binary message (handles pings, ignores control frames).
69    pub fn recv(&self) -> std::io::Result<Vec<u8>> {
70        let mut s = self.0.lock().unwrap();
71        loop {
72            let mut hdr = [0u8; 2];
73            s.read_exact(&mut hdr)?;
74            let opcode = hdr[0] & 0x0f;
75            let masked = hdr[1] & 0x80 != 0;
76            let mut len = (hdr[1] & 0x7f) as usize;
77            if len == 126 { let mut b = [0u8; 2]; s.read_exact(&mut b)?; len = u16::from_be_bytes(b) as usize; }
78            else if len == 127 { let mut b = [0u8; 8]; s.read_exact(&mut b)?; len = u64::from_be_bytes(b) as usize; }
79            let mut mask = [0u8; 4];
80            if masked { s.read_exact(&mut mask)?; }
81            let mut payload = vec![0u8; len];
82            s.read_exact(&mut payload)?;
83            if masked { for i in 0..len { payload[i] ^= mask[i % 4]; } }
84            match opcode {
85                0x2 | 0x1 => return Ok(payload),        // binary / text
86                0x8 => return Err(std::io::Error::new(std::io::ErrorKind::ConnectionAborted, "close")),
87                _ => continue,                          // ping/pong/continuation → ignore
88            }
89        }
90    }
91}
92
93fn mime(path: &str) -> &'static str {
94    if path.ends_with(".wasm") { "application/wasm" }
95    else if path.ends_with(".js") { "text/javascript" }
96    else if path.ends_with(".html") { "text/html" }
97    else { "application/octet-stream" }
98}
99
100/// Serve `site_dir` over HTTP and hand back the first browser tab that connects to `/ws`.
101pub struct Bridge {
102    pub addr: String,
103    worker: Arc<Mutex<Option<Arc<WsConn>>>>,
104}
105impl Bridge {
106    pub fn url(&self) -> String { format!("http://{}/", self.addr) }
107    pub fn take_worker(&self) -> Option<Arc<WsConn>> { self.worker.lock().unwrap().clone() }
108}
109
110/// Start the bridge on an ephemeral localhost port, serving files from `site_dir`.
111pub fn start(site_dir: String) -> Bridge {
112    let listener = TcpListener::bind("127.0.0.1:0").expect("bind bridge");
113    let addr = listener.local_addr().unwrap().to_string();
114    let worker: Arc<Mutex<Option<Arc<WsConn>>>> = Arc::new(Mutex::new(None));
115    let w2 = worker.clone();
116    std::thread::spawn(move || {
117        for stream in listener.incoming().flatten() {
118            let _ = handle_http(stream, &site_dir, &w2);
119        }
120    });
121    Bridge { addr, worker }
122}
123
124fn handle_http(mut s: TcpStream, site_dir: &str, worker: &Arc<Mutex<Option<Arc<WsConn>>>>) -> std::io::Result<()> {
125    // read request headers
126    let mut buf = Vec::new();
127    let mut byte = [0u8; 1];
128    while !buf.ends_with(b"\r\n\r\n") {
129        if s.read(&mut byte)? == 0 { return Ok(()); }
130        buf.push(byte[0]);
131    }
132    let req = String::from_utf8_lossy(&buf);
133    let path = req.lines().next().and_then(|l| l.split_whitespace().nth(1)).unwrap_or("/");
134    if req.to_lowercase().contains("upgrade: websocket") {
135        let key = req.lines().find(|l| l.to_lowercase().starts_with("sec-websocket-key:"))
136            .map(|l| l.split(':').nth(1).unwrap().trim().to_string()).unwrap_or_default();
137        let accept = base64(&sha1(format!("{key}258EAFA5-E914-47DA-95CA-C5AB0DC85B11").as_bytes()));
138        let resp = format!("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n");
139        s.write_all(resp.as_bytes())?;
140        *worker.lock().unwrap() = Some(Arc::new(WsConn(Mutex::new(s))));
141        return Ok(());
142    }
143    let rel = if path == "/" { "worker.html".to_string() } else { path.trim_start_matches('/').to_string() };
144    match std::fs::read(format!("{site_dir}/{rel}")) {
145        Ok(body) => {
146            let hdr = format!("HTTP/1.1 200 OK\r\nContent-Type: {}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", mime(&rel), body.len());
147            s.write_all(hdr.as_bytes())?; s.write_all(&body)?;
148        }
149        Err(_) => { s.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n")?; }
150    }
151    Ok(())
152}
ferric_tensor/ws.rs

ferric_tensor/
ws.rs
