trillium-http 1.1.0

/// Errors that can occur during variable-length integer decoding.
#[derive(Debug, thiserror::Error, PartialEq, Eq, Clone, Copy)]
pub enum IntegerPrefixError {
    /// The encoded integer overflows usize.
    #[error("encoded integer overflows usize")]
    Overflow,

    /// The input ended before the integer was fully decoded (more
    /// continuation bytes were expected).
    #[error("unexpected end of input")]
    UnexpectedEnd,
}

/// Decode a prefix-coded integer per RFC 7541 §5.1.
///
/// `prefix_size` is the number of bits in the first byte used for the
/// integer (the low bits). The remaining high bits belong to the
/// caller (flags, type indicators, etc.).
///
/// Returns the decoded value and the unconsumed remainder of the input.
pub(in crate::headers) fn decode(
    input: &[u8],
    prefix_size: u8,
) -> Result<(usize, &[u8]), IntegerPrefixError> {
    debug_assert!((1..=8).contains(&prefix_size));

    let [first, rest @ ..] = input else {
        return Err(IntegerPrefixError::UnexpectedEnd);
    };

    let prefix_mask = u8::MAX >> (8 - prefix_size);
    let mut value = usize::from(first & prefix_mask);

    if value < usize::from(prefix_mask) {
        return Ok((value, rest));
    }

    // Value continues in subsequent bytes
    let mut shift = 0_u32;
    for (i, &byte) in rest.iter().enumerate() {
        let payload = usize::from(byte & 0x7F);
        let increment = payload
            .checked_shl(shift)
            .ok_or(IntegerPrefixError::Overflow)?;
        value = value
            .checked_add(increment)
            .ok_or(IntegerPrefixError::Overflow)?;
        shift += 7;

        if byte & 0x80 == 0 {
            return Ok((value, &rest[i + 1..]));
        }
    }

    Err(IntegerPrefixError::UnexpectedEnd)
}

/// Encoded length of `value` under a prefix-coded integer with `prefix_size`-bit prefix
/// (RFC 7541 §5.1), in bytes. Used to project wire sizes without materializing the
/// encoding.
pub(in crate::headers) fn encoded_length(value: usize, prefix_size: u8) -> usize {
    debug_assert!((1..=8).contains(&prefix_size));

    let prefix_max = u8::MAX >> (8 - prefix_size);

    if value < usize::from(prefix_max) {
        return 1;
    }

    // Prefix byte plus `ceil(bits_needed / 7)` continuation bytes (each carrying 7 value
    // bits), where `bits_needed` is the bit-width of `remaining = value - prefix_max`. A
    // `remaining` of 0 still needs one terminating byte, so round-up with a clamp to 1.
    let remaining = value - usize::from(prefix_max);
    let bits_needed = usize::BITS - remaining.leading_zeros();
    let continuation = (bits_needed.div_ceil(7)).max(1) as usize;
    1 + continuation
}

