boxmaker/crates/gatewayd/tests/ws_conn.rs

//! A WebSocket connection against a scripted server, plain and TLS: messages in order, pings both
//! ways, a dead peer, closing, and hostile input (M4a spec, section 6). Do not edit.

#[path = "support/tls_server.rs"]
mod tls_server;
#[path = "support/ws_server.rs"]
mod ws_server;

use std::io::Cursor;
use std::sync::mpsc;
use std::time::{Duration, Instant};

use gatewayd::config::ServerUrl;
use gatewayd::net::Connector;
use gatewayd::ws::WsError;
use gatewayd::ws::conn::{Timing, Ws, host_header};
use tls_server::{fixture, server_config};
use ws_server::{ClientFrame, serve_ws};

const SLOW: Timing = Timing {
    ping_every: Duration::from_secs(60),
    dead_after: Duration::from_secs(60),
};

/// Plenty of deterministic "random" bytes: the key, then masks.
fn random() -> Box<Cursor<Vec<u8>>> {
    Box::new(Cursor::new(
        (0..4096u32).map(|i| (i * 37 % 251) as u8).collect(),
    ))
}

fn open(port: u16, tls: bool, timing: Timing) -> Result<Ws, WsError> {
    let url = ServerUrl {
        tls,
        host: "localhost".to_string(),
        port,
    };
    let ca = tls.then(|| fixture("test-ca.pem"));
    let c = Connector::new(url, ca.as_deref()).unwrap();
    Ws::open(&c, "TOKEN", timing, random())
}

/// What the server saw, within 5 s: a missing frame fails the test instead of hanging it.
fn got<T>(rx: &mpsc::Receiver<T>) -> T {
    rx.recv_timeout(Duration::from_secs(5))
        .expect("the server saw nothing within 5 s")
}

/// Poll until a text message or an error, for at most 5 s.
fn next(ws: &mut Ws) -> Result<String, WsError> {
    let until = Instant::now() + Duration::from_secs(5);
    loop {
        if let Some(t) = ws.poll(Duration::from_millis(200))? {
            return Ok(t);
        }
        assert!(Instant::now() < until, "no message within 5 s");
    }
}

#[test]
fn messages_arrive_in_order_plain_and_over_tls() {
    for tls in [false, true] {
        let (tx, rx) = mpsc::channel();
        let config = tls.then(|| server_config("server"));
        let addr = serve_ws(config, None, move |mut p| {
            tx.send(p.request.clone()).unwrap();
            p.text("{\"event\":\"hello\"}");
            p.send(&[0x01, 0x03, b'o', b'n', b'e']);
            p.send(&[0x80, 0x04, b'-', b't', b'w', b'o']);
            p.text("three");
            p.pause(Duration::from_secs(2));
        });
        let mut ws = open(addr.port(), tls, SLOW).unwrap();
        let request = got(&rx);
        assert!(
            request.contains("Authorization: Bearer TOKEN\r\n"),
            "{request}"
        );
        assert!(
            request.starts_with("GET /api/v4/websocket HTTP/1.1\r\n"),
            "{request}"
        );
        assert!(
            request.contains(&format!("Host: localhost:{}\r\n", addr.port())),
            "{request}"
        );
        assert_eq!(next(&mut ws).unwrap(), "{\"event\":\"hello\"}");
        assert_eq!(next(&mut ws).unwrap(), "one-two");
        assert_eq!(next(&mut ws).unwrap(), "three");
    }
}

#[test]
fn a_ping_is_answered_with_the_same_payload() {
    let (tx, rx) = mpsc::channel();
    let addr = serve_ws(None, None, move |mut p| {
        p.frame(0x9, b"are you there");
        tx.send(p.read_frame()).unwrap();
        p.text("after");
        p.pause(Duration::from_secs(2));
    });
    let mut ws = open(addr.port(), false, SLOW).unwrap();
    assert_eq!(next(&mut ws).unwrap(), "after");
    let pong: ClientFrame = got(&rx).expect("a pong");
    assert_eq!(
        (pong.opcode, pong.masked, pong.payload.as_slice()),
        (0xA, true, &b"are you there"[..])
    );
}

