magicka-vm/crates/server/src/ws.rs

//! 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<Opcode> {
        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<u8>,
}

fn read_frame<R: Read>(r: &mut R) -> io::Result<Frame> {
    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<u8>),
    /// 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<R: Read>(stream: &mut R) -> io::Result<Option<Message>> {
    let mut buf: Vec<u8> = Vec::new();
    let mut msg_op: Option<Opcode> = 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: Write>(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: Write>(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: Write>(w: &mut W, payload: &[u8]) -> io::Result<()> {
    write_frame(w, Opcode::Pong, payload)
}

/// Send a close frame.
pub fn write_close<W: Write>(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());
    }
}