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

rust_web_server/websocket/
mod.rs