#[test]
fn we_ping_on_schedule_and_every_frame_is_masked_differently() {
    let (tx, rx) = mpsc::channel();
    let addr = serve_ws(None, None, move |mut p| {
        for _ in 0..3 {
            tx.send(p.read_frame()).unwrap();
            p.frame(0xA, b"");
        }
        p.pause(Duration::from_secs(2));
    });
    let timing = Timing {
        ping_every: Duration::from_millis(100),
        dead_after: Duration::from_secs(5),
    };
    let mut ws = open(addr.port(), false, timing).unwrap();
    ws.send_text("first").unwrap();
    let started = Instant::now();
    while started.elapsed() < Duration::from_millis(350) {
        let _ = ws.poll(Duration::from_millis(50)).unwrap();
    }
    let frames: Vec<ClientFrame> = (0..3).map(|_| got(&rx).unwrap()).collect();
    assert_eq!(
        (frames[0].opcode, frames[0].payload.as_slice()),
        (0x1, &b"first"[..])
    );
    assert_eq!(frames[1].opcode, 0x9, "a ping after ping_every");
    assert_eq!(frames[2].opcode, 0x9);
    assert!(frames.iter().all(|f| f.masked));
    assert_ne!(
        frames[0].mask, frames[1].mask,
        "a fresh mask for every frame"
    );
}

#[test]
fn silence_is_a_dead_peer() {
    let addr = serve_ws(None, None, |p| p.pause(Duration::from_secs(10)));
    let timing = Timing {
        ping_every: Duration::from_secs(60),
        dead_after: Duration::from_millis(300),
    };
    let mut ws = open(addr.port(), false, timing).unwrap();
    let started = Instant::now();
    let err = loop {
        match ws.poll(Duration::from_millis(100)) {
            Ok(_) => assert!(
                started.elapsed() < Duration::from_secs(3),
                "never declared dead"
            ),
            Err(e) => break e,
        }
    };
    assert!(matches!(err, WsError::Dead), "{err}");
    assert!(started.elapsed() < Duration::from_secs(1));
}

#[test]
fn a_peer_that_trickles_is_alive_and_its_message_arrives() {
    let addr = serve_ws(None, None, |mut p| {
        for b in [0x81u8, 0x05, b'd', b'r', b'i', b'p', b's'] {
            p.send(&[b]);
            p.pause(Duration::from_millis(100));
        }
        p.pause(Duration::from_secs(2));
    });
    let timing = Timing {
        ping_every: Duration::from_secs(60),
        dead_after: Duration::from_millis(400),
    };
    let mut ws = open(addr.port(), false, timing).unwrap();
    assert_eq!(next(&mut ws).unwrap(), "drips");
}

#[test]
fn a_close_frame_is_answered_and_ends_the_connection() {
    let (tx, rx) = mpsc::channel();
    let addr = serve_ws(None, None, move |mut p| {
        p.frame(0x8, &[0x03, 0xE8]);
        tx.send(p.read_frame()).unwrap();
    });
    let mut ws = open(addr.port(), false, SLOW).unwrap();
    assert!(matches!(next(&mut ws), Err(WsError::Closed)));
    let reply = got(&rx).expect("a close in reply");
    assert_eq!(
        (reply.opcode, reply.payload.as_slice()),
        (0x8, &[0x03u8, 0xE8][..])
    );
}

#[test]
fn a_dropped_connection_is_closed() {
    let addr = serve_ws(None, None, drop);
    let mut ws = open(addr.port(), false, SLOW).unwrap();
    assert!(matches!(next(&mut ws), Err(WsError::Closed)));
}

#[test]
fn a_hostile_frame_is_an_error_not_a_panic() {
    let addr = serve_ws(None, None, |mut p| {
        p.send(&[0x81, 0xFF, 0x80, 0, 0, 0, 0, 0, 0, 0]);
        p.pause(Duration::from_secs(2));
    });
    let mut ws = open(addr.port(), false, SLOW).unwrap();
    assert!(matches!(next(&mut ws), Err(WsError::Protocol(_))));
}

#[test]
fn a_refused_handshake() {
    let addr = serve_ws(
        None,
        Some("HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n"),
        |_| {},
    );
    assert!(matches!(
        open(addr.port(), false, SLOW),
        Err(WsError::Handshake(_))
    ));
}

#[test]
fn the_host_header_names_the_port_only_when_it_is_not_the_default() {
    let url = |tls, port| ServerUrl {
        tls,
        host: "chat.example".to_string(),
        port,
    };
    assert_eq!(host_header(&url(true, 443)), "chat.example");
    assert_eq!(host_header(&url(false, 80)), "chat.example");
    assert_eq!(host_header(&url(true, 80)), "chat.example:80");
    assert_eq!(host_header(&url(false, 8065)), "chat.example:8065");
}