Skip to main content

commonware_codec/
util.rs

1//! Helpers for reading, validating, and sizing encoded values.
2
3use crate::{Buf, Error, FixedSize, MaxEncodeSize, Read, varint::UInt};
4#[cfg(not(feature = "std"))]
5use alloc::vec::Vec;
6#[cfg(feature = "std")]
7use std::vec::Vec;
8
9/// Checks if the buffer has at least `len` bytes remaining. Returns an [Error::EndOfBuffer] if not.
10#[inline]
11pub fn at_least<B: Buf>(buf: &mut B, len: usize) -> Result<(), Error> {
12    let rem = buf.remaining();
13    if rem < len {
14        return Err(Error::EndOfBuffer);
15    }
16    Ok(())
17}
18
19/// Checks if the buffer has at least `len * item_size` bytes remaining, treating multiplication
20/// overflow as insufficient. Returns an [Error::EndOfBuffer] if not.
21#[inline]
22pub fn at_least_items<B: Buf>(buf: &mut B, len: usize, item_size: usize) -> Result<(), Error> {
23    at_least(buf, len.checked_mul(item_size).ok_or(Error::EndOfBuffer)?)
24}
25
26/// Returns an upper bound on the encoded size of a `Vec` or other length-prefixed sequence of
27/// `len` items that each encode to at most `item_size` bytes, with the length prefix at its
28/// widest.
29///
30/// # Panics
31///
32/// Panics if `len` exceeds `u32::MAX`.
33#[inline]
34pub fn max_items_size(len: usize, item_size: usize) -> usize {
35    u32::try_from(len).expect("len must fit the u32 length prefix");
36    len * item_size + UInt::<u32>::MAX_ENCODE_SIZE
37}
38
39/// Reads `len` values of a [FixedSize] type from the buffer into a vector.
40///
41/// Checks that the buffer contains all `len * SIZE` bytes before allocating or decoding, so a
42/// maliciously large `len` fails fast with [Error::EndOfBuffer]. Intended as a `Read::read_vec`
43/// override for [FixedSize] element types.
44#[inline]
45pub fn read_fixed_vec<T: Read + FixedSize>(
46    buf: &mut impl Buf,
47    len: usize,
48    cfg: &T::Cfg,
49) -> Result<Vec<T>, Error> {
50    at_least_items(buf, len, T::SIZE)?;
51    let mut values = Vec::with_capacity(len);
52    for _ in 0..len {
53        values.push(T::read_cfg(buf, cfg)?);
54    }
55    Ok(values)
56}
57
58/// Ensures the next `size` bytes are all zeroes in the provided buffer, returning an [Error]
59/// otherwise.
60#[inline]
61pub fn ensure_zeros<B: Buf>(buf: &mut B, size: usize) -> Result<(), Error> {
62    at_least(buf, size)?;
63    let mut remaining = size;
64    while remaining > 0 {
65        // Compare (and advance) a chunk at a time rather than a byte at a time. Padding regularly
66        // spans dozens of bytes, and a slice comparison vectorizes.
67        let chunk = buf.chunk();
68        let len = chunk.len().min(remaining);
69        if chunk[..len].iter().any(|&b| b != 0) {
70            return Err(Error::Invalid("codec", "non-zero bytes"));
71        }
72        buf.advance(len);
73        remaining -= len;
74    }
75    Ok(())
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use crate::{Copying, EncodeSize, varint::MAX_U32_VARINT_SIZE};
82    use bytes::Buf as _;
83
84    #[test]
85    fn test_max_items_size() {
86        // The bound reserves a 5-byte length prefix and covers an encoded vector of 4-byte items
87        // on both sides of the one-byte prefix boundary.
88        for (len, expected) in [(0, 5), (1, 9), (127, 513), (128, 517)] {
89            let items = vec![[0u8; 4]; len];
90            let size = max_items_size(len, 4);
91            assert_eq!(size, expected);
92            assert!(items.encode_size() <= size);
93        }
94
95        // Up to `u32::MAX` items fit the length prefix.
96        const LARGEST: usize = u32::MAX as usize;
97        assert_eq!(max_items_size(LARGEST, 1), LARGEST + MAX_U32_VARINT_SIZE);
98    }
99
100    #[test]
101    #[should_panic(expected = "len must fit the u32 length prefix")]
102    fn test_max_items_size_beyond_prefix() {
103        max_items_size(u32::MAX as usize + 1, 0);
104    }
105
106    #[test]
107    fn test_ensure_zeros() {
108        // Consumes exactly `size` bytes of an all-zero region.
109        let mut buf = Copying(&[0u8, 0, 0, 0, 7]);
110        ensure_zeros(&mut buf, 4).unwrap();
111        assert_eq!(buf.remaining(), 1);
112
113        // A zero-length check consumes nothing, even on an empty buffer.
114        let mut buf = Copying(&[]);
115        ensure_zeros(&mut buf, 0).unwrap();
116
117        // A short buffer fails without panicking.
118        let mut buf = Copying(&[0u8, 0]);
119        assert!(matches!(ensure_zeros(&mut buf, 3), Err(Error::EndOfBuffer)));
120
121        // A non-zero byte anywhere in the region fails.
122        for i in 0..4 {
123            let mut bytes = [0u8; 4];
124            bytes[i] = 1;
125            let mut buf = Copying(&bytes);
126            assert!(matches!(
127                ensure_zeros(&mut buf, 4),
128                Err(Error::Invalid(_, _))
129            ));
130        }
131    }
132
133    #[test]
134    fn test_ensure_zeros_across_chunks() {
135        // A chained buffer exposes the region as multiple chunks, exercising the chunk loop.
136        let mut buf = Copying(&[0u8, 0]).chain(Copying(&[0u8, 0, 0]));
137        ensure_zeros(&mut buf, 5).unwrap();
138        assert_eq!(buf.remaining(), 0);
139
140        // A non-zero byte in the second chunk still fails.
141        let mut buf = Copying(&[0u8, 0]).chain(Copying(&[0u8, 2]));
142        assert!(matches!(
143            ensure_zeros(&mut buf, 4),
144            Err(Error::Invalid(_, _))
145        ));
146    }
147}