use super::Error;
use bytes::{BufMut, Bytes, TryGetError, buf::Take};
use commonware_codec::{
Buf, Decode, Encode, EncodeSize, OwnedBytes, ReadExt as _, Write,
varint::{MAX_U32_VARINT_SIZE, UInt},
};
use commonware_runtime::{Blob, Buf as _, IoBufMut, IoBufs, buffer::paged::Writer};
use commonware_utils::{Cached, Widen};
use std::{future::Future, io::Cursor};
use zstd::{
bulk::Decompressor,
zstd_safe::{CCtx, compress_bound, get_frame_content_size},
};
commonware_utils::thread_local_cache!(static COMPRESSOR: CCtx<'static>);
commonware_utils::thread_local_cache!(static DECOMPRESSOR: Decompressor<'static>);
pub(super) trait FrameReader {
fn read_up_to(
&self,
offset: u64,
len: usize,
buf: impl Into<IoBufMut> + Send,
) -> impl Future<Output = Result<(IoBufMut, usize), Error>> + Send;
fn read_at(
&self,
offset: u64,
len: usize,
) -> impl Future<Output = Result<IoBufs, Error>> + Send;
}
impl<B: Blob, Phase: Send + Sync> FrameReader for Writer<B, Phase> {
async fn read_up_to(
&self,
offset: u64,
len: usize,
buf: impl Into<IoBufMut> + Send,
) -> Result<(IoBufMut, usize), Error> {
Self::read_up_to(self, offset, len, buf)
.await
.map_err(Error::Runtime)
}
async fn read_at(&self, offset: u64, len: usize) -> Result<IoBufs, Error> {
Self::read_at(self, offset, len)
.await
.map_err(Error::Runtime)
}
}
#[inline]
pub(super) fn decode_length_prefix(buf: &mut impl Buf) -> Result<(usize, usize), Error> {
let initial = buf.remaining();
let size = UInt::<u32>::read(buf)?.0 as usize;
let varint_len = initial - buf.remaining();
Ok((size, varint_len))
}
pub(super) enum FrameInfo {
Complete {
varint_len: usize,
data_len: usize,
},
Incomplete {
varint_len: usize,
prefix_len: usize,
total_len: usize,
},
}
pub(super) fn find_frame(buf: &mut impl Buf, offset: u64) -> Result<(u64, FrameInfo), Error> {
let available = buf.remaining();
let (size, varint_len) = decode_length_prefix(buf)?;
let next_offset = offset
.checked_add(varint_len as u64)
.ok_or(Error::OffsetOverflow)?
.checked_add(size as u64)
.ok_or(Error::OffsetOverflow)?;
let buffered = available.saturating_sub(varint_len);
let item = if buffered >= size {
FrameInfo::Complete {
varint_len,
data_len: size,
}
} else {
FrameInfo::Incomplete {
varint_len,
prefix_len: buffered,
total_len: size,
}
};
Ok((next_offset, item))
}
pub(super) fn decompress(compressed: &[u8]) -> Result<Vec<u8>, Error> {
let size = get_frame_content_size(compressed)
.ok()
.flatten()
.and_then(|size| usize::try_from(size).ok())
.ok_or(Error::DecompressionFailed)?;
Cached::take(&DECOMPRESSOR, Decompressor::new, |_| Ok(()))
.and_then(|mut decompressor| decompressor.decompress(compressed, size))
.map_err(|_| Error::DecompressionFailed)
}
pub(super) fn decode_item<V: Decode>(
mut item_data: impl Buf,
cfg: &V::Cfg,
compressed: bool,
) -> Result<V, Error> {
if compressed {
let len = item_data.remaining();
let decompressed = if item_data.chunk().len() == len {
let decompressed = decompress(item_data.chunk());
item_data.advance(len);
decompressed?
} else {
decompress(&item_data.copy_to_bytes(len))?
};
V::decode_cfg(decompressed, cfg).map_err(Error::Codec)
} else {
V::decode_cfg(item_data, cfg).map_err(Error::Codec)
}
}
pub(super) struct Limited<B> {
inner: Take<B>,
}
impl<B: Buf> Limited<B> {
pub(super) fn new(inner: B, limit: usize) -> Self {
Self {
inner: inner.take(limit),
}
}
}
impl<B: Buf> Buf for Limited<B> {
fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
self.inner.take_owned(len)
}
}
impl<B: Buf> bytes::Buf for Limited<B> {
fn remaining(&self) -> usize {
self.inner.remaining()
}
fn chunk(&self) -> &[u8] {
self.inner.chunk()
}
fn advance(&mut self, cnt: usize) {
self.inner.advance(cnt);
}
fn copy_to_bytes(&mut self, len: usize) -> Bytes {
self.inner.copy_to_bytes(len)
}
fn copy_to_slice(&mut self, dst: &mut [u8]) {
let len = dst.len();
if len > self.inner.remaining() {
return self.inner.copy_to_slice(dst);
}
self.inner.get_mut().copy_to_slice(dst);
self.inner.set_limit(self.inner.limit() - len);
}
fn try_copy_to_slice(&mut self, dst: &mut [u8]) -> Result<(), TryGetError> {
let len = dst.len();
let available = self.inner.remaining();
if len > available {
return Err(TryGetError {
requested: len,
available,
});
}
self.inner.get_mut().try_copy_to_slice(dst)?;
self.inner.set_limit(self.inner.limit() - len);
Ok(())
}
}
pub(super) async fn read_frame_at<V: Decode>(
reader: &impl FrameReader,
offset: u64,
cfg: &V::Cfg,
compressed: bool,
) -> Result<(u64, u32, V), Error> {
let (buf, available) = reader
.read_up_to(
offset,
MAX_U32_VARINT_SIZE,
IoBufMut::with_capacity(MAX_U32_VARINT_SIZE),
)
.await?;
let buf = buf.freeze();
let mut cursor = buf.slice(..available);
let (next_offset, item_info) = find_frame(&mut cursor, offset)?;
let (item_size, decoded) = match item_info {
FrameInfo::Complete {
varint_len,
data_len,
} => {
let decoded = decode_item::<V>(
buf.slice(varint_len..varint_len + data_len),
cfg,
compressed,
)?;
(data_len as u32, decoded)
}
FrameInfo::Incomplete {
varint_len,
total_len,
..
} if compressed => {
let data_offset = offset
.checked_add(varint_len as u64)
.ok_or(Error::OffsetOverflow)?;
let data = reader.read_at(data_offset, total_len).await?;
(total_len as u32, decode_item::<V>(data, cfg, compressed)?)
}
FrameInfo::Incomplete {
varint_len,
prefix_len,
total_len,
} => {
let prefix = buf.slice(varint_len..varint_len + prefix_len);
let read_offset = offset
.checked_add(varint_len as u64)
.and_then(|offset| offset.checked_add(prefix_len as u64))
.ok_or(Error::OffsetOverflow)?;
let remainder = reader.read_at(read_offset, total_len - prefix_len).await?;
let decoded = decode_item::<V>(prefix.chain(remainder), cfg, compressed)?;
(total_len as u32, decoded)
}
};
Ok((next_offset, item_size, decoded))
}
pub(super) fn compress_into(level: u8, data: &[u8], buf: &mut Vec<u8>) -> Result<usize, Error> {
let start = buf.len();
buf.reserve(compress_bound(data.len()));
let mut compressor = Cached::take(
&COMPRESSOR,
|| CCtx::try_create().ok_or(Error::CompressionFailed),
|_| Ok(()),
)?;
let mut tail = Cursor::new(buf);
tail.set_position(start as u64);
compressor
.compress(&mut tail, data, level.into())
.map_err(|_| Error::CompressionFailed)
}
pub(super) fn encode_frame_into<V: Encode>(
compression: Option<u8>,
item: &V,
buf: &mut Vec<u8>,
) -> Result<u32, Error> {
if let Some(compression) = compression {
return encode_compressed_frame_into(compression, item, buf);
}
let frame = UncompressedFrame::new(item)?;
buf.reserve(frame.frame_len);
frame.write(buf);
Ok(frame.item_len)
}
#[inline(never)]
pub(super) fn encode_compressed_frame_into<V: Encode>(
compression: u8,
item: &V,
buf: &mut Vec<u8>,
) -> Result<u32, Error> {
let encoded = item.encode();
let start = buf.len();
let max_len = compress_bound(encoded.len());
let max_size_len = UInt(u32::try_from(max_len).unwrap_or(u32::MAX)).encode_size();
let max_entry_len = max_size_len
.checked_add(max_len)
.ok_or(Error::OffsetOverflow)?;
buf.reserve(max_entry_len);
buf.resize(start + max_size_len, 0);
let item_len = compress_into(compression, &encoded, buf)
.and_then(|len| u32::try_from(len).map_err(|_| Error::ItemTooLarge(len)))
.inspect_err(|_| buf.truncate(start))?;
let size_len = UInt(item_len).encode_size();
if size_len < max_size_len {
buf.copy_within(start + max_size_len.., start + size_len);
buf.truncate(start + size_len + Widen::widen(item_len));
}
UInt(item_len).write(&mut &mut buf[start..start + size_len]);
Ok(item_len)
}
pub(super) struct UncompressedFrame<'a, V> {
item: &'a V,
pub(super) item_len: u32,
frame_len: usize,
}
impl<'a, V: EncodeSize> UncompressedFrame<'a, V> {
pub(super) fn new(item: &'a V) -> Result<Self, Error> {
let len = item.encode_size();
let item_len = u32::try_from(len).map_err(|_| Error::ItemTooLarge(len))?;
let frame_len = UInt(item_len)
.encode_size()
.checked_add(len)
.ok_or(Error::OffsetOverflow)?;
Ok(Self {
item,
item_len,
frame_len,
})
}
}
impl<V> EncodeSize for UncompressedFrame<'_, V> {
fn encode_size(&self) -> usize {
self.frame_len
}
}
impl<V: Write> Write for UncompressedFrame<'_, V> {
fn write(&self, buf: &mut impl BufMut) {
UInt(self.item_len).write(buf);
self.item.write(buf);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::utils::codec::View;
use commonware_codec::{Copying, Encode, Error as CodecError, Read};
use commonware_utils::test_rng;
use rand::{Rng as _, RngExt as _};
use std::{
cell::Cell,
panic::{AssertUnwindSafe, catch_unwind},
};
use zstd::bulk::compress;
struct TrackingBuf {
inner: Bytes,
copies: usize,
try_copies: usize,
}
impl TrackingBuf {
fn new(inner: Bytes) -> Self {
Self {
inner,
copies: 0,
try_copies: 0,
}
}
}
impl Buf for TrackingBuf {}
impl bytes::Buf for TrackingBuf {
fn remaining(&self) -> usize {
self.inner.remaining()
}
fn chunk(&self) -> &[u8] {
self.inner.chunk()
}
fn advance(&mut self, cnt: usize) {
self.inner.advance(cnt);
}
fn copy_to_bytes(&mut self, len: usize) -> Bytes {
self.inner.copy_to_bytes(len)
}
fn copy_to_slice(&mut self, dst: &mut [u8]) {
self.copies += 1;
self.inner.copy_to_slice(dst);
}
fn try_copy_to_slice(&mut self, dst: &mut [u8]) -> Result<(), TryGetError> {
self.try_copies += 1;
self.inner.try_copy_to_slice(dst)
}
}
fn frame<V: Encode>(compression: Option<u8>, item: &V) -> Vec<u8> {
let mut buf = Vec::new();
encode_frame_into(compression, item, &mut buf).unwrap();
buf
}
#[test]
fn test_limited_buf_forwards_bulk_copies() {
let mut backing = TrackingBuf::new(Bytes::from_static(b"abcdefghijNEXT"));
{
let mut limited = Limited::new(&mut backing, 10);
let mut empty = [];
limited.try_copy_to_slice(&mut empty).unwrap();
let mut prefix = [0; 4];
limited.try_copy_to_slice(&mut prefix).unwrap();
assert_eq!(&prefix, b"abcd");
let mut rest = [0; 6];
limited.copy_to_slice(&mut rest);
assert_eq!(&rest, b"efghij");
assert_eq!(limited.remaining(), 0);
}
assert_eq!(backing.try_copies, 2);
assert_eq!(backing.copies, 1);
assert_eq!(backing.inner.as_ref(), b"NEXT");
}
#[test]
fn test_limited_buf_rejects_overreads_without_mutation() {
let mut backing = TrackingBuf::new(Bytes::from_static(b"dataNEXT"));
let mut limited = Limited::new(&mut backing, 4);
let mut dst = [0xAA; 5];
let err = limited.try_copy_to_slice(&mut dst).unwrap_err();
assert_eq!(err.requested, 5);
assert_eq!(err.available, 4);
assert_eq!(dst, [0xAA; 5]);
assert_eq!(limited.remaining(), 4);
assert_eq!(limited.inner.get_ref().inner.as_ref(), b"dataNEXT");
assert_eq!(limited.inner.get_ref().try_copies, 0);
let mut dst = [0xBB; 5];
assert!(catch_unwind(AssertUnwindSafe(|| limited.copy_to_slice(&mut dst))).is_err());
assert_eq!(dst, [0xBB; 5]);
assert_eq!(limited.remaining(), 4);
assert_eq!(limited.inner.get_ref().inner.as_ref(), b"dataNEXT");
let err = limited.try_get_u64().unwrap_err();
assert_eq!(err.requested, 8);
assert_eq!(err.available, 4);
assert_eq!(limited.remaining(), 4);
let mut short = TrackingBuf::new(Bytes::from_static(b"abc"));
let mut limited = Limited::new(&mut short, usize::MAX);
let err = limited.try_copy_to_slice(&mut [0; 4]).unwrap_err();
assert_eq!(err.requested, 4);
assert_eq!(err.available, 3);
assert_eq!(limited.remaining(), 3);
assert_eq!(limited.inner.get_ref().inner.as_ref(), b"abc");
}
#[test]
fn test_limited_buf_preserves_owned_bytes_and_frame_bounds() {
let source = Bytes::from_static(b"abcdefghNEXT");
let source_ptr = source.as_ptr();
let retained = {
let mut limited = Limited::new(source.clone(), 8);
limited.copy_to_bytes(8)
};
assert_eq!(retained.as_ptr(), source_ptr);
drop(source);
assert_eq!(retained.as_ref(), b"abcdefgh");
let mut source = Bytes::from_static(b"abcdefghNEXT");
assert!(matches!(
decode_item::<u8>(Limited::new(&mut source, 8), &(), false),
Err(Error::Codec(CodecError::ExtraData(7)))
));
assert!(source.as_ref().ends_with(b"NEXT"));
let mut source = Bytes::from_static(b"abcdefghNEXT");
assert!(matches!(
decode_item::<u64>(Limited::new(&mut source, 7), &(), false),
Err(Error::Codec(CodecError::EndOfBuffer))
));
assert_eq!(source.as_ref(), b"abcdefghNEXT");
}
#[test]
fn test_roundtrip_uncompressed() {
let buf = frame(None, &42u64);
let mut cursor = Copying(&buf);
let (next_offset, info) = find_frame(&mut cursor, 0).unwrap();
let FrameInfo::Complete {
varint_len,
data_len,
} = info
else {
panic!("expected complete frame");
};
assert_eq!(varint_len, 1);
assert_eq!(data_len, 8);
assert_eq!(next_offset, 9);
let item: u64 =
decode_item(Copying(&buf[varint_len..varint_len + data_len]), &(), false).unwrap();
assert_eq!(item, 42);
}
#[test]
fn test_roundtrip_compressed() {
let buf = frame(Some(3), &42u64);
let mut cursor = Copying(&buf);
let (_, info) = find_frame(&mut cursor, 0).unwrap();
let FrameInfo::Complete {
varint_len,
data_len,
} = info
else {
panic!("expected complete frame");
};
let item: u64 =
decode_item(Copying(&buf[varint_len..varint_len + data_len]), &(), true).unwrap();
assert_eq!(item, 42);
}
#[test]
fn test_cached_compressor_matches_independent_frames() {
let cases = [
("empty", 0, 3),
("single byte", 1, 1),
("small record", 64, 19),
("grow to 1 KiB", 1024, 3),
("grow to 16 KiB", 16 * 1024, 1),
("row match finder", 64 * 1024, 7),
("cross a 128 KiB block", 130 * 1024, 3),
("small after large", 64, 0),
("clamped level", 1024, 255),
("empty after nonempty", 0, 9),
];
let mut rng = test_rng();
let mut cached_frames = b"existing data".to_vec();
let mut independent_frames = cached_frames.clone();
for (case, len, level) in cases {
let repeated = vec![0xAB; len];
let mut random = vec![0; len];
rng.fill_bytes(&mut random);
let partial = (0..len).map(|_| rng.random_range(0..16u8)).collect();
for (pattern, item) in [
("repeated", repeated),
("random", random),
("partly compressible", partial),
] {
let compressed = zstd::bulk::compress(&item.encode(), level.into()).unwrap();
UInt(compressed.len() as u32).write(&mut independent_frames);
independent_frames.extend_from_slice(&compressed);
let frame_start = cached_frames.len();
let compressed_len =
encode_frame_into(Some(level), &item, &mut cached_frames).unwrap();
assert_eq!(
compressed_len as usize,
compressed.len(),
"{case}, {pattern}, compression level {level}"
);
assert_eq!(
cached_frames, independent_frames,
"{case}, {pattern}, compression level {level}"
);
let mut frame = Copying(&cached_frames[frame_start..]);
let (frame_len, _) = decode_length_prefix(&mut frame).unwrap();
assert_eq!(frame.remaining(), frame_len);
let decoded = decode_item::<Vec<u8>>(frame, &((..).into(), ()), true).unwrap();
assert_eq!(
decoded, item,
"{case}, {pattern}, compression level {level}"
);
}
}
assert!(
COMPRESSOR.with(|slot| slot.borrow().1.is_some()),
"the context must return to the thread's cache"
);
}
#[test]
fn test_accumulation_preserves_existing_contents() {
let mut buf = Vec::new();
encode_frame_into(None, &1u64, &mut buf).unwrap();
let first_frame_len = buf.len();
encode_frame_into(None, &2u64, &mut buf).unwrap();
let mut cursor = Copying(&buf);
let (first_end, _) = find_frame(&mut cursor, 0).unwrap();
assert_eq!(first_end as usize, first_frame_len);
let first: u64 = decode_item(Copying(&buf[1..9]), &(), false).unwrap();
assert_eq!(first, 1);
let mut cursor = Copying(&buf[first_frame_len..]);
let (second_end, _) = find_frame(&mut cursor, first_end).unwrap();
assert_eq!(second_end as usize, buf.len());
let second: u64 = decode_item(Copying(&buf[first_frame_len + 1..]), &(), false).unwrap();
assert_eq!(second, 2);
}
#[test]
fn test_uncompressed_frame_sizes_item_once() {
struct Counted<'a> {
bytes: &'a [u8],
size_calls: Cell<usize>,
}
impl EncodeSize for Counted<'_> {
fn encode_size(&self) -> usize {
self.size_calls.set(self.size_calls.get() + 1);
self.bytes.len()
}
}
impl Write for Counted<'_> {
fn write(&self, buf: &mut impl BufMut) {
buf.put_slice(self.bytes);
}
}
for len in [0, 127, 128, 16_383, 16_384] {
let bytes = vec![7; len];
let item = Counted {
bytes: &bytes,
size_calls: Cell::new(0),
};
let frame = UncompressedFrame::new(&item).unwrap();
assert_eq!(frame.item_len as usize, len);
assert_eq!(frame.encode_size(), UInt(len as u32).encode_size() + len);
let mut encoded = Vec::new();
frame.write(&mut encoded);
let mut expected = Vec::new();
UInt(len as u32).write(&mut expected);
expected.extend_from_slice(&bytes);
assert_eq!(encoded, expected);
assert_eq!(item.size_calls.get(), 1);
}
}
#[test]
fn test_find_frame_zero_length_payload() {
let buf = [0x00u8];
let mut cursor = Copying(&buf);
let (next_offset, info) = find_frame(&mut cursor, 7).unwrap();
let FrameInfo::Complete {
varint_len,
data_len,
} = info
else {
panic!("expected complete frame");
};
assert_eq!((varint_len, data_len), (1, 0));
assert_eq!(next_offset, 8);
}
#[test]
fn test_find_frame_incomplete_payload() {
let buf = [0x05u8, 1, 2, 3];
let mut cursor = Copying(&buf);
let (next_offset, info) = find_frame(&mut cursor, 100).unwrap();
let FrameInfo::Incomplete {
varint_len,
prefix_len,
total_len,
} = info
else {
panic!("expected incomplete frame");
};
assert_eq!((varint_len, prefix_len, total_len), (1, 3, 5));
assert_eq!(next_offset, 106);
assert_eq!(cursor.remaining(), 3);
}
#[test]
fn test_find_frame_payload_boundary() {
let buf = [0x03u8, 1, 2, 3];
let mut cursor = Copying(&buf);
assert!(matches!(
find_frame(&mut cursor, 0).unwrap().1,
FrameInfo::Complete { data_len: 3, .. }
));
let buf = [0x03u8, 1, 2];
let mut cursor = Copying(&buf);
assert!(matches!(
find_frame(&mut cursor, 0).unwrap().1,
FrameInfo::Incomplete {
prefix_len: 2,
total_len: 3,
..
}
));
}
#[test]
fn test_find_frame_empty_buffer() {
let mut cursor = Copying(&[]);
assert!(matches!(find_frame(&mut cursor, 0), Err(Error::Codec(_))));
}
#[test]
fn test_find_frame_truncated_varint() {
let buf = [0x80u8];
let mut cursor = Copying(&buf);
assert!(matches!(find_frame(&mut cursor, 0), Err(Error::Codec(_))));
}
#[test]
fn test_find_frame_varint_exceeds_u32() {
let buf = [0xFFu8, 0xFF, 0xFF, 0xFF, 0x7F];
let mut cursor = Copying(&buf);
assert!(matches!(find_frame(&mut cursor, 0), Err(Error::Codec(_))));
}
#[test]
fn test_find_frame_offset_overflow() {
let buf = frame(None, &42u64);
let mut cursor = Copying(&buf);
assert!(matches!(
find_frame(&mut cursor, u64::MAX),
Err(Error::OffsetOverflow)
));
}
#[test]
fn test_decode_item_rejects_trailing_bytes() {
let buf = [0u8; 9];
assert!(matches!(
decode_item::<u64>(Copying(&buf), &(), false),
Err(Error::Codec(commonware_codec::Error::ExtraData(_)))
));
}
#[test]
fn test_decode_item_view() {
let value: Vec<Bytes> = (0..64).map(|_| Bytes::from(vec![7u8; 17])).collect();
let cfg = ((..).into(), (..).into());
let buf = value.encode();
let range = buf.as_ptr_range();
let decoded = decode_item::<Vec<Bytes>>(buf.clone(), &cfg, false).unwrap();
assert_eq!(decoded, value);
assert!(decoded.iter().all(|b| range.contains(&b.as_ptr())));
let copied = decode_item::<Vec<Bytes>>(Copying(&buf), &cfg, false).unwrap();
assert_eq!(copied, value);
assert!(copied.iter().all(|b| !range.contains(&b.as_ptr())));
let value = vec![View::new(1), View::new(2)];
let buf = Bytes::from(zstd::bulk::compress(&value.encode(), 3).unwrap());
let decoded = decode_item::<Vec<View>>(buf, &((..).into(), ()), true).unwrap();
assert_eq!(decoded.len(), value.len());
for (decoded, expected) in decoded.iter().zip(&value) {
assert_eq!(decoded.bytes, expected.bytes);
decoded.assert_shared();
}
}
#[test]
fn test_decode_item_corrupt_compressed_payload() {
let mut buf = frame(Some(3), &42u64);
buf[1] ^= 0xFF;
assert!(matches!(
decode_item::<u64>(Copying(&buf[1..]), &(), true),
Err(Error::DecompressionFailed)
));
}
#[test]
fn test_decode_item_decompresses_split_payloads_and_consumes_exactly() {
let cfg = ((..).into(), ());
for (len, level) in [(0usize, 3), (1, 1), (4096, 19), (70_000, 3)] {
let item: Vec<u8> = (0..len).map(|i| (i % 7) as u8).collect();
let payload = compress(&item.encode(), level).unwrap();
let mut source = Bytes::from([payload.as_slice(), &[9, 9]].concat());
let decoded =
decode_item::<Vec<u8>>((&mut source).take(payload.len()), &cfg, true).unwrap();
assert_eq!(decoded, item);
assert_eq!(source.as_ref(), &[9, 9]);
let (head, tail) = payload.split_at(payload.len().min(3));
let mut split =
Bytes::copy_from_slice(head).chain(Bytes::from([tail, &[9, 9]].concat()));
let decoded =
decode_item::<Vec<u8>>((&mut split).take(payload.len()), &cfg, true).unwrap();
assert_eq!(decoded, item);
assert_eq!(split.copy_to_bytes(split.remaining()).as_ref(), &[9, 9]);
}
}
#[test]
fn test_compressed_frames_declare_content_size() {
for len in [0usize, 1, 4096] {
let item = vec![7u8; len];
let buf = frame(Some(3), &item);
let (_, varint_len) = decode_length_prefix(&mut Copying(&buf)).unwrap();
assert_eq!(
get_frame_content_size(&buf[varint_len..]).unwrap(),
Some(item.encode_size() as u64)
);
}
}
#[test]
fn test_decompress_reuses_thread_context() {
for len in [70_000usize, 64, 0] {
let data: Vec<u8> = (0..len).map(|i| (i % 7) as u8).collect();
let payload = compress(&data, 3).unwrap();
assert_eq!(get_frame_content_size(&payload).unwrap(), Some(len as u64));
assert_eq!(decompress(&payload).unwrap(), data);
}
assert!(
DECOMPRESSOR.with(|slot| slot.borrow().1.is_some()),
"the context must return to the thread's cache"
);
}
#[test]
fn test_decompress_rejects_frame_without_content_size() {
let payload = zstd::stream::encode_all(42u64.encode().as_ref(), 3).unwrap();
assert_eq!(get_frame_content_size(&payload).unwrap(), None);
assert!(matches!(
decompress(&payload),
Err(Error::DecompressionFailed)
));
}
#[test]
fn test_decompress_rejects_wrong_declared_size() {
let data = vec![7u8; 100];
let payload = compress(&data, 3).unwrap();
assert_eq!(payload[5], 100);
for declared in [99u8, 101] {
let mut corrupted = payload.clone();
corrupted[5] = declared;
assert_eq!(
get_frame_content_size(&corrupted).unwrap(),
Some(u64::from(declared))
);
assert!(matches!(
decompress(&corrupted),
Err(Error::DecompressionFailed)
));
}
}
#[test]
fn test_decompress_context_recovers_after_errors() {
let payload = compress(&42u64.encode(), 3).unwrap();
for len in 0..payload.len() {
assert!(decompress(&payload[..len]).is_err());
assert_eq!(
decode_item::<u64>(Copying(&payload), &(), true).unwrap(),
42
);
}
}
struct Oversized;
impl EncodeSize for Oversized {
fn encode_size(&self) -> usize {
u32::MAX as usize + 1
}
}
impl Write for Oversized {
fn write(&self, _: &mut impl BufMut) {
unreachable!("size check rejects the item before writing")
}
}
impl Read for Oversized {
type Cfg = ();
fn read_cfg(_: &mut impl Buf, _: &Self::Cfg) -> Result<Self, commonware_codec::Error> {
unreachable!("never decoded")
}
}
#[test]
fn test_encode_frame_rejects_oversized_item() {
let mut buf = Vec::new();
assert!(matches!(
encode_frame_into(None, &Oversized, &mut buf),
Err(Error::ItemTooLarge(_))
));
assert!(buf.is_empty());
}
}