rust_web_server/websocket/
mod.rs

1//! WebSocket protocol support — RFC 6455.
2//!
3//! Provides the building blocks for adding WebSocket endpoints to a server:
4//! upgrade detection, the opening handshake, and frame-level read/write.
5//!
6//! # Integration pattern
7//!
8//! WebSocket connections require taking over the raw TCP stream after sending
9//! the 101 response, so they cannot be handled inside a normal
10//! [`Controller::process`](crate::controller::Controller) call (which has no
11//! access to the stream).  The recommended pattern is to bypass
12//! [`Server::run`](crate::server::Server) for upgraded connections and drive
13//! them with your own accept loop:
14//!
15//! ```rust,no_run
16//! use std::net::TcpListener;
17//! use rust_web_server::websocket::{WebSocket, Frame};
18//! use rust_web_server::request::Request;
19//! use rust_web_server::response::Response;
20//!
21//! let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
22//! for stream in listener.incoming() {
23//!     let mut stream = stream.unwrap();
24//!     // peek / read request bytes, decide if it's a WS upgrade
25//!     // (simplified — real code should read the full request first)
26//!     let raw = vec![0u8; 4096];
27//!     // ... parse into Request, then:
28//!     //
29//!     // if WebSocket::is_upgrade_request(&request) {
30//!     //     let response = WebSocket::handshake_response(&request).unwrap();
31//!     //     // write the 101 response to stream, then frame loop
32//!     //     loop {
33//!     //         match WebSocket::read_frame(&mut stream) { ... }
34//!     //     }
35//!     // } else {
36//!     //     // normal Server::process
37//!     // }
38//! }
39//! ```
40
41#[cfg(test)]
42mod tests;
43mod connection;
44
45pub use connection::{Message, WebSocketConnection, WebSocketSender};
46pub(crate) use connection::{take_upgrade, upgrade_head};
47#[cfg(feature = "http2")]
48pub(crate) use connection::hand_over_async;
49pub use connection::{MAX_MESSAGE, UPGRADE_TTL};
50
51use std::io::{Read, Write};
52
53use crate::header::Header;
54use crate::http::VERSION;
55use crate::request::Request;
56use crate::response::{Response, STATUS_CODE_REASON_PHRASE};
57
58/// A decoded WebSocket frame.
59#[derive(Debug, PartialEq, Eq)]
60pub enum Frame {
61    /// A complete UTF-8 text message (opcode 0x1).
62    Text(String),
63    /// A binary message (opcode 0x2).
64    Binary(Vec<u8>),
65    /// A ping control frame (opcode 0x9). Server should respond with Pong.
66    Ping(Vec<u8>),
67    /// A pong control frame (opcode 0xA).
68    Pong(Vec<u8>),
69    /// A close control frame (opcode 0x8) with an optional status code and reason.
70    Close(Option<u16>, String),
71    /// A continuation fragment (opcode 0x0).
72    Continuation { fin: bool, data: Vec<u8> },
73}
74
75/// WebSocket protocol utilities.
76pub struct WebSocket;
77
78impl WebSocket {
79    const MAGIC: &'static str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
80
81    /// The largest frame payload [`WebSocket::read_frame`] accepts (16 MiB);
82    /// a frame claiming more is an error. Use
83    /// [`WebSocket::read_frame_limited`] for another limit.
84    pub const MAX_FRAME_PAYLOAD: u64 = 16 * 1024 * 1024;
85
86    // ── Handshake ─────────────────────────────────────────────────────────────
87
88    /// Returns `true` if `request` is a valid WebSocket upgrade request
89    /// (has `Upgrade: websocket`, `Connection: upgrade`, and `Sec-WebSocket-Key`).
90    pub fn is_upgrade_request(request: &Request) -> bool {
91        let upgrade = request.get_header("Upgrade".to_string());
92        let connection = request.get_header("Connection".to_string());
93        let key = request.get_header("Sec-WebSocket-Key".to_string());
94
95        upgrade.map(|h| h.value.to_lowercase() == "websocket").unwrap_or(false)
96            && connection.map(|h| h.value.to_lowercase().contains("upgrade")).unwrap_or(false)
97            && key.is_some()
98    }
99
100    /// Build the HTTP `101 Switching Protocols` response for a WebSocket
101    /// opening handshake. Returns an error if `Sec-WebSocket-Key` is absent.
102    ///
103    /// Write the raw bytes from [`Response::generate_response`] to the stream,
104    /// then transition to frame-level I/O with [`WebSocket::read_frame`] /
105    /// [`WebSocket::write_frame`].
106    pub fn handshake_response(request: &Request) -> Result<Response, String> {
107        let key_header = request.get_header("Sec-WebSocket-Key".to_string())
108            .ok_or_else(|| "missing Sec-WebSocket-Key header".to_string())?;
109        let accept = Self::accept_key(&key_header.value);
110
111        let mut response = Response {
112            http_version: VERSION.http_1_1.to_string(),
113            status_code: *STATUS_CODE_REASON_PHRASE.n101_switching_protocols.status_code,
114            reason_phrase: STATUS_CODE_REASON_PHRASE.n101_switching_protocols.reason_phrase.to_string(),
115            headers: vec![
116                Header { name: "Upgrade".into(),              value: "websocket".into() },
117                Header { name: "Connection".into(),           value: "Upgrade".into() },
118                Header { name: "Sec-WebSocket-Accept".into(), value: accept.into() },
119            ],
120            content_range_list: vec![],
121            stream_file: None,
122            stream_pipe: None,
123            trailers: None,
124        };
125
126        if let Some(proto) = request.get_header("Sec-WebSocket-Protocol".to_string()) {
127            response.headers.push(Header {
128                name: "Sec-WebSocket-Protocol".into(),
129                value: proto.value.split(',').next().unwrap_or("").trim().to_string().into(),
130            });
131        }
132
133        Ok(response)
134    }
135
136    /// Compute the `Sec-WebSocket-Accept` value from the client's
137    /// `Sec-WebSocket-Key` using SHA-1 and base64 as specified in RFC 6455.
138    pub fn accept_key(client_key: &str) -> String {
139        let mut data = client_key.as_bytes().to_vec();
140        data.extend_from_slice(Self::MAGIC.as_bytes());
141        let hash = sha1(&data);
142        base64_encode(&hash)
143    }
144
145    // ── Frame I/O ─────────────────────────────────────────────────────────────
146
147    /// Read one WebSocket frame from `stream`.
148    ///
149    /// Handles client-to-server masking automatically. Returns an error if the
150    /// stream closes unexpectedly or contains a protocol violation.
151    ///
152    /// A frame whose declared payload exceeds [`WebSocket::MAX_FRAME_PAYLOAD`]
153    /// is refused before any of it is read: the length comes from the peer,
154    /// and allocating it up front (up to 2^64 − 1 bytes) aborted the process.
155    pub fn read_frame(stream: &mut impl Read) -> Result<Frame, String> {
156        Self::read_frame_limited(stream, Self::MAX_FRAME_PAYLOAD)
157    }
158
159    /// [`WebSocket::read_frame`] with `max_payload` bytes as the limit. The
160    /// payload is read as it arrives, so memory follows what the peer
161    /// actually sends, not what it claims.
162    pub fn read_frame_limited(stream: &mut impl Read, max_payload: u64) -> Result<Frame, String> {
163        let RawFrame { fin, opcode, payload, .. } = read_raw_frame(stream, max_payload)?;
164        let frame = match opcode {
165            0x0 => Frame::Continuation { fin, data: payload },
166            0x1 => {
167                let text = String::from_utf8(payload)
168                    .map_err(|_| "text frame contains invalid UTF-8".to_string())?;
169                Frame::Text(text)
170            }
171            0x2 => Frame::Binary(payload),
172            0x8 => {
173                let (code, reason) = close_payload(&payload)?;
174                Frame::Close(code, reason)
175            }
176            0x9 => Frame::Ping(payload),
177            0xA => Frame::Pong(payload),
178            n   => return Err(format!("unknown opcode: 0x{:X}", n)),
179        };
180
181        Ok(frame)
182    }
183
184    /// Write a WebSocket frame to `stream` (server→client, unmasked).
185    pub fn write_frame(stream: &mut impl Write, frame: Frame) -> Result<(), String> {
186        stream.write_all(&encode_frame(frame)).map_err(|e| format!("write error: {}", e))?;
187        stream.flush().map_err(|e| format!("flush error: {}", e))?;
188        Ok(())
189    }
190
191    /// Convenience: send a text message.
192    pub fn send_text(stream: &mut impl Write, text: &str) -> Result<(), String> {
193        Self::write_frame(stream, Frame::Text(text.to_string()))
194    }
195
196    /// Convenience: send a close frame with a status code and reason.
197    pub fn send_close(stream: &mut impl Write, code: u16, reason: &str) -> Result<(), String> {
198        Self::write_frame(stream, Frame::Close(Some(code), reason.to_string()))
199    }
200
201    /// Convenience: reply to a ping with a pong carrying the same payload.
202    pub fn send_pong(stream: &mut impl Write, payload: Vec<u8>) -> Result<(), String> {
203        Self::write_frame(stream, Frame::Pong(payload))
204    }
205}
206
207/// A frame as read: header fields and the unmasked payload.
208pub(crate) struct RawFrame {
209    pub(crate) fin: bool,
210    /// RSV1–3 (no extension is negotiated, so they must be 0).
211    pub(crate) rsv: u8,
212    pub(crate) opcode: u8,
213    pub(crate) masked: bool,
214    pub(crate) payload: Vec<u8>,
215}
216
217/// A close frame's status code and reason (RFC 6455 §5.5.1: the reason is
218/// UTF-8, else the connection fails — it used to be decoded lossily,
219/// growing invalid bytes threefold).
220pub(crate) fn close_payload(payload: &[u8]) -> Result<(Option<u16>, String), String> {
221    let code = if payload.len() >= 2 { Some(u16::from_be_bytes([payload[0], payload[1]])) } else { None };
222    let reason = if payload.len() > 2 {
223        String::from_utf8(payload[2..].to_vec()).map_err(|_| "close frame reason contains invalid UTF-8".to_string())?
224    } else {
225        String::new()
226    };
227    Ok((code, reason))
228}
229
230/// Read one frame, unmasking its payload; a payload over `max_payload` is
231/// refused before any of it is read.
232pub(crate) fn read_raw_frame(stream: &mut impl Read, max_payload: u64) -> Result<RawFrame, String> {
233    let mut header = [0u8; 2];
234    read_exact(stream, &mut header)?;
235
236    let fin = (header[0] & 0x80) != 0;
237    let rsv = (header[0] >> 4) & 0x07;
238    let opcode = header[0] & 0x0F;
239    let masked = (header[1] & 0x80) != 0;
240    let payload_len_byte = (header[1] & 0x7F) as usize;
241
242    let payload_len: u64 = match payload_len_byte {
243        126 => {
244            let mut ext = [0u8; 2];
245            read_exact(stream, &mut ext)?;
246            u16::from_be_bytes(ext) as u64
247        }
248        127 => {
249            let mut ext = [0u8; 8];
250            read_exact(stream, &mut ext)?;
251            u64::from_be_bytes(ext)
252        }
253        n => n as u64,
254    };
255    if payload_len > max_payload {
256        return Err(format!("frame payload of {payload_len} bytes exceeds the {max_payload}-byte limit"));
257    }
258
259    let mask_key = if masked {
260        let mut mk = [0u8; 4];
261        read_exact(stream, &mut mk)?;
262        Some(mk)
263    } else {
264        None
265    };
266
267    let mut payload = Vec::new();
268    if payload_len > 0 {
269        stream
270            .take(payload_len)
271            .read_to_end(&mut payload)
272            .map_err(|e| format!("read error: {}", e))?;
273        if payload.len() as u64 != payload_len {
274            return Err("read error: stream ended inside a frame".to_string());
275        }
276    }
277
278    if let Some(key) = mask_key {
279        for (i, byte) in payload.iter_mut().enumerate() {
280            *byte ^= key[i % 4];
281        }
282    }
283
284    Ok(RawFrame { fin, rsv, opcode, masked, payload })
285}
286
287/// A frame's bytes on the wire (server→client, unmasked).
288pub(crate) fn encode_frame(frame: Frame) -> Vec<u8> {
289    let (opcode, payload, fin) = match frame {
290        Frame::Text(s)          => (0x1u8, s.into_bytes(), true),
291        Frame::Binary(b)        => (0x2, b, true),
292        Frame::Ping(b)          => (0x9, b, true),
293        Frame::Pong(b)          => (0xA, b, true),
294        Frame::Close(code, reason) => {
295            let mut payload = Vec::new();
296            if let Some(c) = code {
297                payload.extend_from_slice(&c.to_be_bytes());
298                payload.extend_from_slice(reason.as_bytes());
299            }
300            (0x8, payload, true)
301        }
302        Frame::Continuation { fin, data } => (0x0, data, fin),
303    };
304
305    let fin_bit: u8 = if fin { 0x80 } else { 0x00 };
306    let payload_len = payload.len();
307    let mut out = Vec::with_capacity(10 + payload_len);
308    out.push(fin_bit | opcode);
309    match payload_len {
310        0..=125 => out.push(payload_len as u8),
311        126..=65535 => {
312            out.push(126u8);
313            out.extend_from_slice(&(payload_len as u16).to_be_bytes());
314        }
315        _ => {
316            out.push(127u8);
317            out.extend_from_slice(&(payload_len as u64).to_be_bytes());
318        }
319    }
320    out.extend_from_slice(&payload);
321    out
322}
323
324// ── Internal utilities ─────────────────────────────────────────────────────────
325
326fn read_exact<R: Read>(r: &mut R, buf: &mut [u8]) -> Result<(), String> {
327    r.read_exact(buf).map_err(|e| format!("read error: {}", e))
328}
329
330/// SHA-1 digest of `data` (FIPS 180-4). Used for the WebSocket handshake only.
331pub(crate) fn sha1(data: &[u8]) -> [u8; 20] {
332    let mut h = [0x67452301u32, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0];
333
334    let msg_len = data.len();
335    let mut msg = data.to_vec();
336    msg.push(0x80);
337    while msg.len() % 64 != 56 {
338        msg.push(0x00);
339    }
340    msg.extend_from_slice(&((msg_len as u64) * 8).to_be_bytes());
341
342    for chunk in msg.chunks(64) {
343        let mut w = [0u32; 80];
344        for i in 0..16 {
345            w[i] = u32::from_be_bytes([chunk[i*4], chunk[i*4+1], chunk[i*4+2], chunk[i*4+3]]);
346        }
347        for i in 16..80 {
348            w[i] = (w[i-3] ^ w[i-8] ^ w[i-14] ^ w[i-16]).rotate_left(1);
349        }
350
351        let (mut a, mut b, mut c, mut d, mut e) = (h[0], h[1], h[2], h[3], h[4]);
352
353        for i in 0..80 {
354            let (f, k) = match i {
355                0..=19  => ((b & c) | (!b & d),          0x5A827999u32),
356                20..=39 => (b ^ c ^ d,                    0x6ED9EBA1u32),
357                40..=59 => ((b & c) | (b & d) | (c & d), 0x8F1BBCDCu32),
358                _       => (b ^ c ^ d,                    0xCA62C1D6u32),
359            };
360            let temp = a.rotate_left(5)
361                .wrapping_add(f)
362                .wrapping_add(e)
363                .wrapping_add(k)
364                .wrapping_add(w[i]);
365            e = d; d = c; c = b.rotate_left(30); b = a; a = temp;
366        }
367
368        h[0] = h[0].wrapping_add(a);
369        h[1] = h[1].wrapping_add(b);
370        h[2] = h[2].wrapping_add(c);
371        h[3] = h[3].wrapping_add(d);
372        h[4] = h[4].wrapping_add(e);
373    }
374
375    let mut out = [0u8; 20];
376    for (i, v) in h.iter().enumerate() {
377        out[i*4..(i+1)*4].copy_from_slice(&v.to_be_bytes());
378    }
379    out
380}
381
382/// Standard base64 encoding (RFC 4648 Table 1, with `=` padding).
383pub(crate) fn base64_encode(data: &[u8]) -> String {
384    const T: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
385    let mut out = Vec::with_capacity((data.len() + 2) / 3 * 4);
386    for chunk in data.chunks(3) {
387        let b0 = chunk[0] as u32;
388        let b1 = if chunk.len() > 1 { chunk[1] as u32 } else { 0 };
389        let b2 = if chunk.len() > 2 { chunk[2] as u32 } else { 0 };
390        let n = (b0 << 16) | (b1 << 8) | b2;
391        out.push(T[((n >> 18) & 0x3F) as usize]);
392        out.push(T[((n >> 12) & 0x3F) as usize]);
393        out.push(if chunk.len() > 1 { T[((n >> 6) & 0x3F) as usize] } else { b'=' });
394        out.push(if chunk.len() > 2 { T[(n & 0x3F) as usize] } else { b'=' });
395    }
396    String::from_utf8(out).unwrap()
397}
rust_web_server/websocket/mod.rs

rust_web_server/websocket/
mod.rs
