deepslate_protocol/

codec.rs

//! Packet framing, compression, and high-level encode/decode functions.
//!
//! The Minecraft wire format is:
//! - Without compression: `[VarInt: packet_length][VarInt: packet_id][payload]`
//! - With compression (below threshold): `[VarInt: packet_length][0x00][VarInt: packet_id][payload]`
//! - With compression (above threshold): `[VarInt: packet_length][VarInt: uncompressed_size][zlib(packet_id + payload)]`

use bytes::{Buf, BufMut};

use crate::types::ProtocolError;
use crate::varint;

/// Maximum allowed frame size (2 MiB). Frames larger than this are rejected
/// to prevent memory exhaustion attacks.
const MAX_FRAME_SIZE: usize = 2 * 1024 * 1024;

/// Maximum allowed uncompressed packet size (8 MiB).
pub const MAX_UNCOMPRESSED_SIZE: usize = 8 * 1024 * 1024;

/// Try to extract a single frame from the given buffer.
///
/// Returns `Ok(Some((varint_size, frame_len)))` if a complete frame is
/// available, `Ok(None)` if more data is needed, or `Err` if the frame is
/// malformed.
///
/// The frame body starts at offset `varint_size` and spans `frame_len` bytes.
/// The total bytes consumed is `varint_size + frame_len`. This function does
/// **not** copy or allocate — callers should use `BytesMut::split_to()` or
/// equivalent zero-copy operations to extract the frame data.
///
/// # Errors
///
/// Returns `ProtocolError::FrameTooLarge` if the frame exceeds `MAX_FRAME_SIZE`.
/// Returns `ProtocolError::VarIntTooLong` if the length prefix is malformed.
#[allow(clippy::cast_sign_loss)]
pub fn try_read_frame(data: &[u8]) -> Result<Option<(usize, usize)>, ProtocolError> {
    let Some((frame_len, varint_size)) = varint::peek_var_int(data)? else {
        return Ok(None);
    };

    if frame_len < 0 {
        return Err(ProtocolError::InvalidData(
            "negative frame length".to_string(),
        ));
    }

    let frame_len = frame_len as usize;
    if frame_len > MAX_FRAME_SIZE {
        return Err(ProtocolError::FrameTooLarge {
            size: frame_len,
            max: MAX_FRAME_SIZE,
        });
    }

    let total = varint_size + frame_len;
    if data.len() < total {
        return Ok(None);
    }

    Ok(Some((varint_size, frame_len)))
}

/// Write a frame (length-prefixed) into the destination buffer.
/// `inner` is the already-encoded packet data (packet ID + payload).
#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
pub fn write_frame(dst: &mut impl BufMut, inner: &[u8]) {
    varint::write_var_int(dst, inner.len() as i32);
    dst.put_slice(inner);
}

/// Wrap packet data with the compression format.
///
/// If `data_len >= threshold`, the data should already be zlib-compressed and
/// `uncompressed_size` should be the original size. If below threshold, pass
/// `uncompressed_size = 0` and provide the raw data.
#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
pub fn write_compressed_frame(dst: &mut impl BufMut, uncompressed_size: i32, payload: &[u8]) {
    let size_varint_len = varint::var_int_bytes(uncompressed_size);
    let frame_len = size_varint_len + payload.len();
    varint::write_var_int(dst, frame_len as i32);
    varint::write_var_int(dst, uncompressed_size);
    dst.put_slice(payload);
}

/// Decode the inner frame data when compression is enabled.
///
/// Returns `(uncompressed_size, payload)` where:
/// - `uncompressed_size == 0` means the payload is NOT compressed
/// - `uncompressed_size > 0` means the payload is zlib-compressed
///
/// The returned payload is a sub-slice of the input — no allocation or copy.
///
/// # Errors
///
/// Returns errors for malformed data.
#[allow(clippy::cast_sign_loss)]
pub fn read_compressed_frame(data: &[u8]) -> Result<(usize, &[u8]), ProtocolError> {
    let mut cursor = data;
    let uncompressed_size = varint::read_var_int(&mut cursor)?;
    if uncompressed_size < 0 {
        return Err(ProtocolError::InvalidData(
            "negative uncompressed size".to_string(),
        ));
    }
    let uncompressed_size = uncompressed_size as usize;
    if uncompressed_size > MAX_UNCOMPRESSED_SIZE {
        return Err(ProtocolError::UncompressedSizeTooLarge {
            size: uncompressed_size,
            max: MAX_UNCOMPRESSED_SIZE,
        });
    }
    Ok((uncompressed_size, cursor))
}

