commonware_codec/
util.rs

1//! Codec utility functions
2
3use crate::{Error, FixedSize, Read};
4#[cfg(not(feature = "std"))]
5use alloc::vec::Vec;
6use bytes::Buf;
7#[cfg(feature = "std")]
8use std::vec::Vec;
9
10/// Checks if the buffer has at least `len` bytes remaining. Returns an [Error::EndOfBuffer] if not.
11#[inline]
12pub fn at_least<B: Buf>(buf: &mut B, len: usize) -> Result<(), Error> {
13    let rem = buf.remaining();
14    if rem < len {
15        return Err(Error::EndOfBuffer);
16    }
17    Ok(())
18}
19
20/// Checks if the buffer has at least `len * item_size` bytes remaining, treating multiplication
21/// overflow as insufficient. Returns an [Error::EndOfBuffer] if not.
22#[inline]
23pub fn at_least_items<B: Buf>(buf: &mut B, len: usize, item_size: usize) -> Result<(), Error> {
24    at_least(buf, len.checked_mul(item_size).ok_or(Error::EndOfBuffer)?)
25}
26
27/// Reads `len` values of a [FixedSize] type from the buffer into a vector.
28///
29/// Checks that the buffer contains all `len * SIZE` bytes before allocating or decoding, so a
30/// maliciously large `len` fails fast with [Error::EndOfBuffer]. Intended as a `Read::read_vec`
31/// override for [FixedSize] element types.
32#[inline]
33pub fn read_fixed_vec<T: Read + FixedSize>(
34    buf: &mut impl Buf,
35    len: usize,
36    cfg: &T::Cfg,
37) -> Result<Vec<T>, Error> {
38    at_least_items(buf, len, T::SIZE)?;
39    let mut values = Vec::with_capacity(len);
40    for _ in 0..len {
41        values.push(T::read_cfg(buf, cfg)?);
42    }
43    Ok(values)
44}
45
46/// Ensures the next `size` bytes are all zeroes in the provided buffer, returning an [Error]
47/// otherwise.
48#[inline]
49pub fn ensure_zeros<B: Buf>(buf: &mut B, size: usize) -> Result<(), Error> {
50    at_least(buf, size)?;
51    let mut remaining = size;
52    while remaining > 0 {
53        // Compare (and advance) a chunk at a time rather than a byte at a time. Padding regularly
54        // spans dozens of bytes, and a slice comparison vectorizes.
55        let chunk = buf.chunk();
56        let len = chunk.len().min(remaining);
57        if chunk[..len].iter().any(|&b| b != 0) {
58            return Err(Error::Invalid("codec", "non-zero bytes"));
59        }
60        buf.advance(len);
61        remaining -= len;
62    }
63    Ok(())
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[test]
71    fn test_ensure_zeros() {
72        // Consumes exactly `size` bytes of an all-zero region.
73        let mut buf = &[0u8, 0, 0, 0, 7][..];
74        ensure_zeros(&mut buf, 4).unwrap();
75        assert_eq!(buf.remaining(), 1);
76
77        // A zero-length check consumes nothing, even on an empty buffer.
78        let mut buf = &[][..];
79        ensure_zeros(&mut buf, 0).unwrap();
80
81        // A short buffer fails without panicking.
82        let mut buf = &[0u8, 0][..];
83        assert!(matches!(ensure_zeros(&mut buf, 3), Err(Error::EndOfBuffer)));
84
85        // A non-zero byte anywhere in the region fails.
86        for i in 0..4 {
87            let mut bytes = [0u8; 4];
88            bytes[i] = 1;
89            let mut buf = &bytes[..];
90            assert!(matches!(
91                ensure_zeros(&mut buf, 4),
92                Err(Error::Invalid(_, _))
93            ));
94        }
95    }
96
97    #[test]
98    fn test_ensure_zeros_across_chunks() {
99        // A chained buffer exposes the region as multiple chunks, exercising the chunk loop.
100        let mut buf = (&[0u8, 0][..]).chain(&[0u8, 0, 0][..]);
101        ensure_zeros(&mut buf, 5).unwrap();
102        assert_eq!(buf.remaining(), 0);
103
104        // A non-zero byte in the second chunk still fails.
105        let mut buf = (&[0u8, 0][..]).chain(&[0u8, 2][..]);
106        assert!(matches!(
107            ensure_zeros(&mut buf, 4),
108            Err(Error::Invalid(_, _))
109        ));
110    }
111}
commonware_codec/util.rs

commonware_codec/
util.rs
