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

//! Wire-level primitives of the MOQT draft this crate implements, over `bytes`
//! buffers.
//!
//! MOQT's integer is the variable-length integer of
//! `draft-ietf-moq-transport-22` §8.1 ([`crate::draft22::varint`]): the number of
//! leading 1 bits of the first byte gives the length, 1 to 9 bytes, and every
//! 64-bit value is representable. It is not QUIC's (RFC 9000 §16), which drafts
//! before -17 used and this module carried before. Byte strings are prefixed by
//! their length as such an integer.

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

/// The largest value the integer carries: every `u64`.
pub const VARINT_MAX: u64 = varint::VARINT_MAX;

/// Encode an unsigned integer into `dst`, in its shortest form.
pub fn encode_varint<B: BufMut>(value: u64, dst: &mut B) -> MoqResult<()> {
    let mut tmp = Vec::with_capacity(9);
    varint::put(&mut tmp, value);
    dst.put_slice(&tmp);
    Ok(())
}

/// Decode an integer from `src`, advancing 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 len = varint::len_from_first(src.chunk()[0]);
    if src.remaining() < len {
        return Err(MoqError::WireDecode(format!("varint: need {} bytes, have {}", len, src.remaining())));
    }
    let mut tmp = [0u8; 9];
    src.copy_to_slice(&mut tmp[..len]);
    let (v, _) = varint::get(&tmp[..len], "varint")?;
    Ok(v)
}

/// Encode a byte slice prefixed by its 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 length-prefixed byte vector.
pub fn decode_bytes<B: Buf>(src: &mut B) -> MoqResult<Bytes> {
    let len = decode_varint(src)?;
    if (src.remaining() as u64) < len {
        return Err(MoqError::WireDecode(format!("bytes: need {} bytes, have {}", len, src.remaining())));
    }
    let len = len as usize;
    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>;
}

impl WireCodec for crate::draft22::ControlMessage {
    fn encode(&self, dst: &mut BytesMut) -> MoqResult<()> {
        dst.put_slice(&crate::draft22::ControlMessage::encode(self)?);
        Ok(())
    }

    fn decode(src: &mut Bytes) -> MoqResult<Self> {
        let mut r = crate::draft22::io::Reader::new(src);
        let m = crate::draft22::ControlMessage::decode(&mut r)?;
        let used = r.position();
        src.advance(used);
        Ok(m)
    }
}

#[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, 127, 128, 16_383, 16_384, 1_073_741_823, 1_073_741_824, (1 << 56) - 1, 1 << 56, VARINT_MAX] {
            roundtrip(v);
        }
    }

    /// Before -17 the integer stopped at 2^62-1 and a larger value was an error.
    /// In -22 every u64 has a form, so there is no overflow to refuse.
    #[test]
    fn every_u64_encodes() {
        for v in [(1u64 << 62) - 1, 1 << 62, u64::MAX] {
            roundtrip(v);
        }
        let mut b = BytesMut::new();
        encode_varint(u64::MAX, &mut b).unwrap();
        assert_eq!(b.len(), 9);
    }

    #[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");
    }
}