simoom/
ws.rs

1//! WebSocket connections after [`crate::Reply::websocket`]: RFC 6455 over HTTP/1.1, and RFC 8441
2//! (extended CONNECT) over HTTP/2. The handler runs on the loop thread; [`WsHandle`] reaches it
3//! from other threads.
4//!
5//! [`WsCore`] is the protocol-independent part (frames, messages, the closing handshake); the
6//! HTTP/1.1 connection runs it as a [`Layer`] ([`WsLayer`]), the HTTP/2 connection on one stream.
7
8use crate::conn::{Io, Layer};
9use crate::handle::{Envelope, Msg, Route};
10use crate::server::Shared;
11use simoom_http1::ws::{self as codec, FrameDecoder, Opcode, Role, WsEvent, WsLimits};
12use std::panic::{catch_unwind, AssertUnwindSafe};
13
14/// A complete message.
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub enum WsMessage<'a> {
17    /// A text message (valid UTF-8).
18    Text(&'a str),
19    /// A binary message.
20    Binary(&'a [u8]),
21}
22
23/// Handles one WebSocket connection, on its loop thread. Must not block.
24pub trait WebSocketHandler: 'static {
25    /// The upgrade completed.
26    fn on_open(&mut self, ws: &mut WebSocket<'_>) {
27        let _ = ws;
28    }
29    /// A whole message arrived (fragments are joined; pings are answered automatically).
30    fn on_message(&mut self, ws: &mut WebSocket<'_>, msg: WsMessage<'_>);
31    /// The connection is closing: the peer's close code (1005 when it sent none, 1006 when the
32    /// connection was lost).
33    fn on_close(&mut self, code: u16, reason: &str) {
34        let _ = (code, reason);
35    }
36}
37
38/// Commands from a [`WsHandle`].
39pub(crate) enum WsCmd {
40    Text(String),
41    Binary(Vec<u8>),
42    Ping(Vec<u8>),
43    Close(u16, String),
44}
45
46/// The connection, inside a [`WebSocketHandler`] callback.
47pub struct WebSocket<'w> {
48    out: &'w mut Vec<u8>,
49    closing: &'w mut bool,
50    route: &'w Route,
51    conn: u64,
52    token: u64,
53    deflate: &'w mut Deflate,
54}
55
56/// The permessage-deflate codec of a WebSocket, when negotiated (feature `ws-deflate`).
57#[cfg(feature = "ws-deflate")]
58pub(crate) type Deflate = Option<simoom_http1::ws::deflate::PerMessage>;
59#[cfg(not(feature = "ws-deflate"))]
60pub(crate) type Deflate = ();
61
62/// Writes one data message, compressed when the extension is on and it is worth it.
63#[inline]
64fn write_message(out: &mut Vec<u8>, deflate: &mut Deflate, op: Opcode, payload: &[u8]) {
65    #[cfg(feature = "ws-deflate")]
66    if let Some(d) = deflate {
67        if d.worth(payload.len()) {
68            let c = d.compress(payload);
69            codec::write_frame_compressed(out, true, op, &c, None);
70            return;
71        }
72    }
73    let _ = deflate;
74    codec::write_frame(out, true, op, payload, None);
75}
76
77/// The `sec-websocket-extensions` answer and the codec for a client's offers (`None, None`
78/// without the feature, the server option, or an acceptable offer).
79#[allow(unused_variables)]
80pub(crate) fn negotiate_deflate(sh: &Shared, offers: Option<&[u8]>) -> (Option<String>, Deflate) {
81    #[cfg(feature = "ws-deflate")]
82    {
83        use simoom_http1::ws::deflate::{negotiate, PerMessage, Side};
84        if let (Some(cfg), Some(o)) = (&sh.ws_deflate, offers) {
85            if let Some((p, resp)) = negotiate(o, cfg) {
86                return (Some(resp), Some(PerMessage::new(Side::Server, &p, cfg)));
87            }
88        }
89        (None, None)
90    }
91    #[cfg(not(feature = "ws-deflate"))]
92    (None, ())
93}
94
95impl std::fmt::Debug for WebSocket<'_> {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        f.debug_struct("WebSocket").field("closing", &*self.closing).finish()
98    }
99}
100
101impl WebSocket<'_> {
102    /// Sends a text message.
103    pub fn send_text(&mut self, text: &str) {
104        if !*self.closing {
105            write_message(self.out, self.deflate, Opcode::Text, text.as_bytes());
106        }
107    }
108
109    /// Sends a binary message.
110    pub fn send_binary(&mut self, data: &[u8]) {
111        if !*self.closing {
112            write_message(self.out, self.deflate, Opcode::Binary, data);
113        }
114    }
115
116    /// Sends a ping (at most 125 bytes).
117    pub fn ping(&mut self, data: &[u8]) {
118        if !*self.closing && data.len() <= 125 {
119            codec::write_frame(self.out, true, Opcode::Ping, data, None);
120        }
121    }
122
123    /// Starts the closing handshake with `code` (for example 1000) and `reason`.
124    pub fn close(&mut self, code: u16, reason: &str) {
125        if !*self.closing {
126            *self.closing = true;
127            if !codec::write_close(self.out, code, reason, None) {
128                codec::write_close(self.out, 1000, "", None);
129            }
130        }
131    }
132
133    /// True once a close was sent.
134    pub fn is_closing(&self) -> bool {
135        *self.closing
136    }
137
138    /// A handle to send from other threads.
139    pub fn handle(&self) -> WsHandle {
140        WsHandle { route: self.route.clone(), conn: self.conn, token: self.token }
141    }
142}
143
144/// Sends on a WebSocket from any thread. Messages to a closed connection are dropped.
145#[derive(Clone, Debug)]
146pub struct WsHandle {
147    route: Route,
148    /// The HTTP/3 connection (0 otherwise).
149    conn: u64,
150    /// 0 on HTTP/1.1; the stream id on HTTP/2 and HTTP/3.
151    token: u64,
152}
153
154impl WsHandle {
155    fn cmd(&self, c: WsCmd) -> bool {
156        self.route.send(self.conn, self.token, Msg::Ws(c))
157    }
158    /// Sends a text message.
159    pub fn send_text(&self, text: impl Into<String>) -> bool {
160        self.cmd(WsCmd::Text(text.into()))
161    }
162    /// Sends a binary message.
163    pub fn send_binary(&self, data: impl Into<Vec<u8>>) -> bool {
164        self.cmd(WsCmd::Binary(data.into()))
165    }
166    /// Sends a ping.
167    pub fn ping(&self, data: impl Into<Vec<u8>>) -> bool {
168        self.cmd(WsCmd::Ping(data.into()))
169    }
170    /// Closes the connection.
171    pub fn close(&self, code: u16, reason: impl Into<String>) -> bool {
172        self.cmd(WsCmd::Close(code, reason.into()))
173    }
174}
175
176/// What a [`WsCore`] call produced: bytes to send, and whether the transport (the HTTP/1.1
177/// connection, or the HTTP/2 stream) should end.
178pub(crate) struct WsIo {
179    pub(crate) out: Vec<u8>,
180    pub(crate) close: bool,
181    /// A handler callback panicked (a close 1011 went out).
182    pub(crate) panicked: bool,
183    pub(crate) now: u64,
184    pub(crate) route: Route,
185    pub(crate) conn: u64,
186    pub(crate) token: u64,
187}
188
189impl WsIo {
190    pub(crate) fn new(route: Route, conn: u64, token: u64, now: u64) -> WsIo {
191        WsIo { out: Vec::new(), close: false, panicked: false, now, route, conn, token }
192    }
193}
194
195/// One WebSocket: the frame decoder, the message being assembled, and the closing handshake.
196pub(crate) struct WsCore {
197    handler: Box<dyn WebSocketHandler>,
198    dec: FrameDecoder,
199    inbuf: Vec<u8>,
200    msg: Vec<u8>,
201    text: bool,
202    max: u64,
203    closing: bool,
204    /// When our close frame went out (0: not yet): the peer has a few seconds to answer.
205    close_sent_at: u64,
206    done: bool,
207    notified: bool,
208    deflate: Deflate,
209    /// The current message is compressed.
210    compressed: bool,
211}
212
213/// How long to wait for the peer's close frame after ours.
214const CLOSE_WAIT_NS: u64 = 5_000_000_000;
215
216impl WsCore {
217    pub(crate) fn new(handler: Box<dyn WebSocketHandler>, max_message: u64, deflate: Deflate) -> WsCore {
218        let limits = WsLimits { max_frame: max_message, max_message };
219        let dec = FrameDecoder::new(Role::Server, limits);
220        #[cfg(feature = "ws-deflate")]
221        let dec = {
222            let mut dec = dec;
223            dec.set_permessage_deflate(deflate.is_some());
224            dec
225        };
226        WsCore { handler, dec, inbuf: Vec::new(), msg: Vec::new(), text: false, max: max_message, closing: false, close_sent_at: 0, done: false, notified: false, deflate, compressed: false }
227    }
228
229    /// Runs a handler callback with a [`WebSocket`] writing into the output.
230    fn with_ws(&mut self, w: &mut WsIo, f: impl FnOnce(&mut dyn WebSocketHandler, &mut WebSocket<'_>)) {
231        let ok = {
232            let mut ws = WebSocket { out: &mut w.out, closing: &mut self.closing, route: &w.route, conn: w.conn, token: w.token, deflate: &mut self.deflate };
233            catch_unwind(AssertUnwindSafe(|| f(&mut *self.handler, &mut ws))).is_ok()
234        };
235        if !ok {
236            w.panicked = true;
237            self.fail(w, 1011);
238        }
239        self.note_close(w);
240    }
241
242    fn note_close(&mut self, w: &WsIo) {
243        if self.closing && self.close_sent_at == 0 {
244            self.close_sent_at = w.now.max(1);
245        }
246    }
247
248    fn fail(&mut self, w: &mut WsIo, code: u16) {
249        if !self.closing {
250            self.closing = true;
251            codec::write_close(&mut w.out, code, "", None);
252        }
253        self.finish(code, "");
254        w.close = true;
255    }
256
257    fn finish(&mut self, code: u16, reason: &str) {
258        self.done = true;
259        if !self.notified {
260            self.notified = true;
261            let h = &mut self.handler;
262            let _ = catch_unwind(AssertUnwindSafe(|| h.on_close(code, reason)));
263        }
264    }
265
266    pub(crate) fn is_done(&self) -> bool {
267        self.done
268    }
269
270    pub(crate) fn start(&mut self, w: &mut WsIo) {
271        self.with_ws(w, |h, ws| h.on_open(ws));
272    }
273
274    pub(crate) fn input(&mut self, w: &mut WsIo, data: &[u8]) {
275        if self.done {
276            return;
277        }
278        self.inbuf.extend_from_slice(data);
279        let mut pos = 0;
280        loop {
281            let mut buf = std::mem::take(&mut self.inbuf);
282            let r = self.dec.decode(&mut buf[pos..]);
283            let step = match r {
284                Err(e) => {
285                    self.inbuf = buf;
286                    self.fail(w, e.close_code());
287                    return;
288                }
289                Ok((n, ev)) => {
290                    pos += n;
291                    match ev {
292                        WsEvent::NeedMore => None,
293                        WsEvent::Data { text, payload, fin } => {
294                            if self.msg.is_empty() {
295                                self.text = text;
296                                self.compressed = self.dec.message_compressed();
297                            }
298                            if (self.msg.len() + payload.len()) as u64 > self.max {
299                                Some(Err(1009))
300                            } else {
301                                self.msg.extend_from_slice(payload);
302                                Some(Ok(fin))
303                            }
304                        }
305                        WsEvent::Ping(p) => {
306                            if !self.closing {
307                                codec::write_frame(&mut w.out, true, Opcode::Pong, p, None);
308                            }
309                            Some(Ok(false))
310                        }
311                        WsEvent::Pong(_) => Some(Ok(false)),
312                        WsEvent::Close { code, reason } => {
313                            let reason = reason.to_string();
314                            if !self.closing {
315                                self.closing = true;
316                                codec::write_close(&mut w.out, code.unwrap_or(1000), "", None);
317                            }
318                            self.inbuf = buf;
319                            self.finish(code.unwrap_or(1005), &reason);
320                            w.close = true;
321                            return;
322                        }
323                    }
324                }
325            };
326            self.inbuf = buf;
327            match step {
328                None => break,
329                Some(Err(code)) => {
330                    self.fail(w, code);
331                    return;
332                }
333                Some(Ok(true)) => {
334                    let msg = std::mem::take(&mut self.msg);
335                    let text = self.text;
336                    #[cfg(feature = "ws-deflate")]
337                    let msg = if self.compressed {
338                        let mut plain = Vec::new();
339                        let r = match &mut self.deflate {
340                            Some(d) => d.decompress(&msg, &mut plain),
341                            None => Err(simoom_http1::ws::WsError::ReservedBits),
342                        };
343                        let r = r.and_then(|_| if text && std::str::from_utf8(&plain).is_err() { Err(simoom_http1::ws::WsError::InvalidUtf8) } else { Ok(()) });
344                        if let Err(e) = r {
345                            self.fail(w, e.close_code());
346                            return;
347                        }
348                        plain
349                    } else {
350                        msg
351                    };
352                    self.with_ws(w, |h, ws| {
353                        let m = if text { WsMessage::Text(std::str::from_utf8(&msg).unwrap_or("")) } else { WsMessage::Binary(&msg) };
354                        h.on_message(ws, m);
355                    });
356                    let mut msg = msg;
357                    msg.clear();
358                    self.msg = msg;
359                    if self.done {
360                        return;
361                    }
362                }
363                Some(Ok(false)) => {}
364            }
365        }
366        self.inbuf.drain(..pos);
367    }
368
369    pub(crate) fn command(&mut self, w: &mut WsIo, cmd: WsCmd) {
370        if self.closing {
371            return;
372        }
373        match cmd {
374            WsCmd::Text(t) => write_message(&mut w.out, &mut self.deflate, Opcode::Text, t.as_bytes()),
375            WsCmd::Binary(b) => write_message(&mut w.out, &mut self.deflate, Opcode::Binary, &b),
376            WsCmd::Ping(p) if p.len() <= 125 => codec::write_frame(&mut w.out, true, Opcode::Ping, &p, None),
377            WsCmd::Ping(_) => {}
378            WsCmd::Close(code, reason) => {
379                self.closing = true;
380                if !codec::write_close(&mut w.out, code, &reason, None) {
381                    codec::write_close(&mut w.out, 1000, "", None);
382                }
383            }
384        }
385        self.note_close(w);
386    }
387
388    /// The transport ended without a closing handshake.
389    pub(crate) fn eof(&mut self, w: &mut WsIo) {
390        self.finish(1006, "");
391        w.close = true;
392    }
393
394    /// The server is stopping: close 1001.
395    pub(crate) fn stop(&mut self, w: &mut WsIo) {
396        if !self.closing {
397            self.closing = true;
398            codec::write_close(&mut w.out, 1001, "", None);
399        }
400        self.finish(1001, "");
401        w.close = true;
402    }
403
404    pub(crate) fn tick(&mut self, w: &mut WsIo) {
405        if self.close_sent_at != 0 && w.now >= self.close_sent_at + CLOSE_WAIT_NS {
406            self.finish(1006, "");
407            w.close = true;
408        }
409    }
410
411    pub(crate) fn deadline(&self) -> u64 {
412        if self.close_sent_at != 0 && !self.done {
413            self.close_sent_at + CLOSE_WAIT_NS
414        } else {
415            u64::MAX
416        }
417    }
418
419    pub(crate) fn closed(&mut self) {
420        self.finish(1006, "");
421    }
422}
423
424impl Drop for WsCore {
425    fn drop(&mut self) {
426        // however the stream or connection went away, the handler hears about it once
427        self.finish(1006, "");
428    }
429}
430
431/// A WebSocket on an HTTP/1.1 connection, after the `101`.
432pub(crate) struct WsLayer {
433    core: WsCore,
434}
435
436impl WsLayer {
437    pub(crate) fn new(handler: Box<dyn WebSocketHandler>, max_message: u64, deflate: Deflate) -> WsLayer {
438        WsLayer { core: WsCore::new(handler, max_message, deflate) }
439    }
440
441    fn run(&mut self, io: &mut Io<'_, '_>, f: impl FnOnce(&mut WsCore, &mut WsIo)) {
442        let mut w = WsIo::new(Route::Conn(io.cx.handle()), 0, 0, io.now);
443        f(&mut self.core, &mut w);
444        if !w.out.is_empty() {
445            io.write(&w.out);
446        }
447        if w.panicked {
448            crate::server::bump(&io.sh.counters(io.li).panics);
449        }
450        if w.close {
451            io.close = true;
452        }
453    }
454}
455
456impl Layer for WsLayer {
457    fn start(&mut self, io: &mut Io<'_, '_>) {
458        self.run(io, |c, w| c.start(w));
459    }
460
461    fn input(&mut self, io: &mut Io<'_, '_>, data: &[u8]) -> usize {
462        if !self.core.is_done() {
463            self.run(io, |c, w| c.input(w, data));
464        }
465        data.len()
466    }
467
468    fn message(&mut self, io: &mut Io<'_, '_>, env: Envelope) {
469        let Msg::Ws(cmd) = env.msg else { return };
470        self.run(io, |c, w| c.command(w, cmd));
471    }
472
473    fn eof(&mut self, io: &mut Io<'_, '_>) {
474        self.run(io, |c, w| c.eof(w));
475    }
476
477    fn stop(&mut self, io: &mut Io<'_, '_>) {
478        self.run(io, |c, w| c.stop(w));
479    }
480
481    fn tick(&mut self, io: &mut Io<'_, '_>) {
482        if io.now >= self.core.deadline() {
483            self.run(io, |c, w| c.tick(w));
484        }
485    }
486
487    fn deadline(&self, _now: u64, _sh: &Shared) -> u64 {
488        self.core.deadline()
489    }
490
491    fn wants_pause(&self) -> bool {
492        false
493    }
494
495    fn closed(&mut self) {
496        self.core.closed();
497    }
498}
simoom/ws.rs

simoom/
ws.rs
