magicka-vm/crates/web_tests/src/lib.rs

//! `web_tests` — a tiny, dependency-free WebSocket *client* used to drive the
//! real server over a real socket in integration tests. It performs the HTTP
//! upgrade, masks client frames (as RFC 6455 requires), and reads server
//! frames. This is the harness behind the Phase H web CI gates.

use std::io::{self, BufRead, BufReader, Read, Write};
use std::net::TcpStream;
use std::time::Duration;

use protocol::{ClientMessage, ServerMessage};

/// A blocking WebSocket client connection to the test server.
pub struct WsClient {
    stream: TcpStream,
    reader: BufReader<TcpStream>,
}

impl WsClient {
    /// Connect, upgrade to WebSocket, and verify the handshake.
    pub fn connect(addr: &str) -> io::Result<WsClient> {
        let stream = TcpStream::connect(addr)?;
        stream.set_read_timeout(Some(Duration::from_secs(5)))?;
        stream.set_nodelay(true).ok();
        let mut reader = BufReader::new(stream.try_clone()?);

        // A fixed client key keeps the handshake assertion deterministic.
        let key = "dGhlIHNhbXBsZSBub25jZQ==";
        let mut s = stream.try_clone()?;
        let req = format!(
            "GET /ws HTTP/1.1\r\nHost: localhost\r\nUpgrade: websocket\r\n\
             Connection: Upgrade\r\nSec-WebSocket-Key: {key}\r\n\
             Sec-WebSocket-Version: 13\r\n\r\n"
        );
        s.write_all(req.as_bytes())?;
        s.flush()?;

        // Read the response head.
        let mut status = String::new();
        reader.read_line(&mut status)?;
        if !status.contains("101") {
            return Err(io::Error::new(io::ErrorKind::Other, format!("no upgrade: {status:?}")));
        }
        let mut saw_accept = false;
        loop {
            let mut line = String::new();
            let n = reader.read_line(&mut line)?;
            if n == 0 || line.trim_end().is_empty() {
                break;
            }
            if line.to_ascii_lowercase().starts_with("sec-websocket-accept:") {
                let got = line.split(':').nth(1).unwrap_or("").trim();
                // Expected accept for the canonical key above.
                if got == "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=" {
                    saw_accept = true;
                }
            }
        }
        if !saw_accept {
            return Err(io::Error::new(io::ErrorKind::Other, "bad Sec-WebSocket-Accept"));
        }
        Ok(WsClient { stream, reader })
    }

    /// Send a typed client message.
    pub fn send(&mut self, msg: &ClientMessage) -> io::Result<()> {
        self.send_raw_text(&msg.encode())
    }

    /// Send arbitrary text as a masked frame (used by fuzz tests).
    pub fn send_raw_text(&mut self, text: &str) -> io::Result<()> {
        self.write_masked(0x1, text.as_bytes())
    }

    /// Send arbitrary bytes as a masked binary frame (fuzz transport).
    pub fn send_raw_bytes(&mut self, bytes: &[u8]) -> io::Result<()> {
        self.write_masked(0x2, bytes)
    }

    fn write_masked(&mut self, opcode: u8, payload: &[u8]) -> io::Result<()> {
        let mask = [0x12u8, 0x34, 0x56, 0x78];
        let mut frame = vec![0x80 | opcode];
        let len = payload.len();
        if len < 126 {
            frame.push(0x80 | len as u8);
        } else if len < 65536 {
            frame.push(0x80 | 126);
            frame.extend_from_slice(&(len as u16).to_be_bytes());
        } else {
            frame.push(0x80 | 127);
            frame.extend_from_slice(&(len as u64).to_be_bytes());
        }
        frame.extend_from_slice(&mask);
        for (i, &b) in payload.iter().enumerate() {
            frame.push(b ^ mask[i % 4]);
        }
        self.stream.write_all(&frame)?;
        self.stream.flush()
    }

    /// Read one server text frame and decode it. Skips control frames.
    pub fn recv(&mut self) -> io::Result<ServerMessage> {
        let text = self.recv_text()?;
        ServerMessage::decode(&text)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("decode: {e}")))
    }

    /// Read one server text frame (raw).
    pub fn recv_text(&mut self) -> io::Result<String> {
        loop {
            let mut hdr = [0u8; 2];
            self.reader.read_exact(&mut hdr)?;
            let opcode = hdr[0] & 0x0f;
            let masked = hdr[1] & 0x80 != 0;
            let len7 = (hdr[1] & 0x7f) as usize;
            let len = match len7 {
                126 => {
                    let mut b = [0u8; 2];
                    self.reader.read_exact(&mut b)?;
                    u16::from_be_bytes(b) as usize
                }
                127 => {
                    let mut b = [0u8; 8];
                    self.reader.read_exact(&mut b)?;
                    u64::from_be_bytes(b) as usize
                }
                n => n,
            };
            // Server frames are not masked, but tolerate it.
            let mask = if masked {
                let mut m = [0u8; 4];
                self.reader.read_exact(&mut m)?;
                Some(m)
            } else {
                None
            };
            let mut payload = vec![0u8; len];
            self.reader.read_exact(&mut payload)?;
            if let Some(m) = mask {
                for (i, b) in payload.iter_mut().enumerate() {
                    *b ^= m[i % 4];
                }
            }
            match opcode {
                0x1 | 0x2 => return Ok(String::from_utf8_lossy(&payload).into_owned()),
                0x8 => return Err(io::Error::new(io::ErrorKind::ConnectionAborted, "closed")),
                _ => continue, // ping/pong/continuation
            }
        }
    }

    /// Receive until a predicate matches, returning that message. Bounded so a
    /// test never hangs.
    pub fn recv_until<F: Fn(&ServerMessage) -> bool>(&mut self, pred: F) -> io::Result<ServerMessage> {
        for _ in 0..256 {
            let m = self.recv()?;
            if pred(&m) {
                return Ok(m);
            }
        }
        Err(io::Error::new(io::ErrorKind::Other, "predicate never matched"))
    }
}

/// Start a fresh server instance on an ephemeral port for a test, returning the
/// address string. Each call binds a new port so tests are isolated.
pub fn spawn_test_server(turn_ms: u64) -> String {
    let addr = server::serve("127.0.0.1:0", server::Config { turn_ms })
        .expect("bind test server");
    format!("127.0.0.1:{}", addr.port())
}