use std::io::Write;
use crate::error::{CompressError, DecompressError};
use crate::{CHUNK_SIZE, STREAM_BUF_SIZE};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressionFormat {
None,
Zrip,
}
impl CompressionFormat {
pub fn as_str(self) -> &'static str {
match self {
CompressionFormat::None => "identity",
CompressionFormat::Zrip => "zrip",
}
}
pub fn parse_header(s: &str) -> Option<Self> {
match s.trim().to_ascii_lowercase().as_str() {
"" | "identity" | "none" => Some(CompressionFormat::None),
"zrip" | "zstd" => Some(CompressionFormat::Zrip),
_ => None,
}
}
}
pub const MAX_FRAME_OUTPUT: usize = CHUNK_SIZE as usize;
pub const MAX_PENDING_FRAME: usize = CHUNK_SIZE as usize + STREAM_BUF_SIZE;
pub const MAX_OUTPUT_PER_CALL: usize = MAX_FRAME_OUTPUT.saturating_mul(8);
pub const ZRIP_DEFAULT_LEVEL: i32 = 1;
pub trait Compressor: Send {
fn format(&self) -> CompressionFormat;
fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError>;
fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError>;
}
pub trait Decompressor: Send {
fn format(&self) -> CompressionFormat;
fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError>;
fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError>;
}
pub fn compressor(format: CompressionFormat) -> Result<Box<dyn Compressor>, CompressError> {
match format {
CompressionFormat::None => Ok(Box::new(PassthroughCompressor)),
CompressionFormat::Zrip => Ok(Box::new(ZripCompressor::new(ZRIP_DEFAULT_LEVEL)?)),
}
}
pub fn decompressor(format: CompressionFormat) -> Box<dyn Decompressor> {
decompressor_with_limit(format, MAX_OUTPUT_PER_CALL)
}
pub fn decompressor_with_limit(
format: CompressionFormat,
max_output_per_call: usize,
) -> Box<dyn Decompressor> {
match format {
CompressionFormat::None => Box::new(PassthroughDecompressor),
CompressionFormat::Zrip => Box::new(ZripDecompressor::with_max_output(max_output_per_call)),
}
}
struct PassthroughCompressor;
impl Compressor for PassthroughCompressor {
fn format(&self) -> CompressionFormat {
CompressionFormat::None
}
fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError> {
out.extend_from_slice(input);
Ok(())
}
fn finish(&mut self, _out: &mut Vec<u8>) -> Result<(), CompressError> {
Ok(())
}
}
struct PassthroughDecompressor;
impl Decompressor for PassthroughDecompressor {
fn format(&self) -> CompressionFormat {
CompressionFormat::None
}
fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError> {
out.extend_from_slice(input);
Ok(())
}
fn finish(&mut self, _out: &mut Vec<u8>) -> Result<(), DecompressError> {
Ok(())
}
}
pub struct ZripCompressor {
encoder: Option<zrip::FrameEncoder<Vec<u8>>>,
dirty: bool,
}
impl ZripCompressor {
pub fn new(level: i32) -> Result<Self, CompressError> {
Ok(ZripCompressor {
encoder: Some(zrip::FrameEncoder::new(Vec::new(), level)?),
dirty: false,
})
}
}
impl Compressor for ZripCompressor {
fn format(&self) -> CompressionFormat {
CompressionFormat::Zrip
}
fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError> {
if input.is_empty() {
return Ok(());
}
let encoder = self
.encoder
.as_mut()
.ok_or_else(|| std::io::Error::other("compressor already finished"))?;
if self.dirty {
let finished = encoder.reset(Vec::new())?;
out.extend_from_slice(&finished);
}
encoder.write_all(input)?;
self.dirty = true;
Ok(())
}
fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError> {
let encoder = self
.encoder
.take()
.ok_or_else(|| std::io::Error::other("compressor already finished"))?;
let tail = encoder.finish()?;
out.extend_from_slice(&tail);
self.dirty = false;
Ok(())
}
}
pub struct ZripDecompressor {
pending: Vec<u8>,
finished: bool,
max_output_per_call: usize,
}
impl ZripDecompressor {
pub fn new() -> Self {
ZripDecompressor::with_max_output(MAX_OUTPUT_PER_CALL)
}
pub fn with_max_output(max_output_per_call: usize) -> Self {
ZripDecompressor {
pending: Vec::with_capacity(STREAM_BUF_SIZE),
finished: false,
max_output_per_call,
}
}
fn drain_frames(&mut self, out: &mut Vec<u8>, max_add: usize) -> Result<(), DecompressError> {
let mut added = 0usize;
loop {
let boundary = frame_boundary(&self.pending);
match boundary {
FrameBoundary::Complete { len, content_size } => {
if content_size.is_some_and(|cs| cs > MAX_FRAME_OUTPUT as u64) {
return Err(DecompressError::TooLarge {
limit: MAX_FRAME_OUTPUT,
});
}
let decoded = {
let frame = &self.pending[..len];
zrip::decompress_with_limit(frame, MAX_FRAME_OUTPUT).map_err(|e| {
DecompressError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
e,
))
})?
};
if added.saturating_add(decoded.len()) > max_add {
return Err(DecompressError::TooLarge { limit: max_add });
}
out.extend_from_slice(&decoded);
added = added.saturating_add(decoded.len());
self.pending.drain(..len);
}
FrameBoundary::Incomplete => return Ok(()),
FrameBoundary::Invalid => {
return Err(DecompressError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"invalid zstd frame data",
)))
}
}
}
}
}
impl Default for ZripDecompressor {
fn default() -> Self {
ZripDecompressor::new()
}
}
impl Decompressor for ZripDecompressor {
fn format(&self) -> CompressionFormat {
CompressionFormat::Zrip
}
fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError> {
if self.finished {
return Err(DecompressError::Io(std::io::Error::other(
"decompressor already finished",
)));
}
if !input.is_empty() {
self.pending.extend_from_slice(input);
}
self.drain_frames(out, self.max_output_per_call)?;
if self.pending.len() > MAX_PENDING_FRAME {
return Err(DecompressError::TooLarge {
limit: MAX_PENDING_FRAME,
});
}
Ok(())
}
fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError> {
if self.finished {
return Ok(());
}
self.finished = true;
self.drain_frames(out, self.max_output_per_call)?;
if !self.pending.is_empty() {
return Err(DecompressError::Truncated(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
format!("{} trailing compressed bytes", self.pending.len()),
)));
}
Ok(())
}
}
const ZSTD_MAGIC: u32 = 0xFD2F_B528;
const SKIPPABLE_MASK: u32 = 0xFFFF_FFF0;
const SKIPPABLE_MAGIC: u32 = 0x184D_2A50;
enum FrameBoundary {
Complete { len: usize, content_size: Option<u64> },
Incomplete,
Invalid,
}
fn frame_boundary(buf: &[u8]) -> FrameBoundary {
if buf.len() < 4 {
return FrameBoundary::Incomplete;
}
let magic = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]);
if magic == ZSTD_MAGIC {
zstd_frame_boundary(buf)
} else if (magic & SKIPPABLE_MASK) == SKIPPABLE_MAGIC {
skippable_frame_boundary(buf)
} else {
FrameBoundary::Invalid
}
}
fn skippable_frame_boundary(buf: &[u8]) -> FrameBoundary {
if buf.len() < 8 {
return FrameBoundary::Incomplete;
}
let skip = u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]) as usize;
let total = 8usize.saturating_add(skip);
if buf.len() >= total {
FrameBoundary::Complete {
len: total,
content_size: Some(0),
}
} else {
FrameBoundary::Incomplete
}
}
fn zstd_frame_boundary(buf: &[u8]) -> FrameBoundary {
if buf.len() < 5 {
return FrameBoundary::Incomplete;
}
let descriptor = buf[4];
if descriptor & 0x18 != 0 {
return FrameBoundary::Invalid;
}
let single_segment = descriptor & 0x20 != 0;
let checksum = descriptor & 0x04 != 0;
let dict_id_flag = descriptor & 0x03;
let fcs_flag = (descriptor >> 6) & 0x03;
let mut hdr_len = 5usize;
if !single_segment {
hdr_len += 1; }
hdr_len += match dict_id_flag {
0 => 0,
1 => 1,
2 => 2,
3 => 4,
_ => unreachable!(),
};
let fcs_size: usize = match fcs_flag {
0 if single_segment => 1,
0 => 0,
1 => 2,
2 => 4,
3 => 8,
_ => unreachable!(),
};
hdr_len += fcs_size;
if buf.len() < hdr_len {
return FrameBoundary::Incomplete;
}
let content_size = if fcs_size > 0 {
let mut v = 0u64;
for (i, &b) in buf[5..5 + fcs_size].iter().enumerate() {
v |= (b as u64) << (8 * i);
}
Some(v)
} else {
None
};
let mut off = hdr_len;
loop {
if buf.len() < off + 3 {
return FrameBoundary::Incomplete;
}
let block_header = u32::from_le_bytes([buf[off], buf[off + 1], buf[off + 2], 0]);
let last = block_header & 0x01 != 0;
let block_type = (block_header >> 1) & 0x03;
let block_size = (block_header >> 3) as usize;
if block_type == 3 {
return FrameBoundary::Invalid;
}
off += 3 + block_size;
if off > MAX_PENDING_FRAME {
return FrameBoundary::Invalid;
}
if last {
break;
}
}
if checksum {
off += 4;
}
if buf.len() < off {
return FrameBoundary::Incomplete;
}
FrameBoundary::Complete {
len: off,
content_size,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn roundtrip_chunks(data: &[u8], feed: &[usize]) {
let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
let mut compressed = Vec::new();
let mut off = 0;
for &n in feed {
let end = (off + n).min(data.len());
if end > off {
c.compress(&data[off..end], &mut compressed).unwrap();
}
off = end;
}
c.finish(&mut compressed).unwrap();
assert!(!compressed.is_empty());
let mut d = ZripDecompressor::new();
let mut plain = Vec::new();
let mut step = 1;
let mut i = 0;
while i < compressed.len() {
let end = (i + step).min(compressed.len());
d.decompress(&compressed[i..end], &mut plain).unwrap();
i = end;
step = step % 5 + 1; }
d.finish(&mut plain).unwrap();
assert_eq!(plain, data, "roundtrip mismatch with feed {feed:?}");
}
#[test]
fn roundtrip_single_chunk() {
let data: Vec<u8> = (0..100_000u32).map(|i| (i % 251) as u8).collect();
roundtrip_chunks(&data, &[data.len()]);
}
#[test]
fn roundtrip_multi_chunk_64k_windows() {
let data: Vec<u8> = (0..300_000u32).map(|i| (i / 7) as u8).collect();
let feed: Vec<usize> = std::iter::repeat(STREAM_BUF_SIZE).take(5).collect();
roundtrip_chunks(&data, &feed);
}
#[test]
fn roundtrip_highly_compressible() {
let data = b"libfw streaming compression test. ".repeat(10_000);
let mut feed = Vec::new();
let mut consumed = 0;
for (_, &size) in [1024usize, 2048, 4096, 8192, 16384].iter().cycle().enumerate() {
feed.push(size);
consumed += size;
if consumed >= data.len() {
break;
}
}
roundtrip_chunks(&data, &feed);
}
#[test]
fn roundtrip_empty_stream() {
let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
let mut compressed = Vec::new();
c.finish(&mut compressed).unwrap();
let mut d = ZripDecompressor::new();
let mut plain = Vec::new();
d.decompress(&compressed, &mut plain).unwrap();
d.finish(&mut plain).unwrap();
assert!(plain.is_empty());
}
#[test]
fn empty_input_chunks_are_noops() {
let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
let mut out = Vec::new();
c.compress(&[], &mut out).unwrap();
c.compress(b"hello", &mut out).unwrap();
c.finish(&mut out).unwrap();
let mut d = ZripDecompressor::new();
let mut plain = Vec::new();
d.decompress(&out, &mut plain).unwrap();
d.finish(&mut plain).unwrap();
assert_eq!(plain, b"hello");
}
#[test]
fn format_header_roundtrip() {
assert_eq!(CompressionFormat::parse_header("zrip"), Some(CompressionFormat::Zrip));
assert_eq!(CompressionFormat::parse_header("ZSTD"), Some(CompressionFormat::Zrip));
assert_eq!(CompressionFormat::parse_header("identity"), Some(CompressionFormat::None));
assert_eq!(CompressionFormat::parse_header(""), Some(CompressionFormat::None));
assert_eq!(CompressionFormat::parse_header("br"), None);
assert_eq!(CompressionFormat::Zrip.as_str(), "zrip");
assert_eq!(CompressionFormat::None.as_str(), "identity");
}
#[test]
fn passthrough_roundtrip() {
let mut c = compressor(CompressionFormat::None).unwrap();
let mut d = decompressor(CompressionFormat::None);
let mut compressed = Vec::new();
let mut plain = Vec::new();
c.compress(b"abc", &mut compressed).unwrap();
c.finish(&mut compressed).unwrap();
d.decompress(&compressed, &mut plain).unwrap();
d.finish(&mut plain).unwrap();
assert_eq!(plain, b"abc");
}
#[test]
fn truncated_stream_is_detected() {
let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
let mut compressed = Vec::new();
c.compress(&vec![7u8; 5000], &mut compressed).unwrap();
c.finish(&mut compressed).unwrap();
compressed.truncate(compressed.len() - 1);
let mut d = ZripDecompressor::new();
let mut plain = Vec::new();
d.decompress(&compressed, &mut plain).unwrap();
assert!(matches!(d.finish(&mut plain), Err(DecompressError::Truncated(_))));
}
#[test]
fn corrupt_stream_is_detected() {
let mut d = ZripDecompressor::new();
let mut plain = Vec::new();
let err = d.decompress(b"this is not a zstd frame at all", &mut plain);
assert!(err.is_err());
}
#[test]
fn per_call_output_budget_rejects_multi_frame_bomb() {
let mut compressed = Vec::new();
for _ in 0..64 {
let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
c.compress(&vec![7u8; STREAM_BUF_SIZE], &mut compressed)
.unwrap();
c.finish(&mut compressed).unwrap();
}
let mut d = ZripDecompressor::with_max_output(MAX_FRAME_OUTPUT);
let mut plain = Vec::new();
let err = d.decompress(&compressed, &mut plain);
assert!(
matches!(err, Err(DecompressError::TooLarge { .. })),
"expected TooLarge, got {err:?}"
);
assert!(plain.len() <= MAX_FRAME_OUTPUT);
}
#[test]
fn generous_default_budget_allows_coalesced_frames() {
let mut compressed = Vec::new();
for _ in 0..8 {
let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
c.compress(&vec![7u8; STREAM_BUF_SIZE], &mut compressed)
.unwrap();
c.finish(&mut compressed).unwrap();
}
let mut d = ZripDecompressor::new();
let mut plain = Vec::new();
d.decompress(&compressed, &mut plain).unwrap();
d.finish(&mut plain).unwrap();
assert_eq!(plain.len(), 8 * STREAM_BUF_SIZE);
}
#[test]
fn zstd_compat_interop() {
let data: Vec<u8> = (0..50_000u32).map(|i| (i % 31) as u8).collect();
let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
let mut compressed = Vec::new();
c.compress(&data, &mut compressed).unwrap();
c.finish(&mut compressed).unwrap();
let decoded = zstd::stream::decode_all(&compressed[..]).unwrap();
assert_eq!(decoded, data);
let zstd_enc = zstd::stream::encode_all(&data[..], 1).unwrap();
let mut d = ZripDecompressor::new();
let mut plain = Vec::new();
d.decompress(&zstd_enc, &mut plain).unwrap();
d.finish(&mut plain).unwrap();
assert_eq!(plain, data);
}
}