Skip to main content

dmx_interface/net/
ws.rs

1//! Minimal WebSocket server (RFC 6455 subset) with JSON envelopes.
2//!
3//! Reference infrastructure for the web contributor. Speaks unfragmented
4//! text frames only (client frames masked, server frames unmasked) and
5//! maps JSON messages onto the config surface:
6//!
7//! ```json
8//! {"type": "config.get", "id": "1"}
9//! {"type": "config.set", "id": "2", "config": {"led_brightness": 128}}
10//! {"type": "system.get", "id": "3"}
11//! {"type": "ping", "id": "4"}
12//! {"type": "reset", "id": "5"}
13//! ```
14//!
15//! Replies: `{"type": "ok", "id": …, "config"?| "system"? …}`,
16//! `{"type": "pong", "id": …}`, `{"type": "err", "id": …, "code", "message"}`.
17//! `config.set` takes a [`ConfigPatch`] — same merge semantics as
18//! `POST /api/config`. Fragmented frames, continuations and binary
19//! messages end the session without a close frame (extending that to a
20//! close code like 1003 is the contributor's call).
21
22use core::fmt::Write as _;
23
24use base64::Engine as _;
25use base64::engine::general_purpose::STANDARD as BASE64;
26use embassy_net::tcp::TcpSocket;
27use embassy_time::{Duration, Instant, Timer, with_timeout};
28use embedded_io_async::{Read as _, Write as _};
29use heapless::String;
30use serde::{Deserialize, Serialize};
31use sha1::{Digest, Sha1};
32
33use crate::config::{Config, ConfigPatch};
34use crate::storage;
35
36/// RFC 6455 magic GUID for the accept key.
37const GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
38
39/// Largest client payload accepted (config patches are well below this).
40const MAX_PAYLOAD: usize = 768;
41
42/// Largest server payload emitted (config JSON answers).
43const MAX_REPLY: usize = 1024;
44
45/// Per-frame idle budget — dead peers free their slot.
46const IDLE_TIMEOUT: Duration = Duration::from_secs(120);
47
48/// Frame opcode: text payload.
49const OPCODE_TEXT: u8 = 0x1;
50/// Frame opcode: close the connection.
51const OPCODE_CLOSE: u8 = 0x8;
52/// Frame opcode: ping probe.
53const OPCODE_PING: u8 = 0x9;
54/// Frame opcode: pong answer.
55const OPCODE_PONG: u8 = 0xA;
56
57/// Run one WebSocket session until close; the socket stays borrowed in.
58pub(super) async fn serve(socket: &mut TcpSocket<'_>, key: &str) -> Result<(), ()> {
59    handshake(socket, key).await?;
60    log::info!("ws: session open");
61
62    let mut payload = [0u8; MAX_PAYLOAD];
63    loop {
64        let frame = with_timeout(IDLE_TIMEOUT, read_frame(socket, &mut payload)).await;
65        let (opcode, len) = if let Ok(result) = frame {
66            result?
67        } else {
68            log::debug!("ws: idle timeout");
69            return Ok(());
70        };
71        let data = payload.get(..len).ok_or(())?;
72        match opcode {
73            OPCODE_TEXT => handle_message(socket, data).await?,
74            OPCODE_PING => write_frame(socket, OPCODE_PONG, data).await?,
75            OPCODE_PONG => {}
76            OPCODE_CLOSE => {
77                write_frame(socket, OPCODE_CLOSE, data).await?;
78                log::info!("ws: session closed by client");
79                return Ok(());
80            }
81            _ => {
82                log::warn!("ws: unsupported opcode {opcode:#x}, closing");
83                return Err(());
84            }
85        }
86    }
87}
88
89/// Complete the HTTP→WebSocket upgrade handshake.
90async fn handshake(socket: &mut TcpSocket<'_>, key: &str) -> Result<(), ()> {
91    let mut hasher = Sha1::new();
92    hasher.update(key.as_bytes());
93    hasher.update(GUID.as_bytes());
94    let digest = hasher.finalize();
95
96    let mut raw = [0u8; 28];
97    let n = BASE64.encode_slice(digest, &mut raw).map_err(|_| ())?;
98    let bytes = raw.get(..n).ok_or(())?;
99    let Some(accept) = core::str::from_utf8(bytes)
100        .ok()
101        .and_then(|text| text.parse::<String<32>>().ok())
102    else {
103        return Err(());
104    };
105
106    let mut response = String::<160>::new();
107    write!(
108        &mut response,
109        "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n"
110    )
111    .map_err(|_| ())?;
112    socket.write_all(response.as_bytes()).await.map_err(|_| ())
113}
114
115/// Read one complete client frame (masked, unfragmented) into `buf`.
116async fn read_frame(socket: &mut TcpSocket<'_>, buf: &mut [u8]) -> Result<(u8, usize), ()> {
117    let mut head = [0u8; 2];
118    socket.read_exact(&mut head).await.map_err(|_| ())?;
119    let first = *head.first().ok_or(())?;
120    let second = *head.get(1).ok_or(())?;
121
122    if first & 0x70 != 0 || first & 0x80 == 0 {
123        // RSV bits set or fragmented — not supported, reject.
124        return Err(());
125    }
126    if second & 0x80 == 0 {
127        // Clients must mask (RFC 6455 §5.1).
128        return Err(());
129    }
130
131    let opcode = first & 0x0f;
132    let len = match second & 0x7f {
133        0..=125 => usize::from(second & 0x7f),
134        126 => {
135            let mut ext = [0u8; 2];
136            socket.read_exact(&mut ext).await.map_err(|_| ())?;
137            usize::from(u16::from_be_bytes(ext))
138        }
139        _ => return Err(()), // 64-bit lengths unsupported
140    };
141    if len > buf.len() {
142        return Err(());
143    }
144
145    let mut mask = [0u8; 4];
146    socket.read_exact(&mut mask).await.map_err(|_| ())?;
147    let target = buf.get_mut(..len).ok_or(())?;
148    socket.read_exact(target).await.map_err(|_| ())?;
149    for (byte, mask_byte) in target.iter_mut().zip(mask.iter().cycle()) {
150        *byte ^= *mask_byte;
151    }
152    Ok((opcode, len))
153}
154
155/// Write one unmasked server frame.
156async fn write_frame(socket: &mut TcpSocket<'_>, opcode: u8, payload: &[u8]) -> Result<(), ()> {
157    let mut head = [0u8; 4];
158    let first = head.first_mut().ok_or(())?;
159    *first = 0x80 | opcode;
160    let len = payload.len();
161    if len <= 125 {
162        let second = head.get_mut(1).ok_or(())?;
163        *second = u8::try_from(len).map_err(|_| ())?;
164        socket
165            .write_all(head.get(..2).ok_or(())?)
166            .await
167            .map_err(|_| ())?;
168    } else {
169        let second = head.get_mut(1).ok_or(())?;
170        *second = 126;
171        let be = u16::try_from(len).map_err(|_| ())?.to_be_bytes();
172        head.get_mut(2..4).ok_or(())?.copy_from_slice(&be);
173        socket.write_all(&head).await.map_err(|_| ())?;
174    }
175    socket.write_all(payload).await.map_err(|_| ())
176}
177
178/// First-pass envelope: `{"type": "...", "id": "..."}` (extra fields ignored).
179///
180/// Internally tagged enums need serde's alloc buffering, which `no_std`
181/// `serde-json-core` lacks — so requests are parsed twice from the same
182/// slice: the head first, then a typed payload for `config.set`.
183#[derive(Deserialize)]
184struct MessageHead {
185    /// Message discriminator.
186    #[serde(rename = "type")]
187    kind: String<16>,
188    /// Correlation id echoed in the reply.
189    #[serde(default)]
190    id: String<16>,
191}
192
193/// Payload of a `config.set` message.
194#[derive(Deserialize)]
195struct ConfigSetPayload {
196    /// Partial config (web: `DeepPartial<Config>`).
197    config: ConfigPatch,
198}
199
200/// Server reply envelope — flat struct (same JSON as a tagged enum).
201#[derive(Serialize)]
202struct Reply {
203    /// Message discriminator: `ok`, `err`, `pong`.
204    #[serde(rename = "type")]
205    kind: &'static str,
206    /// Correlation id from the request.
207    id: String<16>,
208    /// Full config (`config.get`).
209    #[serde(skip_serializing_if = "Option::is_none")]
210    config: Option<Config>,
211    /// Runtime info (`system.get`).
212    #[serde(skip_serializing_if = "Option::is_none")]
213    system: Option<SystemInfo>,
214    /// Machine-readable error code (`err`).
215    #[serde(skip_serializing_if = "Option::is_none")]
216    code: Option<String<24>>,
217    /// Human-readable error detail (`err`).
218    #[serde(skip_serializing_if = "Option::is_none")]
219    message: Option<String<64>>,
220}
221
222/// Runtime information payload (`system.get`).
223#[derive(Serialize)]
224struct SystemInfo {
225    /// Milliseconds since boot.
226    uptime_ms: u64,
227    /// Free internal heap bytes.
228    heap_free: usize,
229    /// Used internal heap bytes.
230    heap_used: usize,
231    /// Station/AP MAC address, colon-separated uppercase hex.
232    mac: String<18>,
233}
234
235impl Reply {
236    /// `ok` without payload.
237    fn ok(id: String<16>) -> Self {
238        Self {
239            kind: "ok",
240            id,
241            config: None,
242            system: None,
243            code: None,
244            message: None,
245        }
246    }
247
248    /// `ok` carrying the full config.
249    fn ok_config(id: String<16>, config: Config) -> Self {
250        Self {
251            kind: "ok",
252            id,
253            config: Some(config),
254            system: None,
255            code: None,
256            message: None,
257        }
258    }
259
260    /// `ok` carrying runtime info.
261    fn ok_system(id: String<16>, system: SystemInfo) -> Self {
262        Self {
263            kind: "ok",
264            id,
265            config: None,
266            system: Some(system),
267            code: None,
268            message: None,
269        }
270    }
271
272    /// `pong` liveness reply.
273    fn pong(id: String<16>) -> Self {
274        Self {
275            kind: "pong",
276            id,
277            config: None,
278            system: None,
279            code: None,
280            message: None,
281        }
282    }
283
284    /// `err` with code and message.
285    fn err(id: String<16>, code: &str, message: &str) -> Self {
286        let mut code_buf = String::<24>::new();
287        let _ = code_buf.push_str(code);
288        let mut message_buf = String::<64>::new();
289        let _ = message_buf.push_str(message);
290        Self {
291            kind: "err",
292            id,
293            config: None,
294            system: None,
295            code: Some(code_buf),
296            message: Some(message_buf),
297        }
298    }
299}
300
301/// Parse and execute one JSON message, then send the reply.
302async fn handle_message(socket: &mut TcpSocket<'_>, data: &[u8]) -> Result<(), ()> {
303    let (head, _used): (MessageHead, usize) = match serde_json_core::from_slice(data) {
304        Ok(parsed) => parsed,
305        Err(err) => {
306            log::warn!("ws: bad message: {err}");
307            return send_reply(
308                socket,
309                &Reply::err(String::new(), "bad_request", "unparseable"),
310            )
311            .await;
312        }
313    };
314
315    let reply = match head.kind.as_str() {
316        "config.get" => Reply::ok_config(head.id, storage::load()),
317        "config.set" => match serde_json_core::from_slice::<ConfigSetPayload>(data) {
318            Ok((payload, _used)) => {
319                let mut config = storage::load();
320                if !config.merge_patch(payload.config) {
321                    log::warn!("ws: config patch rejected");
322                    Reply::err(head.id, "bad_request", "invalid patch")
323                } else if storage::save(&config) {
324                    config.apply();
325                    log::info!("ws: config saved");
326                    Reply::ok(head.id)
327                } else {
328                    log::warn!("ws: config save failed");
329                    Reply::err(head.id, "save_failed", "nvs write failed")
330                }
331            }
332            Err(err) => {
333                log::warn!("ws: bad config.set payload: {err}");
334                Reply::err(head.id, "bad_request", "bad config payload")
335            }
336        },
337        "system.get" => Reply::ok_system(
338            head.id,
339            SystemInfo {
340                uptime_ms: Instant::now().as_millis(),
341                heap_free: esp_alloc::HEAP.free(),
342                heap_used: esp_alloc::HEAP.used(),
343                mac: format_mac(&crate::hardware::efuse::mac_bytes()),
344            },
345        ),
346        "reset" => {
347            log::warn!("ws: factory reset requested");
348            if storage::clear() {
349                let reply = Reply::ok(head.id);
350                if send_reply(socket, &reply).await.is_err() {
351                    log::warn!("ws: reset reply undeliverable — rebooting anyway");
352                }
353                Timer::after(Duration::from_millis(100)).await;
354                esp_hal::system::software_reset()
355            } else {
356                log::error!("ws: factory reset failed to clear NVS — not rebooting");
357                Reply::err(head.id, "reset_failed", "nvs clear failed")
358            }
359        }
360        "ping" => Reply::pong(head.id),
361        other => {
362            log::warn!("ws: unknown message type '{other}'");
363            Reply::err(head.id, "bad_request", "unknown type")
364        }
365    };
366
367    send_reply(socket, &reply).await
368}
369
370/// Serialize and write one reply frame.
371async fn send_reply(socket: &mut TcpSocket<'_>, reply: &Reply) -> Result<(), ()> {
372    let mut body = [0u8; MAX_REPLY];
373    let n = serde_json_core::to_slice(reply, &mut body).map_err(|_| ())?;
374    write_frame(socket, OPCODE_TEXT, body.get(..n).ok_or(())?).await
375}
376
377/// Format MAC bytes as `AA:BB:CC:DD:EE:FF`.
378fn format_mac(mac: &[u8; 6]) -> String<18> {
379    let mut out = String::new();
380    for (index, byte) in mac.iter().enumerate() {
381        if index > 0 {
382            let _ = out.push(':');
383        }
384        let _ = write!(&mut out, "{byte:02X}");
385    }
386    out
387}