1use 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
36const GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
38
39const MAX_PAYLOAD: usize = 768;
41
42const MAX_REPLY: usize = 1024;
44
45const IDLE_TIMEOUT: Duration = Duration::from_secs(120);
47
48const OPCODE_TEXT: u8 = 0x1;
50const OPCODE_CLOSE: u8 = 0x8;
52const OPCODE_PING: u8 = 0x9;
54const OPCODE_PONG: u8 = 0xA;
56
57pub(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
89async 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
115async 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 return Err(());
125 }
126 if second & 0x80 == 0 {
127 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(()), };
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
155async 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#[derive(Deserialize)]
184struct MessageHead {
185 #[serde(rename = "type")]
187 kind: String<16>,
188 #[serde(default)]
190 id: String<16>,
191}
192
193#[derive(Deserialize)]
195struct ConfigSetPayload {
196 config: ConfigPatch,
198}
199
200#[derive(Serialize)]
202struct Reply {
203 #[serde(rename = "type")]
205 kind: &'static str,
206 id: String<16>,
208 #[serde(skip_serializing_if = "Option::is_none")]
210 config: Option<Config>,
211 #[serde(skip_serializing_if = "Option::is_none")]
213 system: Option<SystemInfo>,
214 #[serde(skip_serializing_if = "Option::is_none")]
216 code: Option<String<24>>,
217 #[serde(skip_serializing_if = "Option::is_none")]
219 message: Option<String<64>>,
220}
221
222#[derive(Serialize)]
224struct SystemInfo {
225 uptime_ms: u64,
227 heap_free: usize,
229 heap_used: usize,
231 mac: String<18>,
233}
234
235impl Reply {
236 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 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 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 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 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
301async 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
370async 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
377fn 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}