liminal-server 0.24.9

//! The pre-admission inbound bound on the WebSocket transport: before the
//! connection's `Connect` succeeds, a frame or reassembled message larger than
//! the bound is refused by name with at most the bound held, and once the
//! `Connect` succeeds the acceptor's pinned bound applies.

use std::net::TcpStream;
use std::time::Duration;

use liminal::protocol::{Frame as LiminalFrame, ProtocolVersion, encode, encoded_len};
use tungstenite::Message;
use tungstenite::protocol::frame::coding::{CloseCode, Data, OpCode};
use tungstenite::protocol::frame::{CloseFrame, Frame};
use tungstenite::protocol::{Role, WebSocket};

use super::super::super::pre_admission::{PRE_ADMISSION_FRAME_BOUND, PreAdmissionOversize};
use super::SliceStep;
use super::tests::scheduler_free_process;

/// A guard on every blocking read: it fires only when the signal never comes.
const GUARD: Duration = Duration::from_secs(60);

/// One fragment's payload. Four fragments fill the bound exactly; the fifth
/// is the first byte past it.
const FRAGMENT: usize = PRE_ADMISSION_FRAME_BOUND / 4;

fn client_socket(stream: TcpStream) -> Result<WebSocket<TcpStream>, String> {
    stream
        .set_read_timeout(Some(GUARD))
        .map_err(|error| format!("client guard: {error}"))?;
    Ok(WebSocket::from_raw_socket(stream, Role::Client, None))
}

/// Writes `payloads` as one binary message split into one frame per payload.
fn write_message(
    mut client: WebSocket<TcpStream>,
    payloads: &[usize],
) -> Result<WebSocket<TcpStream>, String> {
    for (index, length) in payloads.iter().enumerate() {
        let opcode = if index == 0 {
            OpCode::Data(Data::Binary)
        } else {
            OpCode::Data(Data::Continue)
        };
        let is_final = index + 1 == payloads.len();
        client
            .write(Message::Frame(Frame::message(
                vec![0_u8; *length],
                opcode,
                is_final,
            )))
            .map_err(|error| format!("write fragment {index}: {error}"))?;
    }
    client.flush().map_err(|error| format!("flush: {error}"))?;
    Ok(client)
}

/// Reads until the server's close frame and returns it.
fn read_close(client: &mut WebSocket<TcpStream>) -> Result<CloseFrame, String> {
    loop {
        match client.read() {
            Ok(Message::Close(Some(frame))) => return Ok(frame),
            Ok(Message::Close(None)) => return Err("close frame carried no code".to_owned()),
            Ok(_) => {}
            Err(error) => return Err(format!("client read before close: {error}")),
        }
    }
}

/// Sends `payloads` as one message to a fresh unadmitted connection, services
/// the connection until it stops, and returns the close frame it sent.
fn refuse(payloads: &[usize]) -> Result<CloseFrame, String> {
    let (mut process, client) = scheduler_free_process()?;
    let server_stream = process
        .socket
        .as_ref()
        .ok_or("process has no socket")?
        .get_ref();
    server_stream
        .set_read_timeout(Some(GUARD))
        .map_err(|error| format!("server guard: {error}"))?;
    let client = client_socket(client)?;
    let (step, client) = std::thread::scope(|scope| {
        let writer = scope.spawn(|| write_message(client, payloads));
        let step = process.service_socket(1);
        (step, writer.join())
    });
    let mut client = client.map_err(|_| "writer thread panicked".to_owned())??;
    if !matches!(step, SliceStep::Stop(_)) {
        return Err("the oversize message did not stop the connection".to_owned());
    }
    read_close(&mut client)
}

#[test]
fn fragmented_message_over_the_bound_before_connect_is_refused_by_name() -> Result<(), String> {
    // Five fragments cross the bound; a sixth, final fragment would complete
    // the message if the bound did not stop it first.
    let close = refuse(&[FRAGMENT; 6])?;

    // The server's own figure in the close reason is the size the fifth
    // fragment would have made: the refusal fired on the first fragment past
    // the bound, before the sixth was read. The bytes the transport library
    // holds are not observable from outside it.
    let expected = PreAdmissionOversize {
        attempted: FRAGMENT * 5,
        bound: PRE_ADMISSION_FRAME_BOUND,
    };
    assert_eq!(close.code, CloseCode::Size, "{close:?}");
    assert_eq!(close.reason.as_str(), expected.to_string());
    Ok(())
}

#[test]
fn single_frame_declaring_more_than_the_bound_before_connect_is_refused_from_its_header()
-> Result<(), String> {
    let close = refuse(&[PRE_ADMISSION_FRAME_BOUND + 1])?;

    let expected = PreAdmissionOversize {
        attempted: PRE_ADMISSION_FRAME_BOUND + 1,
        bound: PRE_ADMISSION_FRAME_BOUND,
    };
    assert_eq!(close.code, CloseCode::Size, "{close:?}");
    assert_eq!(close.reason.as_str(), expected.to_string());
    Ok(())
}

#[test]
fn a_successful_connect_installs_the_acceptor_bound() -> Result<(), String> {
    let (mut process, _client) = scheduler_free_process()?;
    let limit = |process: &super::WebSocketConnectionProcess| {
        process
            .socket
            .as_ref()
            .map(|socket| socket.get_config().max_message_size)
    };
    assert_eq!(limit(&process), Some(Some(PRE_ADMISSION_FRAME_BOUND)));

    let connect = LiminalFrame::Connect {
        flags: 0,
        min_version: ProtocolVersion::new(1, 0),
        max_version: ProtocolVersion::new(1, 0),
        auth_token: Vec::new(),
    };
    let mut bytes = vec![0_u8; encoded_len(&connect).map_err(|error| error.to_string())?];
    let written = encode(&connect, &mut bytes).map_err(|error| error.to_string())?;
    bytes.truncate(written);
    let step = process.handle_inbound_message(1, Message::Binary(bytes.into()));

    assert!(step.is_none(), "the connect was accepted");
    assert!(process.state.authenticated);
    assert_eq!(limit(&process), Some(Some(process.admitted_bound)));
    Ok(())
}

#[test]
fn both_transports_refuse_with_the_same_name() -> Result<(), String> {
    let refusal = PreAdmissionOversize {
        attempted: PRE_ADMISSION_FRAME_BOUND + 1,
        bound: PRE_ADMISSION_FRAME_BOUND,
    };
    let close = super::pre_admission_close_frame(refusal);
    let LiminalFrame::ConnectError {
        message: Some(tcp_message),
        ..
    } = super::super::super::pre_admission::connect_error_frame(refusal)
    else {
        return Err("the byte-stream refusal is a ConnectError carrying a message".to_owned());
    };

    assert_eq!(close.code, CloseCode::Size);
    assert_eq!(close.reason.as_str(), tcp_message);
    assert!(tcp_message.starts_with(super::super::super::pre_admission::REFUSAL_NAME));
    assert!(tcp_message.starts_with("FrameTooLarge"));
    Ok(())
}