travsr-plugin-protocol 0.7.0

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

/// Hard cap on a single frame's payload (256 MiB). A plugin is an untrusted
/// peer over the wire: the 4-byte length prefix is attacker-controlled and a
/// hostile/buggy plugin could send `0xFFFF_FFFF` to make the daemon allocate
/// 4 GiB and OOM. We refuse any frame larger than this before allocating.
/// 256 MiB comfortably exceeds the largest realistic ParseOutput/InvokeResponse.
pub const MAX_FRAME_LEN: usize = 256 * 1024 * 1024;

pub fn encode_message<T: serde::Serialize>(msg: &T) -> io::Result<Vec<u8>> {
    let payload =
        serde_json::to_vec(msg).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
    if payload.len() > MAX_FRAME_LEN {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "frame payload {} bytes exceeds MAX_FRAME_LEN ({MAX_FRAME_LEN} bytes)",
                payload.len()
            ),
        ));
    }
    let len = u32::try_from(payload.len())
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "payload too large"))?;
    let mut frame = Vec::with_capacity(4 + payload.len());
    frame.extend_from_slice(&len.to_be_bytes());
    frame.extend_from_slice(&payload);
    Ok(frame)
}

pub fn decode_message<T: serde::de::DeserializeOwned>(reader: &mut impl Read) -> io::Result<T> {
    let mut len_buf = [0u8; 4];
    reader.read_exact(&mut len_buf)?;
    let len = u32::from_be_bytes(len_buf) as usize;
    // Reject oversized frames BEFORE allocating — prevents a hostile plugin from
    // triggering a multi-gigabyte allocation via a forged length prefix (DoS).
    if len > MAX_FRAME_LEN {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("frame length {len} exceeds MAX_FRAME_LEN ({MAX_FRAME_LEN} bytes)"),
        ));
    }
    let mut payload = vec![0u8; len];
    reader.read_exact(&mut payload)?;
    serde_json::from_slice(&payload).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
}

pub fn write_message<T: serde::Serialize>(writer: &mut impl Write, msg: &T) -> io::Result<()> {
    let frame = encode_message(msg)?;
    writer.write_all(&frame)?;
    writer.flush()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::ParseRequest;
    use std::io::Cursor;
    use std::path::PathBuf;

    #[test]
    fn round_trip_parse_request() {
        let req = ParseRequest {
            path: PathBuf::from("src/main.ts"),
            vname_path: "src/main.ts".into(),
            corpus: "github.com/acme/foo".into(),
            package: "acme".into(),
            source: None,
        };
        let payload = serde_json::to_vec(&req).unwrap();
        let encoded = encode_message(&req).unwrap();
        // Frame must be exactly 4-byte length prefix + payload.
        assert_eq!(encoded.len(), 4 + payload.len());
        // First 4 bytes must be the big-endian payload length.
        let len = u32::from_be_bytes(encoded[..4].try_into().unwrap()) as usize;
        assert_eq!(len, payload.len());
        let mut cursor = Cursor::new(&encoded);
        let decoded: ParseRequest = decode_message(&mut cursor).unwrap();
        assert_eq!(decoded.path, req.path);
        assert_eq!(decoded.vname_path, req.vname_path);
        assert_eq!(decoded.corpus, req.corpus);
    }

    #[test]
    fn decode_rejects_oversized_frame_without_allocating() {
        // Forge a length prefix of u32::MAX (~4 GiB) with no payload behind it.
        // decode_message must reject it on the length check, never attempting the
        // multi-gigabyte allocation, and never blocking on read_exact for a body.
        let mut framed = (u32::MAX).to_be_bytes().to_vec();
        framed.extend_from_slice(b"only a few bytes follow");
        let mut cursor = Cursor::new(framed);
        let err = decode_message::<ParseRequest>(&mut cursor).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
        assert!(
            err.to_string().contains("MAX_FRAME_LEN"),
            "expected MAX_FRAME_LEN rejection, got: {err}"
        );
    }

    #[test]
    fn decode_accepts_frame_at_the_boundary() {
        // A payload exactly at the cap must still decode (it is a valid JSON string).
        let big = "x".repeat(1024);
        let encoded = encode_message(&big).unwrap();
        let mut cursor = Cursor::new(encoded);
        let decoded: String = decode_message(&mut cursor).unwrap();
        assert_eq!(decoded, big);
    }
}