/// Encode a prefix-coded integer per RFC 7541 §5.1, appending the encoded bytes to `buf`.
///
/// The first byte appended contains the integer value in its low `prefix_size` bits; the
/// caller is responsible for OR-ing in any flags in the high bits.
#[allow(clippy::cast_possible_truncation)] // all casts are bounded: value < prefix_max <= 255, remaining < 128
pub(in crate::headers) fn encode_into(value: usize, prefix_size: u8, buf: &mut Vec<u8>) {
    debug_assert!((1..=8).contains(&prefix_size));

    let prefix_max = u8::MAX >> (8 - prefix_size);

    if value < usize::from(prefix_max) {
        buf.push(value as u8);
        return;
    }

    buf.push(prefix_max);
    let mut remaining = value - usize::from(prefix_max);

    while remaining >= 128 {
        buf.push(remaining as u8 | 0x80);
        remaining >>= 7;
    }

    buf.push(remaining as u8);
}

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

    // RFC 7541 C.1.1: encoding 10 with 5-bit prefix
    #[test]
    fn decode_10_prefix_5() {
        let (value, rest) = decode(&[0b000_01010], 5).unwrap();
        assert_eq!(value, 10);
        assert!(rest.is_empty());
    }

    // RFC 7541 C.1.2: encoding 1337 with 5-bit prefix
    #[test]
    fn decode_1337_prefix_5() {
        let (value, rest) = decode(&[0b000_11111, 0b10011010, 0b00001010], 5).unwrap();
        assert_eq!(value, 1337);
        assert!(rest.is_empty());
    }

    // RFC 7541 C.1.3: encoding 42 with 8-bit prefix
    #[test]
    fn decode_42_prefix_8() {
        let (value, rest) = decode(&[42], 8).unwrap();
        assert_eq!(value, 42);
        assert!(rest.is_empty());
    }

    // High bits in first byte should be ignored
    #[test]
    fn decode_ignores_high_bits() {
        let (value, rest) = decode(&[0b111_01010], 5).unwrap();
        assert_eq!(value, 10);
        assert!(rest.is_empty());
    }

    // Trailing bytes should be returned unconsumed
    #[test]
    fn decode_returns_remainder() {
        let (value, rest) = decode(&[10, 0xAA, 0xBB], 8).unwrap();
        assert_eq!(value, 10);
        assert_eq!(rest, &[0xAA, 0xBB]);
    }

    #[test]
    fn decode_empty_input() {
        assert_eq!(decode(&[], 5), Err(IntegerPrefixError::UnexpectedEnd));
    }

    #[test]
    fn decode_truncated_continuation() {
        // Prefix bits all set (needs continuation) but no more bytes
        assert_eq!(
            decode(&[0b000_11111], 5),
            Err(IntegerPrefixError::UnexpectedEnd)
        );
    }

    fn encode(value: usize, prefix_size: u8) -> Vec<u8> {
        let mut buf = Vec::new();
        encode_into(value, prefix_size, &mut buf);
        buf
    }

    // RFC 7541 C.1.1: encoding 10 with 5-bit prefix
    #[test]
    fn encode_10_prefix_5() {
        assert_eq!(encode(10, 5), vec![0b0_1010]);
    }

    // RFC 7541 C.1.2: encoding 1337 with 5-bit prefix
    #[test]
    fn encode_1337_prefix_5() {
        assert_eq!(encode(1337, 5), vec![0b1_1111, 0b1001_1010, 0b0000_1010]);
    }

    // RFC 7541 C.1.3: encoding 42 with 8-bit prefix
    #[test]
    fn encode_42_prefix_8() {
        assert_eq!(encode(42, 8), vec![42]);
    }

    #[test]
    fn encode_into_appends() {
        let mut buf = vec![0xAA, 0xBB];
        encode_into(42, 8, &mut buf);
        assert_eq!(buf, vec![0xAA, 0xBB, 42]);
    }

    #[test]
    fn encoded_length_matches_encode() {
        for n in [1u8, 5, 7, 8] {
            for value in 0..300 {
                assert_eq!(
                    encoded_length(value, n),
                    encode(value, n).len(),
                    "prefix_size={n}, value={value}",
                );
            }
        }
    }

    #[test]
    fn encoded_length_matches_encode_large() {
        for &value in &[
            0usize,
            1,
            126,
            127,
            128,
            255,
            256,
            16383,
            16384,
            1_000_000,
            usize::MAX >> 1,
        ] {
            for n in [1u8, 5, 7, 8] {
                assert_eq!(
                    encoded_length(value, n),
                    encode(value, n).len(),
                    "prefix_size={n}, value={value}",
                );
            }
        }
    }

    #[test]
    fn roundtrip_small() {
        for n in [1, 5, 7, 8] {
            for value in 0..300 {
                let encoded = encode(value, n);
                let (decoded, rest) = decode(&encoded, n).unwrap();
                assert_eq!(decoded, value, "prefix_size={n}, value={value}");
                assert!(rest.is_empty());
            }
        }
    }

    #[test]
    fn roundtrip_large() {
        for &value in &[
            0,
            1,
            126,
            127,
            128,
            255,
            256,
            16383,
            16384,
            1_000_000,
            usize::MAX >> 1,
        ] {
            let encoded = encode(value, 5);
            let (decoded, rest) = decode(&encoded, 5).unwrap();
            assert_eq!(decoded, value);
            assert!(rest.is_empty());
        }
    }
}