/// Encode a typed packet into raw bytes (packet ID + fields), without framing.
///
/// Clears `buf` and writes the packet ID followed by the encoded fields into
/// it. The caller should supply a reusable buffer so that its capacity is
/// retained across calls, avoiding per-packet allocations.
pub fn encode_packet_data(buf: &mut Vec<u8>, packet_id: i32, encode_fn: impl FnOnce(&mut Vec<u8>)) {
    buf.clear();
    varint::write_var_int(buf, packet_id);
    encode_fn(buf);
}

/// Read a packet ID from frame data.
///
/// # Errors
///
/// Returns `ProtocolError` if the packet ID `VarInt` is malformed.
pub fn read_packet_id(buf: &mut impl Buf) -> Result<i32, ProtocolError> {
    varint::read_var_int(buf)
}

#[cfg(test)]
mod tests {
    use super::*;
    use proptest::prelude::*;

    proptest! {
        #[test]
        fn write_frame_roundtrip_prop(payload in prop::collection::vec(any::<u8>(), 0..4096)) {
            let mut buf = Vec::new();
            write_frame(&mut buf, &payload);
            let (varint_size, frame_len) = try_read_frame(&buf).unwrap().unwrap();
            prop_assert_eq!(varint_size + frame_len, buf.len());
            prop_assert_eq!(&buf[varint_size..varint_size + frame_len], &payload[..]);
        }

        #[test]
        fn compressed_frame_roundtrip_prop(
            uncompressed_size in 0..MAX_UNCOMPRESSED_SIZE,
            payload in prop::collection::vec(any::<u8>(), 0..4096)
        ) {
            #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
            let uncompressed_size_i32 = uncompressed_size as i32;
            let mut buf = Vec::new();
            write_compressed_frame(&mut buf, uncompressed_size_i32, &payload);
            let (varint_size, frame_len) = try_read_frame(&buf).unwrap().unwrap();
            let frame_data = &buf[varint_size..varint_size + frame_len];
            let (read_uncompressed_size, data) = read_compressed_frame(frame_data).unwrap();
            prop_assert_eq!(read_uncompressed_size, uncompressed_size);
            prop_assert_eq!(data, &payload[..]);
        }
    }

    #[test]
    fn test_try_read_frame_complete() {
        // Frame: length=3, data=[0x01, 0x02, 0x03]
        let data = vec![0x03, 0x01, 0x02, 0x03];
        let (varint_size, frame_len) = try_read_frame(&data).unwrap().unwrap();
        assert_eq!(varint_size + frame_len, 4);
        assert_eq!(
            &data[varint_size..varint_size + frame_len],
            &[0x01, 0x02, 0x03]
        );
    }

    #[test]
    fn test_try_read_frame_incomplete() {
        let data = vec![0x03, 0x01]; // Says 3 bytes but only 1 available
        assert!(try_read_frame(&data).unwrap().is_none());
    }

    #[test]
    fn test_try_read_frame_empty() {
        assert!(try_read_frame(&[]).unwrap().is_none());
    }

    #[test]
    fn test_write_frame_roundtrip() {
        let inner = vec![0x00, 0x48, 0x65, 0x6C, 0x6C, 0x6F]; // packet_id=0, "Hello"
        let mut buf = Vec::new();
        write_frame(&mut buf, &inner);
        let (varint_size, frame_len) = try_read_frame(&buf).unwrap().unwrap();
        assert_eq!(varint_size + frame_len, buf.len());
        assert_eq!(&buf[varint_size..varint_size + frame_len], &inner[..]);
    }

    #[test]
    fn test_compressed_frame_uncompressed() {
        let mut buf = Vec::new();
        let payload = vec![0x00, 0x48, 0x69];
        write_compressed_frame(&mut buf, 0, &payload);

        let (varint_size, frame_len) = try_read_frame(&buf).unwrap().unwrap();
        let frame_data = &buf[varint_size..varint_size + frame_len];
        let (uncompressed_size, data) = read_compressed_frame(frame_data).unwrap();
        assert_eq!(uncompressed_size, 0);
        assert_eq!(data, &payload[..]);
    }
}