rust_web_server/websocket/
connection.rs

1//! WebSocket endpoints in an ordinary handler: [`WebSocket::upgrade`]
2//! returns the `101 Switching Protocols` response, and once the server has
3//! sent it, the connection goes to the handler as a [`WebSocketConnection`]
4//! on a thread of its own (tokio's blocking pool in async builds, a new
5//! thread in the `http1` build). HTTP/1.1 only: an HTTP/2 or HTTP/3 request
6//! can't upgrade and gets `426 Upgrade Required`.
7//!
8//! How the handler reaches the server: the response carries an internal
9//! header with a random token, and the handler waits under that token in a
10//! process-wide table. The server removes the header, takes the handler and
11//! hands it the socket (with any bytes the client sent after the request).
12//! A response that never reaches a client (replaced by a middleware) leaves
13//! its handler behind for at most [`UPGRADE_TTL`]. An upstream can't use
14//! the header: tokens are random, and unknown ones do nothing.
15
16use std::collections::HashMap;
17use std::hash::{BuildHasher, Hasher};
18use std::io::{Read, Write};
19use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
20use std::sync::{Arc, Mutex, OnceLock};
21use std::time::{Duration, Instant};
22
23use super::{close_payload, encode_frame, read_raw_frame, Frame, RawFrame, WebSocket};
24use crate::core::New;
25use crate::header::Header;
26use crate::request::Request;
27use crate::response::{Response, STATUS_CODE_REASON_PHRASE};
28
29/// The internal header carrying the token from handler to server.
30const UPGRADE_HEADER: &str = "X-Rws-WebSocket-Upgrade";
31/// How long a handler waits for its response to be sent.
32pub const UPGRADE_TTL: Duration = Duration::from_secs(60);
33/// The default limit on one message (all its fragments): 16 MiB.
34pub const MAX_MESSAGE: usize = 16 * 1024 * 1024;
35
36type Handler = Box<dyn FnOnce(WebSocketConnection) + Send>;
37
38/// A complete message from the client: fragments joined, text checked to
39/// be UTF-8.
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub enum Message {
42    Text(String),
43    Binary(Vec<u8>),
44}
45
46/// An upgraded connection. [`recv`](Self::recv) returns the client's
47/// messages, answering pings and the closing handshake itself; send with
48/// the methods here or a [`WebSocketSender`] from [`sender`](Self::sender),
49/// which another thread can use while this one waits in `recv`.
50pub struct WebSocketConnection {
51    reader: Box<dyn Read + Send>,
52    sender: WebSocketSender,
53    max_message: usize,
54    closed: bool,
55}
56
57/// Sends to one WebSocket client; cheap to clone, usable from any thread.
58#[derive(Clone)]
59pub struct WebSocketSender {
60    writer: Arc<Mutex<Box<dyn Write + Send>>>,
61    closed: Arc<AtomicBool>,
62}
63
64impl WebSocketSender {
65    fn send(&self, frame: Frame) -> Result<(), String> {
66        let closing = matches!(frame, Frame::Close(..));
67        if self.closed.load(Ordering::Acquire) {
68            return Err("the WebSocket is closed".to_string());
69        }
70        if closing {
71            self.closed.store(true, Ordering::Release);
72        }
73        let bytes = encode_frame(frame);
74        let mut writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
75        writer.write_all(&bytes).and_then(|_| writer.flush()).map_err(|e| format!("write error: {e}"))
76    }
77
78    /// Send a text message.
79    pub fn send_text(&self, text: &str) -> Result<(), String> {
80        self.send(Frame::Text(text.to_string()))
81    }
82
83    /// Send a binary message.
84    pub fn send_binary(&self, data: &[u8]) -> Result<(), String> {
85        self.send(Frame::Binary(data.to_vec()))
86    }
87
88    /// Send a ping (at most 125 bytes); the client's pong is consumed by
89    /// [`WebSocketConnection::recv`].
90    pub fn ping(&self, payload: &[u8]) -> Result<(), String> {
91        if payload.len() > 125 {
92            return Err("a ping payload is at most 125 bytes".to_string());
93        }
94        self.send(Frame::Ping(payload.to_vec()))
95    }
96
97    /// Start the closing handshake with `code` (1000 = normal) and `reason`;
98    /// nothing more can be sent after it.
99    pub fn close(&self, code: u16, reason: &str) -> Result<(), String> {
100        self.send(Frame::Close(Some(code), reason.to_string()))
101    }
102
103    /// Whether a close frame was sent (by either side's handshake).
104    pub fn is_closed(&self) -> bool {
105        self.closed.load(Ordering::Acquire)
106    }
107}
108
109/// A close code the client may send (RFC 6455 §7.4).
110fn valid_close_code(code: u16) -> bool {
111    matches!(code, 1000..=1003 | 1007..=1011 | 3000..=4999)
112}
113
114impl WebSocketConnection {
115    /// A connection over `reader` (client → server) and `writer`.
116    pub(crate) fn new(reader: Box<dyn Read + Send>, writer: Box<dyn Write + Send>) -> Self {
117        WebSocketConnection {
118            reader,
119            sender: WebSocketSender { writer: Arc::new(Mutex::new(writer)), closed: Arc::new(AtomicBool::new(false)) },
120            max_message: MAX_MESSAGE,
121            closed: false,
122        }
123    }
124
125    /// Refuse messages over `bytes` (default [`MAX_MESSAGE`]) with close
126    /// code 1009.
127    pub fn max_message_size(mut self, bytes: usize) -> Self {
128        self.max_message = bytes;
129        self
130    }
131
132    /// A sender for this connection, for other threads.
133    pub fn sender(&self) -> WebSocketSender {
134        self.sender.clone()
135    }
136
137    /// Send a text message.
138    pub fn send_text(&self, text: &str) -> Result<(), String> {
139        self.sender.send_text(text)
140    }
141
142    /// Send a binary message.
143    pub fn send_binary(&self, data: &[u8]) -> Result<(), String> {
144        self.sender.send_binary(data)
145    }
146
147    /// Start the closing handshake; [`recv`](Self::recv) then returns
148    /// `Ok(None)` when the client answers.
149    pub fn close(&self, code: u16, reason: &str) -> Result<(), String> {
150        self.sender.close(code, reason)
151    }
152
153    /// Fail the connection: a close frame with `code`, then `Err(reason)`.
154    fn fail<T>(&mut self, code: u16, reason: String) -> Result<T, String> {
155        self.closed = true;
156        let _ = self.sender.close(code, "");
157        Err(reason)
158    }
159
160    /// The next message; `Ok(None)` once the connection closed cleanly
161    /// (the closing handshake done, whoever started it). Pings are
162    /// answered and pongs dropped on the way. A protocol error — an
163    /// unmasked frame, a bad fragment, invalid UTF-8, an oversized message
164    /// — closes the connection with the matching code and is returned as
165    /// `Err`, as is the client disappearing without a close frame.
166    pub fn recv(&mut self) -> Result<Option<Message>, String> {
167        if self.closed {
168            return Ok(None);
169        }
170        let mut message: Option<(u8, Vec<u8>)> = None;
171        loop {
172            let assembled = message.as_ref().map_or(0, |(_, data)| data.len());
173            let limit = self.max_message.saturating_sub(assembled) as u64;
174            let RawFrame { fin, rsv, opcode, masked, payload } = match read_raw_frame(&mut self.reader, limit) {
175                Ok(frame) => frame,
176                Err(e) if e.contains("exceeds") => return self.fail(1009, format!("message over {} bytes", self.max_message)),
177                Err(e) => {
178                    self.closed = true;
179                    return Err(e);
180                }
181            };
182            if rsv != 0 {
183                return self.fail(1002, "reserved bits set (no extension was negotiated)".to_string());
184            }
185            if !masked {
186                return self.fail(1002, "a client frame wasn't masked".to_string());
187            }
188            if opcode >= 0x8 && (!fin || payload.len() > 125) {
189                return self.fail(1002, "a control frame was fragmented or over 125 bytes".to_string());
190            }
191            match opcode {
192                0x8 => {
193                    let (code, _) = match close_payload(&payload) {
194                        Ok(close) => close,
195                        Err(e) => return self.fail(1007, e),
196                    };
197                    if payload.len() == 1 || code.is_some_and(|c| !valid_close_code(c)) {
198                        return self.fail(1002, format!("invalid close code {code:?}"));
199                    }
200                    self.closed = true;
201                    // Answer the client's close, unless this ends ours.
202                    if !self.sender.is_closed() {
203                        let _ = self.sender.close(code.unwrap_or(1000), "");
204                    }
205                    return Ok(None);
206                }
207                0x9 => {
208                    let _ = self.sender.send(Frame::Pong(payload));
209                }
210                0xA => {}
211                0x1 | 0x2 if message.is_some() => return self.fail(1002, "a new message began inside a fragmented one".to_string()),
212                0x1 | 0x2 => message = Some((opcode, payload)),
213                0x0 => match message.as_mut() {
214                    Some((_, data)) => data.extend_from_slice(&payload),
215                    None => return self.fail(1002, "a continuation frame with no message".to_string()),
216                },
217                other => return self.fail(1002, format!("unknown opcode 0x{other:X}")),
218            }
219            if fin && opcode < 0x8 {
220                let Some((kind, data)) = message.take() else { continue };
221                return match kind {
222                    0x1 => match String::from_utf8(data) {
223                        Ok(text) => Ok(Some(Message::Text(text))),
224                        Err(_) => self.fail(1007, "a text message isn't UTF-8".to_string()),
225                    },
226                    _ => Ok(Some(Message::Binary(data))),
227                };
228            }
229        }
230    }
231}
232
233// ── Upgrade ───────────────────────────────────────────────────────────────────
234
235/// A handler waiting for its connection.
236pub(crate) struct Upgrade(Handler);
237
238impl Upgrade {
239    /// Run the handler on `reader` / `writer` (blocks for the connection's
240    /// life).
241    pub(crate) fn run(self, reader: Box<dyn Read + Send>, writer: Box<dyn Write + Send>) {
242        (self.0)(WebSocketConnection::new(reader, writer))
243    }
244}
245
246fn waiting() -> &'static Mutex<HashMap<String, (Instant, Handler)>> {
247    static WAITING: OnceLock<Mutex<HashMap<String, (Instant, Handler)>>> = OnceLock::new();
248    WAITING.get_or_init(Default::default)
249}
250
251/// 128 random bits as hex: SipHash with the process's random keys over a
252/// counter and the time.
253fn token() -> String {
254    static KEYS: OnceLock<std::collections::hash_map::RandomState> = OnceLock::new();
255    static COUNTER: AtomicU64 = AtomicU64::new(0);
256    let keys = KEYS.get_or_init(Default::default);
257    let n = COUNTER.fetch_add(1, Ordering::Relaxed);
258    let nanos = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_or(0, |d| d.as_nanos());
259    let half = |salt: u8| {
260        let mut hasher = keys.build_hasher();
261        hasher.write_u8(salt);
262        hasher.write_u64(n);
263        hasher.write_u128(nanos);
264        hasher.finish()
265    };
266    format!("{:016x}{:016x}", half(1), half(2))
267}
268
269impl WebSocket {
270    /// Answer a WebSocket upgrade request: returns the `101` response to
271    /// return from the handler, and runs `handler` with the connection once
272    /// it's sent. A request that isn't a valid upgrade (or comes over HTTP/2
273    /// or HTTP/3) gets `426 Upgrade Required` and `handler` never runs.
274    ///
275    /// ```rust,no_run
276    /// use rust_web_server::app::App;
277    /// use rust_web_server::websocket::{Message, WebSocket};
278    ///
279    /// let app = App::with_state(()).get("/echo", |request, _params, _conn, _state| {
280    ///     WebSocket::upgrade(request, |mut ws| {
281    ///         while let Ok(Some(message)) = ws.recv() {
282    ///             let _ = match message {
283    ///                 Message::Text(text) => ws.send_text(&text),
284    ///                 Message::Binary(data) => ws.send_binary(&data),
285    ///             };
286    ///         }
287    ///     })
288    /// });
289    /// ```
290    pub fn upgrade<F>(request: &Request, handler: F) -> Response
291    where
292        F: FnOnce(WebSocketConnection) + Send + 'static,
293    {
294        let version = request.header("Sec-WebSocket-Version").map(|h| h.value.trim().to_string()).unwrap_or_default();
295        let mut response = match WebSocket::handshake_response(request) {
296            Ok(response) if Self::is_upgrade_request(request) && version == "13" && request.method == "GET" => response,
297            _ => {
298                let mut refused = Response::new();
299                refused.status_code = *STATUS_CODE_REASON_PHRASE.n426_upgrade_required.status_code;
300                refused.reason_phrase = STATUS_CODE_REASON_PHRASE.n426_upgrade_required.reason_phrase.to_string();
301                refused.headers.push(Header { name: "Upgrade".into(), value: "websocket".into() });
302                refused.headers.push(Header { name: "Sec-WebSocket-Version".into(), value: "13".into() });
303                return refused;
304            }
305        };
306        let token = token();
307        let mut waiting = waiting().lock().unwrap_or_else(|e| e.into_inner());
308        let now = Instant::now();
309        waiting.retain(|_, (since, _)| now.duration_since(*since) < UPGRADE_TTL);
310        waiting.insert(token.clone(), (now, Box::new(handler)));
311        response.headers.push(Header { name: UPGRADE_HEADER.into(), value: token.into() });
312        response
313    }
314}
315
316/// The handler waiting for `response`, if it's a `101` from
317/// [`WebSocket::upgrade`]; the internal header is removed either way.
318pub(crate) fn take_upgrade(response: &mut Response) -> Option<Upgrade> {
319    let at = response.headers.iter().position(|h| h.name.eq_ignore_ascii_case(UPGRADE_HEADER))?;
320    let token = response.headers.remove(at).value;
321    if response.status_code != 101 {
322        return None;
323    }
324    let (_, handler) = waiting().lock().unwrap_or_else(|e| e.into_inner()).remove(token.as_ref())?;
325    Some(Upgrade(handler))
326}
327
328/// The `101` response head as sent: the status line and the response's
329/// headers — no body, no `Content-Length` (RFC 9110 §8.6).
330pub(crate) fn upgrade_head(response: &Response) -> Vec<u8> {
331    let mut head = format!("HTTP/1.1 101 {}\r\n", response.reason_phrase);
332    for header in &response.headers {
333        if header.name.eq_ignore_ascii_case("content-length") || header.name.eq_ignore_ascii_case("transfer-encoding") {
334            continue;
335        }
336        head.push_str(&header.name);
337        head.push_str(": ");
338        head.push_str(&header.value);
339        head.push_str("\r\n");
340    }
341    head.push_str("\r\n");
342    head.into_bytes()
343}
344
345/// Blocking `Read`/`Write` over one half of an async stream, for a handler
346/// on the blocking pool.
347#[cfg(feature = "http2")]
348struct Blocking<T> {
349    inner: T,
350    handle: tokio::runtime::Handle,
351}
352
353#[cfg(feature = "http2")]
354impl<T: tokio::io::AsyncRead + Unpin> Read for Blocking<T> {
355    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
356        self.handle.block_on(tokio::io::AsyncReadExt::read(&mut self.inner, buf))
357    }
358}
359
360#[cfg(feature = "http2")]
361impl<T: tokio::io::AsyncWrite + Unpin> Write for Blocking<T> {
362    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
363        self.handle.block_on(tokio::io::AsyncWriteExt::write(&mut self.inner, buf))
364    }
365
366    fn flush(&mut self) -> std::io::Result<()> {
367        self.handle.block_on(tokio::io::AsyncWriteExt::flush(&mut self.inner))
368    }
369}
370
371/// Give an async connection (plain or TLS) to `upgrade` on the blocking
372/// pool, `pending` (bytes already read past the request) first.
373#[cfg(feature = "http2")]
374pub(crate) fn hand_over_async<S>(stream: S, pending: Vec<u8>, upgrade: Upgrade)
375where
376    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static,
377{
378    let handle = tokio::runtime::Handle::current();
379    let (read, write) = tokio::io::split(stream);
380    let reader = std::io::Cursor::new(pending).chain(Blocking { inner: read, handle: handle.clone() });
381    let writer = Blocking { inner: write, handle };
382    tokio::task::spawn_blocking(move || upgrade.run(Box::new(reader), Box::new(writer)));
383}
rust_web_server/websocket/connection.rs

rust_web_server/websocket/
connection.rs
