#[cfg(feature = "compression")]
use lz4_flex;
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Debug, Error, Clone, PartialEq)]
pub enum ByteStorageError {
#[error("input exceeds maximum size")]
InputTooLarge,
#[error("decompression ratio exceeds safety limit")]
DecompressionBomb,
#[error("integrity check failed")]
ChecksumMismatch,
#[error("compression failed")]
CompressionFailed,
#[error("decompression failed")]
DecompressionFailed,
#[error("size validation failed")]
SizeValidationFailed,
#[error("serialization failed: {0}")]
SerializationFailed(String),
#[error("deserialization failed: {0}")]
DeserializationFailed(String),
}
const MAX_UNCOMPRESSED_SIZE: usize = 512 * 1024 * 1024; const MAX_COMPRESSED_SIZE: usize = 512 * 1024 * 1024; const MAX_COMPRESSION_RATIO: u64 = 1000;
#[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
const LARGE_PAYLOAD: usize = 64 * 1024;
#[cfg(feature = "compression")]
const LITERAL_BLOCK_THRESHOLD: usize = 64;
#[cfg(feature = "compression")]
const _: () = assert!(LITERAL_BLOCK_THRESHOLD - 1 <= 15 + 254);
#[cfg(feature = "compression")]
fn compress_block(data: &[u8]) -> Vec<u8> {
if data.len() >= LITERAL_BLOCK_THRESHOLD {
return lz4_flex::compress(data);
}
let len = data.len();
let mut block = Vec::with_capacity(len + 2);
if len < 15 {
block.push((len as u8) << 4);
} else {
block.push(0xF0);
block.push((len - 15) as u8);
}
block.extend_from_slice(data);
block
}
#[derive(Serialize, Deserialize)]
pub struct StorageEnvelope {
#[serde(with = "serde_bytes")]
pub compressed_data: Vec<u8>,
pub checksum: [u8; 8],
pub original_size: u32,
pub format: String,
}
impl StorageEnvelope {
#[cfg(all(feature = "compression", feature = "checksum"))]
pub fn new(data: &[u8], format: String) -> Result<Self, ByteStorageError> {
if data.len() > MAX_UNCOMPRESSED_SIZE {
return Err(ByteStorageError::InputTooLarge);
}
let original_size = data.len() as u32;
let compressed_data = compress_block(data);
if compressed_data.len() > MAX_COMPRESSED_SIZE {
return Err(ByteStorageError::InputTooLarge);
}
let checksum = crate::checksum::checksum(data);
Ok(StorageEnvelope {
compressed_data,
checksum,
original_size,
format,
})
}
#[cfg(all(feature = "compression", feature = "checksum"))]
pub fn extract(&self) -> Result<Vec<u8>, ByteStorageError> {
extract_payload(&self.compressed_data, &self.checksum, self.original_size)
}
#[cfg(feature = "messagepack")]
pub fn encoded_len(&self) -> usize {
fn uint_len(v: u32) -> usize {
match v {
0..=0x7F => 1,
0x80..=0xFF => 2,
0x100..=0xFFFF => 3,
_ => 5,
}
}
fn bin_header_len(len: usize) -> usize {
match len {
0..=0xFF => 2,
0x100..=0xFFFF => 3,
_ => 5,
}
}
fn str_header_len(len: usize) -> usize {
match len {
0..=31 => 1,
32..=0xFF => 2,
0x100..=0xFFFF => 3,
_ => 5,
}
}
let checksum_len = 1 + self
.checksum
.iter()
.map(|&b| uint_len(b.into()))
.sum::<usize>();
let fixed = 1 + bin_header_len(self.compressed_data.len())
+ checksum_len
+ uint_len(self.original_size)
+ str_header_len(self.format.len());
fixed
.saturating_add(self.compressed_data.len())
.saturating_add(self.format.len())
}
#[cfg(feature = "messagepack")]
pub fn encode_into(&self, out: &mut [u8]) -> Result<(), ByteStorageError> {
let len = self.capped_encoded_len()?;
if out.len() != len {
return Err(ByteStorageError::SerializationFailed(format!(
"output buffer is {} bytes, the envelope encodes to {len}",
out.len()
)));
}
let mut rest = &mut out[..];
rmp_serde::encode::write(&mut rest, self)
.map_err(|e| ByteStorageError::SerializationFailed(e.to_string()))?;
if !rest.is_empty() {
return Err(ByteStorageError::SerializationFailed(format!(
"envelope encoded {} bytes short of its computed length",
rest.len()
)));
}
Ok(())
}
#[cfg(feature = "messagepack")]
fn capped_encoded_len(&self) -> Result<usize, ByteStorageError> {
let len = self.encoded_len();
if len > MAX_COMPRESSED_SIZE {
return Err(ByteStorageError::InputTooLarge);
}
Ok(len)
}
#[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
fn encode(&self) -> Result<Vec<u8>, ByteStorageError> {
let len = self.capped_encoded_len()?;
if len >= LARGE_PAYLOAD {
return rmp_serde::to_vec(self)
.map_err(|e| ByteStorageError::SerializationFailed(e.to_string()));
}
let mut out = Vec::with_capacity(len);
rmp_serde::encode::write(&mut out, self)
.map_err(|e| ByteStorageError::SerializationFailed(e.to_string()))?;
debug_assert_eq!(out.len(), len, "encoded_len must be exact");
Ok(out)
}
}
#[cfg(all(feature = "compression", feature = "checksum"))]
fn extract_payload(
compressed_data: &[u8],
checksum: &[u8; 8],
original_size: u32,
) -> Result<Vec<u8>, ByteStorageError> {
#[cfg(test)]
crate::read_allocation_probe::reserve_first_control(compressed_data, original_size);
check_decompression_bound(compressed_data.len(), original_size)?;
let decompressed = lz4_flex::decompress(compressed_data, original_size as usize)
.map_err(|_| ByteStorageError::DecompressionFailed)?;
if !crate::checksum::verify_checksum(&decompressed, checksum) {
return Err(ByteStorageError::ChecksumMismatch);
}
if decompressed.len() != original_size as usize {
return Err(ByteStorageError::SizeValidationFailed);
}
Ok(decompressed)
}
#[cfg(all(feature = "compression", feature = "checksum"))]
fn check_decompression_bound(
compressed_len: usize,
original_size: u32,
) -> Result<(), ByteStorageError> {
if compressed_len > MAX_COMPRESSED_SIZE {
return Err(ByteStorageError::InputTooLarge);
}
if original_size as usize > MAX_UNCOMPRESSED_SIZE {
return Err(ByteStorageError::InputTooLarge);
}
let compressed_size = compressed_len as u64;
if compressed_size == 0 {
return Err(ByteStorageError::DecompressionBomb);
}
let max_allowed_original = MAX_COMPRESSION_RATIO
.checked_mul(compressed_size)
.ok_or(ByteStorageError::DecompressionBomb)?;
if (original_size as u64) > max_allowed_original {
return Err(ByteStorageError::DecompressionBomb);
}
Ok(())
}
pub struct ByteStorage {
default_format: String,
}
impl ByteStorage {
pub fn new(default_format: Option<String>) -> Self {
ByteStorage {
default_format: default_format.unwrap_or_else(|| "msgpack".to_string()),
}
}
#[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
pub fn store(&self, data: &[u8], format: Option<String>) -> Result<Vec<u8>, ByteStorageError> {
if data.len() > MAX_UNCOMPRESSED_SIZE {
return Err(ByteStorageError::InputTooLarge);
}
let format = format.unwrap_or_else(|| self.default_format.clone());
let envelope = StorageEnvelope::new(data, format)?;
envelope.encode()
}
#[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
pub fn retrieve(&self, envelope_bytes: &[u8]) -> Result<(Vec<u8>, String), ByteStorageError> {
if envelope_bytes.len() > MAX_COMPRESSED_SIZE {
return Err(ByteStorageError::InputTooLarge);
}
let envelope = decode_envelope(envelope_bytes)?;
let data = envelope.extract()?;
Ok((data, envelope.format))
}
#[cfg(feature = "compression")]
pub fn estimate_compression(&self, data: &[u8]) -> Result<f64, ByteStorageError> {
if data.len() > MAX_UNCOMPRESSED_SIZE {
return Err(ByteStorageError::InputTooLarge);
}
let compressed = compress_block(data);
Ok(data.len() as f64 / compressed.len() as f64)
}
#[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
pub fn validate(&self, envelope_bytes: &[u8]) -> bool {
if envelope_bytes.len() > MAX_COMPRESSED_SIZE {
return false; }
match decode_envelope(envelope_bytes) {
Ok(envelope) => envelope.extract().is_ok(),
Err(_) => false,
}
}
pub fn max_uncompressed_size(&self) -> usize {
MAX_UNCOMPRESSED_SIZE
}
pub fn max_compressed_size(&self) -> usize {
MAX_COMPRESSED_SIZE
}
pub fn max_compression_ratio(&self) -> u64 {
MAX_COMPRESSION_RATIO
}
}
impl Default for ByteStorage {
fn default() -> Self {
Self::new(None)
}
}
#[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
fn decode_envelope(
envelope_bytes: &[u8],
) -> Result<borrowed::StorageEnvelope<'_>, ByteStorageError> {
const MAX_DEPTH: usize = 100;
crate::check_msgpack_structure(envelope_bytes, MAX_DEPTH).map_err(|what| {
ByteStorageError::DeserializationFailed(format!("decode pre-scan: {what}"))
})?;
rmp_serde::from_slice(envelope_bytes)
.map_err(|e| ByteStorageError::DeserializationFailed(e.to_string()))
}
#[cfg(all(feature = "compression", feature = "checksum", feature = "messagepack"))]
mod borrowed {
use std::borrow::Cow;
use serde::de::{Deserializer, Error, SeqAccess, Visitor};
use serde::Deserialize;
use super::{extract_payload, ByteStorageError, LARGE_PAYLOAD};
#[derive(Deserialize)]
pub(super) struct StorageEnvelope<'a> {
#[serde(borrow, deserialize_with = "borrow_payload")]
pub(super) compressed_data: Cow<'a, [u8]>,
pub(super) checksum: [u8; 8],
pub(super) original_size: u32,
pub(super) format: String,
}
impl StorageEnvelope<'_> {
pub(super) fn extract(&self) -> Result<Vec<u8>, ByteStorageError> {
extract_payload(&self.compressed_data, &self.checksum, self.original_size)
}
}
fn borrow_payload<'de, D: Deserializer<'de>>(
deserializer: D,
) -> Result<Cow<'de, [u8]>, D::Error> {
struct PayloadVisitor;
impl<'de> Visitor<'de> for PayloadVisitor {
type Value = Cow<'de, [u8]>;
fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
formatter.write_str("byte array")
}
fn visit_borrowed_bytes<E: Error>(self, v: &'de [u8]) -> Result<Self::Value, E> {
Ok(borrow_below_large(v))
}
fn visit_borrowed_str<E: Error>(self, v: &'de str) -> Result<Self::Value, E> {
Ok(borrow_below_large(v.as_bytes()))
}
fn visit_bytes<E: Error>(self, v: &[u8]) -> Result<Self::Value, E> {
Ok(Cow::Owned(v.to_vec()))
}
fn visit_str<E: Error>(self, v: &str) -> Result<Self::Value, E> {
Ok(Cow::Owned(v.as_bytes().to_vec()))
}
fn visit_seq<V: SeqAccess<'de>>(self, mut seq: V) -> Result<Self::Value, V::Error> {
let mut bytes = Vec::with_capacity(seq.size_hint().unwrap_or(0).min(4096));
while let Some(b) = seq.next_element()? {
bytes.push(b);
}
Ok(Cow::Owned(bytes))
}
}
deserializer.deserialize_byte_buf(PayloadVisitor)
}
fn borrow_below_large(v: &[u8]) -> Cow<'_, [u8]> {
if v.len() < LARGE_PAYLOAD {
Cow::Borrowed(v)
} else {
Cow::Owned(v.to_vec())
}
}
}
#[cfg(all(
test,
feature = "compression",
feature = "checksum",
feature = "messagepack"
))]
mod tests {
use super::*;
#[test]
fn envelope_embeds_canonical_checksum() {
let data = b"DRY-guard payload";
let envelope = StorageEnvelope::new(data, "test".to_string()).unwrap();
assert_eq!(envelope.checksum, crate::checksum::checksum(data));
}
#[test]
fn test_storage_envelope_roundtrip() {
let data = b"Hello, World! This is test data for compression.".to_vec();
let envelope = StorageEnvelope::new(&data, "test".to_string()).unwrap();
let extracted = envelope.extract().unwrap();
assert_eq!(data, extracted);
}
#[test]
fn test_compression_works() {
let data = vec![b'a'; LITERAL_BLOCK_THRESHOLD]; let envelope = StorageEnvelope::new(&data, "test".to_string()).unwrap();
assert!(envelope.compressed_data.len() < data.len());
}
#[test]
fn below_threshold_writes_a_literals_only_block() {
for len in [0, 1, 14, 15, LITERAL_BLOCK_THRESHOLD - 1] {
let data = vec![b'a'; len];
let envelope = StorageEnvelope::new(&data, "test".to_string()).unwrap();
let block = &envelope.compressed_data;
let token = block[0];
assert_eq!(token & 0x0F, 0, "len {len}: token has a match length");
let header = if len < 15 {
assert_eq!(usize::from(token >> 4), len, "len {len}: literal length");
1
} else {
assert_eq!(token >> 4, 15, "len {len}: literal length nibble");
assert_eq!(usize::from(block[1]) + 15, len, "len {len}: length byte");
2
};
assert_eq!(&block[header..], &data[..], "len {len}: literals");
assert_eq!(lz4_flex::decompress(block, len).unwrap(), data, "len {len}");
}
}
#[test]
fn at_threshold_writes_lz4_flex_output() {
for data in [
vec![b'a'; LITERAL_BLOCK_THRESHOLD],
(0..LITERAL_BLOCK_THRESHOLD as u8).collect::<Vec<u8>>(),
] {
let envelope = StorageEnvelope::new(&data, "test".to_string()).unwrap();
assert_eq!(envelope.compressed_data, lz4_flex::compress(&data));
}
}
#[test]
fn below_threshold_incompressible_input_matches_lz4_flex() {
for len in 0..LITERAL_BLOCK_THRESHOLD {
let data: Vec<u8> = (0..len as u8).collect();
assert_eq!(
compress_block(&data),
lz4_flex::compress(&data),
"len {len}"
);
}
}
#[test]
fn store_roundtrips_around_and_above_threshold() {
let storage = ByteStorage::new(None);
let mut state: u32 = 0x9E37_79B9;
let mut next = move || {
state ^= state << 13;
state ^= state >> 17;
state ^= state << 5;
state as u8
};
for len in [
LITERAL_BLOCK_THRESHOLD - 1,
LITERAL_BLOCK_THRESHOLD,
LITERAL_BLOCK_THRESHOLD + 1,
100,
1000,
4103,
] {
let inputs = [
vec![b'z'; len],
(0..len).map(|i| (i % 251) as u8).collect::<Vec<u8>>(),
(0..len).map(|_| next()).collect::<Vec<u8>>(),
];
for data in inputs {
let stored = storage.store(&data, None).unwrap();
assert_eq!(storage.retrieve(&stored).unwrap().0, data, "len {len}");
}
}
}
#[test]
fn estimate_compression_reports_what_store_writes() {
let storage = ByteStorage::new(None);
for len in [
0,
32,
LITERAL_BLOCK_THRESHOLD - 1,
LITERAL_BLOCK_THRESHOLD,
1024,
] {
let data = vec![b'a'; len];
let envelope = StorageEnvelope::new(&data, "test".to_string()).unwrap();
let expected = len as f64 / envelope.compressed_data.len() as f64;
assert_eq!(
storage.estimate_compression(&data).unwrap(),
expected,
"len {len}"
);
}
}
#[test]
fn test_checksum_validation() {
let mut envelope = StorageEnvelope::new(b"test", "test".to_string()).unwrap();
envelope.checksum[0] = !envelope.checksum[0];
assert!(matches!(
envelope.extract(),
Err(ByteStorageError::ChecksumMismatch)
));
}
#[test]
fn test_raw_persistence_roundtrip() {
let storage = ByteStorage::new(None);
let test_data = b"test data for persistence";
let stored = storage.store(test_data, None).unwrap();
let (retrieved_data, format) = storage.retrieve(&stored).unwrap();
assert_eq!(test_data, retrieved_data.as_slice());
assert_eq!("msgpack", format);
}
#[test]
fn test_size_limits_input() {
let storage = ByteStorage::new(None);
let large_data = vec![0u8; MAX_UNCOMPRESSED_SIZE + 1];
let result = storage.store(&large_data, None);
assert!(matches!(result, Err(ByteStorageError::InputTooLarge)));
}
#[test]
fn test_size_limits_envelope() {
let max_data = vec![0u8; MAX_UNCOMPRESSED_SIZE];
let envelope_result = StorageEnvelope::new(&max_data, "test".to_string());
assert!(envelope_result.is_ok());
}
#[test]
fn test_compression_ratio_bomb_protection() {
let malicious_envelope = StorageEnvelope {
compressed_data: vec![0u8; 1000], checksum: [0u8; 8], original_size: 200 * 1024 * 1024, format: "test".to_string(),
};
let result = malicious_envelope.extract();
assert!(matches!(result, Err(ByteStorageError::DecompressionBomb)));
}
#[test]
fn test_decompression_bomb_zero_compressed_size() {
let malicious_envelope = StorageEnvelope {
compressed_data: vec![], checksum: [0u8; 8],
original_size: 1000, format: "test".to_string(),
};
let result = malicious_envelope.extract();
assert!(
matches!(result, Err(ByteStorageError::DecompressionBomb)),
"Zero compressed size should be rejected as decompression bomb"
);
}
#[test]
fn test_decompression_bomb_extreme_ratio() {
let malicious_envelope = StorageEnvelope {
compressed_data: vec![0u8; 1], checksum: [0u8; 8],
original_size: 2000, format: "test".to_string(),
};
let result = malicious_envelope.extract();
assert!(
matches!(result, Err(ByteStorageError::DecompressionBomb)),
"Extreme ratio should be rejected as bomb: {:?}",
result
);
}
#[test]
fn test_decompression_u32_max_original_size() {
let malicious_envelope = StorageEnvelope {
compressed_data: vec![0u8; 1000],
checksum: [0u8; 8],
original_size: u32::MAX, format: "test".to_string(),
};
let result = malicious_envelope.extract();
assert!(
matches!(result, Err(ByteStorageError::InputTooLarge)),
"u32::MAX should be rejected as InputTooLarge (exceeds 512MB limit): {:?}",
result
);
}
#[test]
fn test_decompression_exactly_at_threshold() {
let envelope = StorageEnvelope {
compressed_data: vec![0u8; 100], checksum: [0u8; 8],
original_size: 100_000, format: "test".to_string(),
};
let result = envelope.extract();
assert!(
!matches!(result, Err(ByteStorageError::DecompressionBomb)),
"Exactly 1000:1 ratio should pass bomb check: {:?}",
result
);
assert!(
result.is_err(),
"Invalid data should still fail after ratio check"
);
}
#[test]
fn test_decompression_just_over_threshold() {
let malicious_envelope = StorageEnvelope {
compressed_data: vec![0u8; 100], checksum: [0u8; 8],
original_size: 100_001, format: "test".to_string(),
};
let result = malicious_envelope.extract();
assert!(
matches!(result, Err(ByteStorageError::DecompressionBomb)),
"Just over 1000:1 ratio should be rejected as bomb"
);
}
#[test]
fn test_decompression_bomb_integer_boundary() {
let envelope = StorageEnvelope {
compressed_data: vec![0u8; 1_000_000], checksum: [0u8; 8],
original_size: 1_000_000_000, format: "test".to_string(),
};
let result = envelope.extract();
assert!(
matches!(result, Err(ByteStorageError::InputTooLarge)),
"Should fail size check before ratio check: {:?}",
result
);
}
#[test]
fn test_extract_rejects_oversized_compressed_data() {
let envelope = StorageEnvelope {
compressed_data: vec![0u8; MAX_COMPRESSED_SIZE + 1],
checksum: [0u8; 8],
original_size: 1,
format: "test".to_string(),
};
assert_eq!(envelope.extract(), Err(ByteStorageError::InputTooLarge));
}
#[test]
fn test_envelope_size_validation() {
let storage = ByteStorage::new(None);
let oversized_envelope = vec![0u8; MAX_COMPRESSED_SIZE + 1];
let result = storage.retrieve(&oversized_envelope);
assert!(matches!(result, Err(ByteStorageError::InputTooLarge)));
}
#[test]
fn test_security_limits_getters() {
let storage = ByteStorage::new(None);
assert_eq!(storage.max_uncompressed_size(), MAX_UNCOMPRESSED_SIZE);
assert_eq!(storage.max_compressed_size(), MAX_COMPRESSED_SIZE);
assert_eq!(storage.max_compression_ratio(), 1000u64);
}
#[test]
fn test_compression_estimate_security() {
let storage = ByteStorage::new(None);
let large_data = vec![0u8; MAX_UNCOMPRESSED_SIZE + 1];
let result = storage.estimate_compression(&large_data);
assert!(matches!(result, Err(ByteStorageError::InputTooLarge)));
}
#[test]
fn test_validate_security() {
let storage = ByteStorage::new(None);
let large_envelope = vec![0u8; MAX_COMPRESSED_SIZE + 1];
let result = storage.validate(&large_envelope);
assert!(!result); }
#[test]
fn test_edge_case_exactly_at_limits() {
let storage = ByteStorage::new(None);
let max_size_data = vec![1u8; MAX_UNCOMPRESSED_SIZE];
let result = storage.store(&max_size_data, None);
assert!(result.is_ok());
}
fn decode_corpus() -> Vec<(String, Vec<u8>)> {
let fixture: serde_json::Value =
serde_json::from_str(include_str!("../tests/vectors/wire-format.json")).unwrap();
let mut corpus = Vec::new();
for group in ["vectors", "reject_vectors", "payload_reject_vectors"] {
for v in fixture[group].as_array().unwrap() {
let hex = v["envelope_hex"].as_str().unwrap();
corpus.push((
v["name"].as_str().unwrap().to_string(),
hex::decode(hex).unwrap(),
));
}
}
let checksum = [0x98, 0, 1, 2, 3, 4, 5, 6, 0xCC, 0xFF];
for (name, payload) in [
("payload_str", &[0xA3, b'a', b'b', b'c'][..]),
("payload_int", &[0x07][..]),
("payload_nil", &[0xC0][..]),
("payload_map", &[0x80][..]),
("payload_bool", &[0xC3][..]),
("payload_ext", &[0xD4, 0x01, 0x00][..]),
("payload_float", &[0xCA, 0, 0, 0, 0][..]),
("legacy_bad_element", &[0x92, 0x01, 0xA1, b'x'][..]),
("payload_truncated_bin", &[0xC4, 0x09, 0x00][..]),
] {
let mut env = vec![0x94];
env.extend_from_slice(payload);
env.extend_from_slice(&checksum);
env.extend_from_slice(&[0x03, 0xA4, b't', b'e', b's', b't']);
corpus.push((name.to_string(), env));
}
let mut named =
rmp_serde::to_vec_named(&StorageEnvelope::new(b"named", "f".into()).unwrap()).unwrap();
corpus.push(("map_form".into(), named.clone()));
named[0] = 0x83;
corpus.push(("map_form_missing_field".into(), named));
corpus.push(("not_an_array".into(), vec![0xA2, b'h', b'i']));
corpus.push(("empty".into(), vec![]));
corpus
}
#[test]
fn borrowed_decode_matches_owned_decode() {
for (name, bytes) in decode_corpus() {
let owned = rmp_serde::from_slice::<StorageEnvelope>(&bytes);
let borrowed = rmp_serde::from_slice::<borrowed::StorageEnvelope>(&bytes);
match (owned, borrowed) {
(Ok(o), Ok(b)) => {
assert_eq!(o.compressed_data, &*b.compressed_data, "[{name}] payload");
assert_eq!(o.checksum, b.checksum, "[{name}] checksum");
assert_eq!(o.original_size, b.original_size, "[{name}] original_size");
assert_eq!(o.format, b.format, "[{name}] format");
}
(Err(o), Err(b)) => assert_eq!(o.to_string(), b.to_string(), "[{name}] error"),
(o, b) => panic!("[{name}] owned {:?} vs borrowed {:?}", o.err(), b.err()),
}
}
}
#[test]
fn bin_payload_is_borrowed_and_legacy_payload_is_owned() {
let fixture: serde_json::Value =
serde_json::from_str(include_str!("../tests/vectors/wire-format.json")).unwrap();
use std::borrow::Cow;
let mut seen = (0, 0);
for v in fixture["vectors"].as_array().unwrap() {
let bytes = hex::decode(v["envelope_hex"].as_str().unwrap()).unwrap();
let envelope = decode_envelope(&bytes).unwrap();
let name = v["name"].as_str().unwrap();
if name.ends_with("_bin") {
seen.0 += 1;
assert!(
matches!(envelope.compressed_data, Cow::Borrowed(_)),
"[{name}] bin payload copied"
);
} else {
seen.1 += 1;
assert!(
matches!(envelope.compressed_data, Cow::Owned(_)),
"[{name}] legacy payload"
);
}
}
assert!(
seen.0 > 0 && seen.1 > 0,
"fixture has both encodings: {seen:?}"
);
let storage = ByteStorage::new(None);
for (len, borrowed) in [(LARGE_PAYLOAD - 512, true), (LARGE_PAYLOAD + 512, false)] {
let data = xorshift_bytes(len);
let stored = storage.store(&data, None).unwrap();
let envelope = decode_envelope(&stored).unwrap();
let payload = envelope.compressed_data.len();
assert_eq!(
payload < LARGE_PAYLOAD,
borrowed,
"len {len}: payload {payload}"
);
assert_eq!(
matches!(envelope.compressed_data, Cow::Borrowed(_)),
borrowed,
"len {len}: payload {payload}"
);
assert_eq!(storage.retrieve(&stored).unwrap().0, data, "len {len}");
}
}
fn xorshift_bytes(len: usize) -> Vec<u8> {
let mut state: u32 = 0x9E37_79B9;
(0..len)
.map(|_| {
state ^= state << 13;
state ^= state >> 17;
state ^= state << 5;
state as u8
})
.collect()
}
fn envelope_with(
payload_len: usize,
checksum: [u8; 8],
size: u32,
format_len: usize,
) -> StorageEnvelope {
StorageEnvelope {
compressed_data: vec![0; payload_len],
checksum,
original_size: size,
format: "f".repeat(format_len),
}
}
#[test]
fn encoded_len_matches_rmp_serde() {
let lengths = [0, 1, 31, 32, 255, 256, 65_535, 65_536, 70_000];
let sizes = [0, 127, 128, 255, 256, 65_535, 65_536, u32::MAX];
let checksums = [
[0x00; 8],
[0x7F; 8],
[0x80; 8],
[0xFF; 8],
[0, 0x80, 1, 0xFF, 0x7F, 2, 0xC0, 3],
];
for payload_len in lengths {
for format_len in lengths {
let envelope = envelope_with(payload_len, checksums[4], 1, format_len);
let expected = rmp_serde::to_vec(&envelope).unwrap().len();
assert_eq!(
envelope.encoded_len(),
expected,
"payload {payload_len}, format {format_len}"
);
}
}
for size in sizes {
for checksum in checksums {
let envelope = envelope_with(3, checksum, size, 7);
let expected = rmp_serde::to_vec(&envelope).unwrap().len();
assert_eq!(
envelope.encoded_len(),
expected,
"size {size}, checksum {checksum:?}"
);
}
}
}
#[test]
fn store_and_encode_into_write_rmp_serde_bytes() {
let storage = ByteStorage::new(None);
for (data, large) in [
(vec![], false),
(vec![b'a'; 63], false),
(vec![b'a'; 64], false),
((0..70_000).map(|i| (i * 7919 % 251) as u8).collect(), false),
(xorshift_bytes(LARGE_PAYLOAD - 512), false),
(xorshift_bytes(LARGE_PAYLOAD), true),
(xorshift_bytes(300_000), true),
] {
let envelope = StorageEnvelope::new(&data, "msgpack".to_string()).unwrap();
let expected = rmp_serde::to_vec(&envelope).unwrap();
let stored = storage.store(&data, None).unwrap();
assert_eq!(stored, expected, "store, len {}", data.len());
assert_eq!(
stored.len() >= LARGE_PAYLOAD,
large,
"gate side, len {}: envelope {}",
data.len(),
stored.len()
);
if stored.len() < LARGE_PAYLOAD {
assert_eq!(
stored.capacity(),
stored.len(),
"store capacity, len {}",
data.len()
);
} else {
assert_eq!(
stored.capacity(),
expected.capacity(),
"store capacity, len {}",
data.len()
);
}
let mut out = vec![0u8; envelope.encoded_len()];
envelope.encode_into(&mut out).unwrap();
assert_eq!(out, expected, "encode_into, len {}", data.len());
}
}
#[test]
fn encode_into_rejects_a_wrong_sized_buffer() {
let envelope = StorageEnvelope::new(b"payload", "msgpack".to_string()).unwrap();
let len = envelope.encoded_len();
for wrong in [0, len - 1, len + 1] {
let mut out = vec![0u8; wrong];
assert!(
matches!(
envelope.encode_into(&mut out),
Err(ByteStorageError::SerializationFailed(_))
),
"buffer of {wrong} bytes for a {len}-byte envelope"
);
}
}
#[test]
fn encode_into_applies_the_envelope_size_cap() {
let envelope = envelope_with(MAX_COMPRESSED_SIZE, [0; 8], 1, 7);
assert!(envelope.encoded_len() > MAX_COMPRESSED_SIZE);
assert_eq!(
envelope.encode_into(&mut []),
Err(ByteStorageError::InputTooLarge)
);
assert_eq!(
envelope.encode().err(),
Some(ByteStorageError::InputTooLarge)
);
}
#[test]
fn test_zero_size_edge_case() {
let storage = ByteStorage::new(None);
let empty_data = vec![];
let stored = storage.store(&empty_data, None).unwrap();
let (retrieved_data, format) = storage.retrieve(&stored).unwrap();
assert_eq!(empty_data, retrieved_data);
assert_eq!("msgpack", format);
}
}
#[cfg(kani)]
mod kani_proofs {
use super::*;
#[kani::proof]
#[kani::unwind(10)] fn verify_checksum_detects_corruption() {
let checksum_a: [u8; 8] = kani::any();
let mut checksum_b = checksum_a;
let byte_index: usize = kani::any();
let bit_index: usize = kani::any();
kani::assume(byte_index < 8);
kani::assume(bit_index < 8);
checksum_b[byte_index] ^= 1 << bit_index;
assert_ne!(checksum_a, checksum_b);
}
const PINNED_MAX_COMPRESSED: u64 = 512 * 1024 * 1024;
const PINNED_MAX_UNCOMPRESSED: u64 = 512 * 1024 * 1024;
const PINNED_MAX_RATIO: u64 = 1000;
#[kani::proof]
fn verify_decompression_bound_size_caps() {
let compressed_len: usize = kani::any();
let original_size: u32 = kani::any();
let over_cap = compressed_len as u64 > PINNED_MAX_COMPRESSED
|| original_size as u64 > PINNED_MAX_UNCOMPRESSED;
let result = check_decompression_bound(compressed_len, original_size);
assert_eq!(
over_cap,
matches!(result, Err(ByteStorageError::InputTooLarge))
);
}
#[kani::proof]
fn verify_decompression_bound_ratio() {
let compressed_len: usize = kani::any();
let original_size: u32 = kani::any();
kani::assume(compressed_len as u64 <= PINNED_MAX_COMPRESSED);
kani::assume(original_size as u64 <= PINNED_MAX_UNCOMPRESSED);
let is_bomb =
compressed_len == 0 || original_size as u64 > PINNED_MAX_RATIO * compressed_len as u64;
let result = check_decompression_bound(compressed_len, original_size);
if is_bomb {
assert!(matches!(result, Err(ByteStorageError::DecompressionBomb)));
} else {
assert!(result.is_ok());
}
}
}