boxmaker/docs/plans/M4a/files/crates/gatewayd/tests/ws_handshake.rs

//! Base64 and the WebSocket opening handshake (RFC 4648; RFC 6455, section 4). Do not edit.

use std::io::{Cursor, Read, Write};

use gatewayd::http::Head;
use gatewayd::ws::WsError;
use gatewayd::ws::handshake::{
    accept_for, base64, check_response, handshake, new_key, request_text,
};

#[test]
fn base64_vectors() {
    for (input, want) in [
        ("", ""),
        ("f", "Zg=="),
        ("fo", "Zm8="),
        ("foo", "Zm9v"),
        ("foob", "Zm9vYg=="),
        ("fooba", "Zm9vYmE="),
        ("foobar", "Zm9vYmFy"),
    ] {
        assert_eq!(base64(input.as_bytes()), want, "{input:?}");
    }
    assert_eq!(base64(&[0xff, 0xfe, 0xfd, 0x00, 0x3f]), "//79AD8=");
    assert_eq!(
        base64(&(0u8..=15).collect::<Vec<_>>()),
        "AAECAwQFBgcICQoLDA0ODw=="
    );
}

#[test]
fn the_rfc_example_accept() {
    assert_eq!(
        accept_for("dGhlIHNhbXBsZSBub25jZQ=="),
        "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
    );
}

#[test]
fn a_key_is_sixteen_random_bytes() {
    let mut random = Cursor::new((0u8..=15).collect::<Vec<_>>());
    assert_eq!(new_key(&mut random).unwrap(), "AAECAwQFBgcICQoLDA0ODw==");
    let mut short = Cursor::new(vec![1u8; 15]);
    assert!(
        new_key(&mut short).is_err(),
        "too few random bytes is an error, not a weak key"
    );
}

#[test]
fn the_request_is_exactly_this() {
    assert_eq!(
        request_text("a.example", "/api/v4/websocket", "KEY==", "TOKEN"),
        "GET /api/v4/websocket HTTP/1.1\r\nHost: a.example\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
         Sec-WebSocket-Key: KEY==\r\nSec-WebSocket-Version: 13\r\nAuthorization: Bearer TOKEN\r\n\r\n"
    );
}

fn head(status: u16, headers: &[(&str, &str)]) -> Head {
    Head {
        status,
        headers: headers
            .iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect(),
    }
}

const KEY: &str = "dGhlIHNhbXBsZSBub25jZQ==";
const ACCEPT: &str = "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=";

#[test]
fn only_a_proper_upgrade_is_accepted() {
    let good = [
        ("Upgrade", "websocket"),
        ("Connection", "Upgrade"),
        ("Sec-WebSocket-Accept", ACCEPT),
    ];
    assert!(check_response(&head(101, &good), KEY).is_ok());
    let loose = [
        ("upgrade", "WebSocket"),
        ("connection", "keep-alive, Upgrade"),
        ("sec-websocket-accept", ACCEPT),
    ];
    assert!(
        check_response(&head(101, &loose), KEY).is_ok(),
        "names and tokens are case-insensitive"
    );
    let cases: [(u16, &[(&str, &str)]); 7] = [
        (200, &good),
        (401, &good),
        (
            101,
            &[("Connection", "Upgrade"), ("Sec-WebSocket-Accept", ACCEPT)],
        ),
        (
            101,
            &[
                ("Upgrade", "h2c"),
                ("Connection", "Upgrade"),
                ("Sec-WebSocket-Accept", ACCEPT),
            ],
        ),
        (
            101,
            &[("Upgrade", "websocket"), ("Sec-WebSocket-Accept", ACCEPT)],
        ),
        (101, &[("Upgrade", "websocket"), ("Connection", "Upgrade")]),
        (
            101,
            &[
                ("Upgrade", "websocket"),
                ("Connection", "Upgrade"),
                ("Sec-WebSocket-Accept", "s3pplmbitxaq9kygzzhzrbk+xoo="),
            ],
        ),
    ];
    for (status, headers) in cases {
        assert!(
            matches!(
                check_response(&head(status, headers), KEY),
                Err(WsError::Handshake(_))
            ),
            "{status} {headers:?}"
        );
    }
}

/// A server side scripted as bytes; records what the client wrote.
struct Scripted {
    input: Cursor<Vec<u8>>,
    output: Vec<u8>,
}

impl Read for Scripted {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.input.read(buf)
    }
}

impl Write for Scripted {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.output.write(buf)
    }
    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

#[test]
fn a_whole_handshake_leaves_the_first_frame_unread() {
    let key_bytes: Vec<u8> = (0u8..=15).collect();
    let key = base64(&key_bytes);
    let reply = format!(
        "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {}\r\n\r\n",
        accept_for(&key)
    );
    let mut bytes = reply.into_bytes();
    bytes.extend_from_slice(&[0x81, 0x02, b'h', b'i']);
    let mut s = Scripted {
        input: Cursor::new(bytes),
        output: Vec::new(),
    };
    handshake(
        &mut s,
        "a.example",
        "/api/v4/websocket",
        "TOKEN",
        &mut Cursor::new(key_bytes),
    )
    .unwrap();
    assert_eq!(
        String::from_utf8(s.output).unwrap(),
        request_text("a.example", "/api/v4/websocket", &key, "TOKEN")
    );
    let mut rest = Vec::new();
    s.input.read_to_end(&mut rest).unwrap();
    assert_eq!(rest, [0x81, 0x02, b'h', b'i']);
}

#[test]
fn a_refused_or_broken_handshake_is_a_handshake_error() {
    for reply in [
        &b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n"[..],
        b"HTTP/1.1 101 Switching",
        b"not http at all\r\n\r\n",
        b"",
    ] {
        let mut s = Scripted {
            input: Cursor::new(reply.to_vec()),
            output: Vec::new(),
        };
        let got = handshake(&mut s, "h", "/p", "t", &mut Cursor::new(vec![7u8; 16]));
        assert!(
            matches!(got, Err(WsError::Handshake(_))),
            "{:?}: {got:?}",
            String::from_utf8_lossy(reply)
        );
    }
}