boxmaker/docs/plans/M1/files/crates/proto/tests/frame.rs

//! Tests for the frame codec. Do not edit these or the fixtures.

use proto::{
    CallId, Envelope, FrameError, MAX_FRAME, Message, SessionId, ToolRequest, read_frame,
    write_frame,
};
use std::io::{Cursor, Read};

fn fixture_bytes(name: &str) -> Vec<u8> {
    let path = format!("{}/tests/fixtures/frame/{name}", env!("CARGO_MANIFEST_DIR"));
    std::fs::read(&path).unwrap_or_else(|e| panic!("{path}: {e}"))
}

fn request(arguments: String) -> Envelope {
    let body = ToolRequest {
        session: SessionId::new("mm-thread-42").unwrap(),
        call: CallId(3),
        tool: "read_file".to_string(),
        arguments,
    };
    Envelope {
        v: 1,
        id: 7,
        r#final: true,
        msg: Message::ToolRequest(body),
    }
}

/// Serves `head`, then panics if anyone reads further.
struct HeaderOnly {
    head: Cursor<Vec<u8>>,
}

impl Read for HeaderOnly {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let n = self.head.read(buf)?;
        assert!(
            n > 0,
            "the reader asked for the body of a frame it should have rejected"
        );
        Ok(n)
    }
}

/// Hands out one byte per call, to catch codecs that assume `read` fills the buffer.
struct OneByteAtATime(Cursor<Vec<u8>>);

impl Read for OneByteAtATime {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let end = buf.len().min(1);
        self.0.read(&mut buf[..end])
    }
}

#[test]
fn max_frame_is_one_mebibyte() {
    assert_eq!(MAX_FRAME, 1_048_576);
}

#[test]
fn write_matches_the_fixture_byte_for_byte() {
    let mut out = Vec::new();
    write_frame(&mut out, &request(r#"{"path":"/etc/hosts"}"#.to_string())).unwrap();
    assert_eq!(out, fixture_bytes("tool_request.bin"));
    assert_eq!(&out[..4], &[0, 0, 0, 159]);
}

#[test]
fn read_decodes_the_fixture() {
    let mut input = Cursor::new(fixture_bytes("tool_request.bin"));
    let env = read_frame(&mut input).unwrap();
    assert_eq!(env, request(r#"{"path":"/etc/hosts"}"#.to_string()));
}

#[test]
fn two_frames_in_a_row_then_closed() {
    let mut bytes = fixture_bytes("tool_request.bin");
    bytes.extend(fixture_bytes("tool_request.bin"));
    let mut input = Cursor::new(bytes);
    assert!(read_frame(&mut input).is_ok());
    assert!(read_frame(&mut input).is_ok());
    assert!(matches!(read_frame(&mut input), Err(FrameError::Closed)));
}

#[test]
fn short_reads_are_handled() {
    let mut input = OneByteAtATime(Cursor::new(fixture_bytes("tool_request.bin")));
    assert!(read_frame(&mut input).is_ok());
}

#[test]
fn zero_length_is_empty() {
    let mut input = Cursor::new(vec![0, 0, 0, 0]);
    assert!(matches!(read_frame(&mut input), Err(FrameError::Empty)));
}

#[test]
fn oversize_length_is_rejected_without_reading_the_body() {
    let n = (MAX_FRAME as u32) + 1;
    let mut input = HeaderOnly {
        head: Cursor::new(n.to_be_bytes().to_vec()),
    };
    match read_frame(&mut input) {
        Err(FrameError::TooLarge(got)) => assert_eq!(got, MAX_FRAME + 1),
        other => panic!("expected TooLarge, got {other:?}"),
    }
    let mut input = HeaderOnly {
        head: Cursor::new(vec![0xff, 0xff, 0xff, 0xff]),
    };
    assert!(matches!(
        read_frame(&mut input),
        Err(FrameError::TooLarge(4_294_967_295))
    ));
}

#[test]
fn exactly_max_frame_is_accepted_by_the_length_check() {
    // A body of MAX_FRAME bytes of spaces is not valid JSON, so the error must be Json, not TooLarge.
    let mut bytes = (MAX_FRAME as u32).to_be_bytes().to_vec();
    bytes.extend(std::iter::repeat_n(b' ', MAX_FRAME));
    assert!(matches!(
        read_frame(&mut Cursor::new(bytes)),
        Err(FrameError::Json(_))
    ));
}

#[test]
fn truncated_header_and_body_are_io_errors() {
    let full = fixture_bytes("tool_request.bin");
    let mut header_cut = Cursor::new(full[..2].to_vec());
    assert!(matches!(
        read_frame(&mut header_cut),
        Err(FrameError::Io(_))
    ));
    let mut body_cut = Cursor::new(full[..full.len() - 1].to_vec());
    assert!(matches!(read_frame(&mut body_cut), Err(FrameError::Io(_))));
}

#[test]
fn garbage_body_is_a_json_error() {
    let mut bytes = vec![0, 0, 0, 3];
    bytes.extend(b"{{{");
    assert!(matches!(
        read_frame(&mut Cursor::new(bytes)),
        Err(FrameError::Json(_))
    ));
}

#[test]
fn other_protocol_versions_are_rejected() {
    let mut env = request("{}".to_string());
    env.v = 2;
    let mut out = Vec::new();
    write_frame(&mut out, &env).unwrap();
    assert!(matches!(
        read_frame(&mut Cursor::new(out)),
        Err(FrameError::BadVersion(2))
    ));
}

#[test]
fn oversize_envelopes_are_not_written() {
    let mut out = Vec::new();
    let err = write_frame(&mut out, &request("x".repeat(MAX_FRAME))).unwrap_err();
    assert!(matches!(err, FrameError::TooLarge(n) if n > MAX_FRAME));
    assert!(
        out.is_empty(),
        "nothing may be written when the envelope is too large"
    );
}

#[test]
fn frame_error_is_a_std_error_with_a_message() {
    let e: Box<dyn std::error::Error> = Box::new(FrameError::Empty);
    assert!(!e.to_string().is_empty());
}