//! Minimal WebSocket (RFC 6455) support over `std::net`, no external crates. //! Implements just what the game needs: the upgrade handshake (SHA1 + base64), //! masked client-frame reading with fragment reassembly, and unmasked //! server-frame writing. All reads are length-checked so a hostile frame //! returns an `Err`, never a panic or unbounded allocation. use std::io::{self, Read, Write}; const WS_GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; /// Reject any single message larger than this (defensive bound). pub const MAX_MESSAGE: usize = 1 << 20; // 1 MiB // --------------------------------------------------------------------------- // SHA-1 (FIPS 180-1). Used only for the handshake accept key. // --------------------------------------------------------------------------- fn sha1(data: &[u8]) -> [u8; 20] { let mut h: [u32; 5] = [0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0]; let ml = (data.len() as u64) * 8; let mut msg = data.to_vec(); msg.push(0x80); while msg.len() % 64 != 56 { msg.push(0); } msg.extend_from_slice(&ml.to_be_bytes()); for chunk in msg.chunks_exact(64) { let mut w = [0u32; 80]; for (i, wi) in w.iter_mut().enumerate().take(16) { *wi = u32::from_be_bytes([ chunk[i * 4], chunk[i * 4 + 1], chunk[i * 4 + 2], chunk[i * 4 + 3], ]); } for i in 16..80 { w[i] = (w[i - 3] ^ w[i - 8] ^ w[i - 14] ^ w[i - 16]).rotate_left(1); } let (mut a, mut b, mut c, mut d, mut e) = (h[0], h[1], h[2], h[3], h[4]); for (i, &wi) in w.iter().enumerate() { let (f, k) = match i { 0..=19 => ((b & c) | ((!b) & d), 0x5A827999u32), 20..=39 => (b ^ c ^ d, 0x6ED9EBA1), 40..=59 => ((b & c) | (b & d) | (c & d), 0x8F1BBCDC), _ => (b ^ c ^ d, 0xCA62C1D6), }; let tmp = a .rotate_left(5) .wrapping_add(f) .wrapping_add(e) .wrapping_add(k) .wrapping_add(wi); e = d; d = c; c = b.rotate_left(30); b = a; a = tmp; } h[0] = h[0].wrapping_add(a); h[1] = h[1].wrapping_add(b); h[2] = h[2].wrapping_add(c); h[3] = h[3].wrapping_add(d); h[4] = h[4].wrapping_add(e); } let mut out = [0u8; 20]; for (i, hi) in h.iter().enumerate() { out[i * 4..i * 4 + 4].copy_from_slice(&hi.to_be_bytes()); } out } // --------------------------------------------------------------------------- // base64 (standard alphabet). // --------------------------------------------------------------------------- fn base64(data: &[u8]) -> String { const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; let mut out = String::new(); for chunk in data.chunks(3) { let b = [ chunk[0], *chunk.get(1).unwrap_or(&0), *chunk.get(2).unwrap_or(&0), ]; let n = ((b[0] as u32) << 16) | ((b[1] as u32) << 8) | (b[2] as u32); out.push(ALPHABET[((n >> 18) & 63) as usize] as char); out.push(ALPHABET[((n >> 12) & 63) as usize] as char); if chunk.len() > 1 { out.push(ALPHABET[((n >> 6) & 63) as usize] as char); } else { out.push('='); } if chunk.len() > 2 { out.push(ALPHABET[(n & 63) as usize] as char); } else { out.push('='); } } out } /// Compute the `Sec-WebSocket-Accept` value for a client key. pub fn accept_key(client_key: &str) -> String { let mut concat = client_key.to_string(); concat.push_str(WS_GUID); base64(&sha1(concat.as_bytes())) } // --------------------------------------------------------------------------- // Frames. // --------------------------------------------------------------------------- #[derive(Clone, Copy, PartialEq, Eq, Debug)] pub enum Opcode { Continuation, Text, Binary, Close, Ping, Pong, } impl Opcode { fn from_u8(v: u8) -> Option { Some(match v { 0x0 => Opcode::Continuation, 0x1 => Opcode::Text, 0x2 => Opcode::Binary, 0x8 => Opcode::Close, 0x9 => Opcode::Ping, 0xA => Opcode::Pong, _ => return None, }) } } struct Frame { fin: bool, opcode: Opcode, payload: Vec, } fn read_frame(r: &mut R) -> io::Result { let mut hdr = [0u8; 2]; r.read_exact(&mut hdr)?; let fin = hdr[0] & 0x80 != 0; let opcode = Opcode::from_u8(hdr[0] & 0x0f) .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bad opcode"))?; let masked = hdr[1] & 0x80 != 0; let len7 = (hdr[1] & 0x7f) as usize; let len = match len7 { 126 => { let mut b = [0u8; 2]; r.read_exact(&mut b)?; u16::from_be_bytes(b) as usize } 127 => { let mut b = [0u8; 8]; r.read_exact(&mut b)?; u64::from_be_bytes(b) as usize } n => n, }; if len > MAX_MESSAGE { return Err(io::Error::new(io::ErrorKind::InvalidData, "frame too large")); } // Per RFC, client frames MUST be masked. let mask = if masked { let mut m = [0u8; 4]; r.read_exact(&mut m)?; Some(m) } else { None }; let mut payload = vec![0u8; len]; r.read_exact(&mut payload)?; if let Some(m) = mask { for (i, b) in payload.iter_mut().enumerate() { *b ^= m[i % 4]; } } Ok(Frame { fin, opcode, payload }) } /// A complete application message read from the socket. Control frames are /// surfaced rather than answered inline so that *all* socket writes can be /// funneled through a single writer (avoiding interleaved frames when a server /// is both broadcasting and answering pings). pub enum Message { Text(String), /// A ping with its payload; the caller must reply with a pong. Ping(Vec), /// A pong (informational). Pong, /// The peer requested close. Close, } /// Read one full WebSocket message, reassembling fragments. Returns `Ok(None)` /// on a clean EOF. Reads only — never writes to the socket. pub fn read_message(stream: &mut R) -> io::Result> { let mut buf: Vec = Vec::new(); let mut msg_op: Option = None; loop { let frame = match read_frame(stream) { Ok(f) => f, Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), Err(e) => return Err(e), }; match frame.opcode { Opcode::Close => return Ok(Some(Message::Close)), Opcode::Ping => return Ok(Some(Message::Ping(frame.payload))), Opcode::Pong => return Ok(Some(Message::Pong)), Opcode::Text | Opcode::Binary => { if msg_op.is_some() { return Err(io::Error::new(io::ErrorKind::InvalidData, "interleaved frame")); } msg_op = Some(frame.opcode); buf.extend_from_slice(&frame.payload); } Opcode::Continuation => { if msg_op.is_none() { return Err(io::Error::new(io::ErrorKind::InvalidData, "stray continuation")); } buf.extend_from_slice(&frame.payload); } } if buf.len() > MAX_MESSAGE { return Err(io::Error::new(io::ErrorKind::InvalidData, "message too large")); } if frame.fin { // We only surface text to the application; binary is decoded lossily. let s = String::from_utf8_lossy(&buf).into_owned(); return Ok(Some(Message::Text(s))); } } } fn write_frame(w: &mut W, opcode: Opcode, payload: &[u8]) -> io::Result<()> { let op = match opcode { Opcode::Continuation => 0x0, Opcode::Text => 0x1, Opcode::Binary => 0x2, Opcode::Close => 0x8, Opcode::Ping => 0x9, Opcode::Pong => 0xA, }; let mut frame = vec![0x80 | op]; let len = payload.len(); if len < 126 { frame.push(len as u8); } else if len < 65536 { frame.push(126); frame.extend_from_slice(&(len as u16).to_be_bytes()); } else { frame.push(127); frame.extend_from_slice(&(len as u64).to_be_bytes()); } frame.extend_from_slice(payload); w.write_all(&frame)?; w.flush() } /// Send a text message (server frames are never masked). pub fn write_text(w: &mut W, text: &str) -> io::Result<()> { write_frame(w, Opcode::Text, text.as_bytes()) } /// Send a pong frame echoing a ping payload. pub fn write_pong(w: &mut W, payload: &[u8]) -> io::Result<()> { write_frame(w, Opcode::Pong, payload) } /// Send a close frame. pub fn write_close(w: &mut W) -> io::Result<()> { write_frame(w, Opcode::Close, &[]) } #[cfg(test)] mod tests { use super::*; #[test] fn rfc_example_accept_key() { // The canonical example from RFC 6455 section 1.3. assert_eq!( accept_key("dGhlIHNhbXBsZSBub25jZQ=="), "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=" ); } #[test] fn sha1_known_vector() { // "abc" -> a9993e364706816aba3e25717850c26c9cd0d89d let d = sha1(b"abc"); let hex: String = d.iter().map(|b| format!("{b:02x}")).collect(); assert_eq!(hex, "a9993e364706816aba3e25717850c26c9cd0d89d"); } #[test] fn base64_roundtrip_lengths() { assert_eq!(base64(b""), ""); assert_eq!(base64(b"f"), "Zg=="); assert_eq!(base64(b"fo"), "Zm8="); assert_eq!(base64(b"foo"), "Zm9v"); assert_eq!(base64(b"foobar"), "Zm9vYmFy"); } #[test] fn masked_text_frame_roundtrips_through_reader() { use std::io::Cursor; // Build a masked client text frame for "hi". let payload = b"hi"; let mask = [0x01, 0x02, 0x03, 0x04]; let mut frame = vec![0x81, 0x80 | payload.len() as u8]; frame.extend_from_slice(&mask); for (i, &b) in payload.iter().enumerate() { frame.push(b ^ mask[i % 4]); } // Cursor implements Read+Write (write goes nowhere useful but pong path // is not exercised here). let mut cur = Cursor::new(frame); match read_message(&mut cur).unwrap() { Some(Message::Text(s)) => assert_eq!(s, "hi"), _ => panic!("expected text"), } } #[test] fn oversized_frame_is_rejected() { use std::io::Cursor; // Declares a 127-length (8-byte) payload of u64::MAX — must error, not OOM. let mut frame = vec![0x81, 0x80 | 127]; frame.extend_from_slice(&u64::MAX.to_be_bytes()); frame.extend_from_slice(&[0, 0, 0, 0]); // partial mask let mut cur = Cursor::new(frame); assert!(read_message(&mut cur).is_err()); } }