open-standards/joule-code/joule-code-rs/crates/joule-moq/src/wire.rs

//! Wire-level codec primitives for MoQT.
//!
//! MoQT uses QUIC-style variable-length integers (RFC 9000 §16): the two most
//! significant bits of the first byte encode the integer length (1, 2, 4, or
//! 8 bytes), the remaining bits hold the value. Strings and byte vectors are
//! length-prefixed by varint.

use crate::error::{MoqError, MoqResult};
use bytes::{Buf, BufMut, Bytes, BytesMut};

/// Maximum value representable as a QUIC varint.
pub const VARINT_MAX: u64 = (1 << 62) - 1;

/// Encode an unsigned integer as a QUIC varint into `dst`.
pub fn encode_varint<B: BufMut>(value: u64, dst: &mut B) -> MoqResult<()> {
    if value > VARINT_MAX {
        return Err(MoqError::VarintOverflow);
    }
    if value < (1 << 6) {
        dst.put_u8(value as u8);
    } else if value < (1 << 14) {
        dst.put_u16((value as u16) | (0b01 << 14));
    } else if value < (1 << 30) {
        dst.put_u32((value as u32) | (0b10 << 30));
    } else {
        dst.put_u64(value | (0b11u64 << 62));
    }
    Ok(())
}

/// Decode a QUIC varint from `src`. Advances the cursor.
pub fn decode_varint<B: Buf>(src: &mut B) -> MoqResult<u64> {
    if !src.has_remaining() {
        return Err(MoqError::WireDecode("varint: empty buffer".into()));
    }
    let first = src.chunk()[0];
    let len_bits = first >> 6;
    let len = 1usize << len_bits;
    if src.remaining() < len {
        return Err(MoqError::WireDecode(format!(
            "varint: need {} bytes, have {}",
            len,
            src.remaining()
        )));
    }
    let value = match len {
        1 => (src.get_u8() & 0x3f) as u64,
        2 => (src.get_u16() & 0x3fff) as u64,
        4 => (src.get_u32() & 0x3fff_ffff) as u64,
        8 => src.get_u64() & 0x3fff_ffff_ffff_ffff,
        _ => unreachable!("len_bits is 2 bits"),
    };
    Ok(value)
}

/// Encode a byte slice prefixed by its varint length.
pub fn encode_bytes<B: BufMut>(slice: &[u8], dst: &mut B) -> MoqResult<()> {
    encode_varint(slice.len() as u64, dst)?;
    dst.put_slice(slice);
    Ok(())
}

/// Decode a varint-length-prefixed byte vector.
pub fn decode_bytes<B: Buf>(src: &mut B) -> MoqResult<Bytes> {
    let len = decode_varint(src)? as usize;
    if src.remaining() < len {
        return Err(MoqError::WireDecode(format!(
            "bytes: need {} bytes, have {}",
            len,
            src.remaining()
        )));
    }
    let mut buf = BytesMut::with_capacity(len);
    buf.resize(len, 0);
    src.copy_to_slice(&mut buf);
    Ok(buf.freeze())
}

/// A trait for types that can be serialized to / deserialized from the
/// MoQT wire format.
pub trait WireCodec: Sized {
    fn encode(&self, dst: &mut BytesMut) -> MoqResult<()>;
    fn decode(src: &mut Bytes) -> MoqResult<Self>;
}

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

    fn roundtrip(value: u64) {
        let mut buf = BytesMut::new();
        encode_varint(value, &mut buf).unwrap();
        let mut bytes = buf.freeze();
        let decoded = decode_varint(&mut bytes).unwrap();
        assert_eq!(decoded, value);
        assert_eq!(bytes.remaining(), 0);
    }

    #[test]
    fn varint_boundaries() {
        for v in [0u64, 1, 63, 64, 16_383, 16_384, 1_073_741_823, 1_073_741_824, VARINT_MAX] {
            roundtrip(v);
        }
    }

    #[test]
    fn varint_overflow() {
        let mut buf = BytesMut::new();
        let result = encode_varint(VARINT_MAX + 1, &mut buf);
        assert!(matches!(result, Err(MoqError::VarintOverflow)));
    }

    #[test]
    fn bytes_roundtrip() {
        let mut buf = BytesMut::new();
        encode_bytes(b"hello", &mut buf).unwrap();
        let mut bytes = buf.freeze();
        let decoded = decode_bytes(&mut bytes).unwrap();
        assert_eq!(&decoded[..], b"hello");
    }
}