1use std::io::Write;
37
38use crate::error::{CompressError, DecompressError};
39use crate::{CHUNK_SIZE, STREAM_BUF_SIZE};
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum CompressionFormat {
44 None,
46 Zrip,
48}
49
50impl CompressionFormat {
51 pub fn as_str(self) -> &'static str {
54 match self {
55 CompressionFormat::None => "identity",
56 CompressionFormat::Zrip => "zrip",
57 }
58 }
59
60 pub fn parse_header(s: &str) -> Option<Self> {
62 match s.trim().to_ascii_lowercase().as_str() {
63 "" | "identity" | "none" => Some(CompressionFormat::None),
64 "zrip" | "zstd" => Some(CompressionFormat::Zrip),
65 _ => None,
66 }
67 }
68}
69
70pub const MAX_FRAME_OUTPUT: usize = CHUNK_SIZE as usize;
75
76pub const MAX_PENDING_FRAME: usize = CHUNK_SIZE as usize + STREAM_BUF_SIZE;
81
82pub const MAX_OUTPUT_PER_CALL: usize = MAX_FRAME_OUTPUT.saturating_mul(8);
92
93pub const ZRIP_DEFAULT_LEVEL: i32 = 1;
95
96pub trait Compressor: Send {
99 fn format(&self) -> CompressionFormat;
101
102 fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError>;
108
109 fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError>;
113}
114
115pub trait Decompressor: Send {
118 fn format(&self) -> CompressionFormat;
120
121 fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError>;
126
127 fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError>;
132}
133
134pub fn compressor(format: CompressionFormat) -> Result<Box<dyn Compressor>, CompressError> {
136 match format {
137 CompressionFormat::None => Ok(Box::new(PassthroughCompressor)),
138 CompressionFormat::Zrip => Ok(Box::new(ZripCompressor::new(ZRIP_DEFAULT_LEVEL)?)),
139 }
140}
141
142pub fn decompressor(format: CompressionFormat) -> Box<dyn Decompressor> {
144 decompressor_with_limit(format, MAX_OUTPUT_PER_CALL)
145}
146
147pub fn decompressor_with_limit(
154 format: CompressionFormat,
155 max_output_per_call: usize,
156) -> Box<dyn Decompressor> {
157 match format {
158 CompressionFormat::None => Box::new(PassthroughDecompressor),
159 CompressionFormat::Zrip => Box::new(ZripDecompressor::with_max_output(max_output_per_call)),
160 }
161}
162
163struct PassthroughCompressor;
168
169impl Compressor for PassthroughCompressor {
170 fn format(&self) -> CompressionFormat {
171 CompressionFormat::None
172 }
173
174 fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError> {
175 out.extend_from_slice(input);
176 Ok(())
177 }
178
179 fn finish(&mut self, _out: &mut Vec<u8>) -> Result<(), CompressError> {
180 Ok(())
181 }
182}
183
184struct PassthroughDecompressor;
185
186impl Decompressor for PassthroughDecompressor {
187 fn format(&self) -> CompressionFormat {
188 CompressionFormat::None
189 }
190
191 fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError> {
192 out.extend_from_slice(input);
193 Ok(())
194 }
195
196 fn finish(&mut self, _out: &mut Vec<u8>) -> Result<(), DecompressError> {
197 Ok(())
198 }
199}
200
201pub struct ZripCompressor {
207 encoder: Option<zrip::FrameEncoder<Vec<u8>>>,
208 dirty: bool,
210}
211
212impl ZripCompressor {
213 pub fn new(level: i32) -> Result<Self, CompressError> {
215 Ok(ZripCompressor {
216 encoder: Some(zrip::FrameEncoder::new(Vec::new(), level)?),
217 dirty: false,
218 })
219 }
220}
221
222impl Compressor for ZripCompressor {
223 fn format(&self) -> CompressionFormat {
224 CompressionFormat::Zrip
225 }
226
227 fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError> {
228 if input.is_empty() {
229 return Ok(());
230 }
231 let encoder = self
232 .encoder
233 .as_mut()
234 .ok_or_else(|| std::io::Error::other("compressor already finished"))?;
235 if self.dirty {
237 let finished = encoder.reset(Vec::new())?;
238 out.extend_from_slice(&finished);
239 }
240 encoder.write_all(input)?;
241 self.dirty = true;
242 Ok(())
243 }
244
245 fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError> {
246 let encoder = self
247 .encoder
248 .take()
249 .ok_or_else(|| std::io::Error::other("compressor already finished"))?;
250 let tail = encoder.finish()?;
251 out.extend_from_slice(&tail);
252 self.dirty = false;
253 Ok(())
254 }
255}
256
257pub struct ZripDecompressor {
259 pending: Vec<u8>,
261 finished: bool,
262 max_output_per_call: usize,
265}
266
267impl ZripDecompressor {
268 pub fn new() -> Self {
271 ZripDecompressor::with_max_output(MAX_OUTPUT_PER_CALL)
272 }
273
274 pub fn with_max_output(max_output_per_call: usize) -> Self {
276 ZripDecompressor {
277 pending: Vec::with_capacity(STREAM_BUF_SIZE),
278 finished: false,
279 max_output_per_call,
280 }
281 }
282
283 fn drain_frames(&mut self, out: &mut Vec<u8>, max_add: usize) -> Result<(), DecompressError> {
286 let mut added = 0usize;
287 loop {
288 let boundary = frame_boundary(&self.pending);
289 match boundary {
290 FrameBoundary::Complete { len, content_size } => {
291 if content_size.is_some_and(|cs| cs > MAX_FRAME_OUTPUT as u64) {
292 return Err(DecompressError::TooLarge {
293 limit: MAX_FRAME_OUTPUT,
294 });
295 }
296 let decoded = {
297 let frame = &self.pending[..len];
298 zrip::decompress_with_limit(frame, MAX_FRAME_OUTPUT).map_err(|e| {
299 DecompressError::Io(std::io::Error::new(
300 std::io::ErrorKind::InvalidData,
301 e,
302 ))
303 })?
304 };
305 if added.saturating_add(decoded.len()) > max_add {
308 return Err(DecompressError::TooLarge { limit: max_add });
309 }
310 out.extend_from_slice(&decoded);
311 added = added.saturating_add(decoded.len());
312 self.pending.drain(..len);
313 }
314 FrameBoundary::Incomplete => return Ok(()),
315 FrameBoundary::Invalid => {
316 return Err(DecompressError::Io(std::io::Error::new(
317 std::io::ErrorKind::InvalidData,
318 "invalid zstd frame data",
319 )))
320 }
321 }
322 }
323 }
324}
325
326impl Default for ZripDecompressor {
327 fn default() -> Self {
328 ZripDecompressor::new()
329 }
330}
331
332impl Decompressor for ZripDecompressor {
333 fn format(&self) -> CompressionFormat {
334 CompressionFormat::Zrip
335 }
336
337 fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError> {
338 if self.finished {
339 return Err(DecompressError::Io(std::io::Error::other(
340 "decompressor already finished",
341 )));
342 }
343 if !input.is_empty() {
344 self.pending.extend_from_slice(input);
345 }
346 self.drain_frames(out, self.max_output_per_call)?;
347 if self.pending.len() > MAX_PENDING_FRAME {
348 return Err(DecompressError::TooLarge {
349 limit: MAX_PENDING_FRAME,
350 });
351 }
352 Ok(())
353 }
354
355 fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError> {
356 if self.finished {
357 return Ok(());
358 }
359 self.finished = true;
360 self.drain_frames(out, self.max_output_per_call)?;
361 if !self.pending.is_empty() {
362 return Err(DecompressError::Truncated(std::io::Error::new(
363 std::io::ErrorKind::UnexpectedEof,
364 format!("{} trailing compressed bytes", self.pending.len()),
365 )));
366 }
367 Ok(())
368 }
369}
370
371const ZSTD_MAGIC: u32 = 0xFD2F_B528;
376const SKIPPABLE_MASK: u32 = 0xFFFF_FFF0;
377const SKIPPABLE_MAGIC: u32 = 0x184D_2A50;
378
379enum FrameBoundary {
380 Complete { len: usize, content_size: Option<u64> },
383 Incomplete,
385 Invalid,
387}
388
389fn frame_boundary(buf: &[u8]) -> FrameBoundary {
391 if buf.len() < 4 {
392 return FrameBoundary::Incomplete;
393 }
394 let magic = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]);
395 if magic == ZSTD_MAGIC {
396 zstd_frame_boundary(buf)
397 } else if (magic & SKIPPABLE_MASK) == SKIPPABLE_MAGIC {
398 skippable_frame_boundary(buf)
399 } else {
400 FrameBoundary::Invalid
401 }
402}
403
404fn skippable_frame_boundary(buf: &[u8]) -> FrameBoundary {
406 if buf.len() < 8 {
407 return FrameBoundary::Incomplete;
408 }
409 let skip = u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]) as usize;
410 let total = 8usize.saturating_add(skip);
411 if buf.len() >= total {
412 FrameBoundary::Complete {
413 len: total,
414 content_size: Some(0),
415 }
416 } else {
417 FrameBoundary::Incomplete
418 }
419}
420
421fn zstd_frame_boundary(buf: &[u8]) -> FrameBoundary {
426 if buf.len() < 5 {
427 return FrameBoundary::Incomplete;
428 }
429 let descriptor = buf[4];
430 if descriptor & 0x18 != 0 {
432 return FrameBoundary::Invalid;
433 }
434 let single_segment = descriptor & 0x20 != 0;
435 let checksum = descriptor & 0x04 != 0;
436 let dict_id_flag = descriptor & 0x03;
437 let fcs_flag = (descriptor >> 6) & 0x03;
438
439 let mut hdr_len = 5usize;
440 if !single_segment {
441 hdr_len += 1; }
443 hdr_len += match dict_id_flag {
444 0 => 0,
445 1 => 1,
446 2 => 2,
447 3 => 4,
448 _ => unreachable!(),
449 };
450 let fcs_size: usize = match fcs_flag {
451 0 if single_segment => 1,
452 0 => 0,
453 1 => 2,
454 2 => 4,
455 3 => 8,
456 _ => unreachable!(),
457 };
458 hdr_len += fcs_size;
459
460 if buf.len() < hdr_len {
461 return FrameBoundary::Incomplete;
462 }
463 let content_size = if fcs_size > 0 {
464 let mut v = 0u64;
465 for (i, &b) in buf[5..5 + fcs_size].iter().enumerate() {
466 v |= (b as u64) << (8 * i);
467 }
468 Some(v)
469 } else {
470 None
471 };
472
473 let mut off = hdr_len;
475 loop {
476 if buf.len() < off + 3 {
477 return FrameBoundary::Incomplete;
478 }
479 let block_header = u32::from_le_bytes([buf[off], buf[off + 1], buf[off + 2], 0]);
482 let last = block_header & 0x01 != 0;
483 let block_type = (block_header >> 1) & 0x03;
484 let block_size = (block_header >> 3) as usize;
485 if block_type == 3 {
486 return FrameBoundary::Invalid;
488 }
489 off += 3 + block_size;
490 if off > MAX_PENDING_FRAME {
491 return FrameBoundary::Invalid;
492 }
493 if last {
494 break;
495 }
496 }
497 if checksum {
498 off += 4;
499 }
500 if buf.len() < off {
501 return FrameBoundary::Incomplete;
502 }
503 FrameBoundary::Complete {
504 len: off,
505 content_size,
506 }
507}
508
509#[cfg(test)]
510mod tests {
511 use super::*;
512
513 fn roundtrip_chunks(data: &[u8], feed: &[usize]) {
514 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
515 let mut compressed = Vec::new();
516 let mut off = 0;
517 for &n in feed {
518 let end = (off + n).min(data.len());
519 if end > off {
520 c.compress(&data[off..end], &mut compressed).unwrap();
521 }
522 off = end;
523 }
524 c.finish(&mut compressed).unwrap();
525 assert!(!compressed.is_empty());
526
527 let mut d = ZripDecompressor::new();
529 let mut plain = Vec::new();
530 let mut step = 1;
531 let mut i = 0;
532 while i < compressed.len() {
533 let end = (i + step).min(compressed.len());
534 d.decompress(&compressed[i..end], &mut plain).unwrap();
535 i = end;
536 step = step % 5 + 1; }
538 d.finish(&mut plain).unwrap();
539 assert_eq!(plain, data, "roundtrip mismatch with feed {feed:?}");
540 }
541
542 #[test]
543 fn roundtrip_single_chunk() {
544 let data: Vec<u8> = (0..100_000u32).map(|i| (i % 251) as u8).collect();
545 roundtrip_chunks(&data, &[data.len()]);
546 }
547
548 #[test]
549 fn roundtrip_multi_chunk_64k_windows() {
550 let data: Vec<u8> = (0..300_000u32).map(|i| (i / 7) as u8).collect();
552 let feed: Vec<usize> = std::iter::repeat(STREAM_BUF_SIZE).take(5).collect();
553 roundtrip_chunks(&data, &feed);
554 }
555
556 #[test]
557 fn roundtrip_highly_compressible() {
558 let data = b"libfw streaming compression test. ".repeat(10_000);
559 let mut feed = Vec::new();
561 let mut consumed = 0;
562 for (_, &size) in [1024usize, 2048, 4096, 8192, 16384].iter().cycle().enumerate() {
563 feed.push(size);
564 consumed += size;
565 if consumed >= data.len() {
566 break;
567 }
568 }
569 roundtrip_chunks(&data, &feed);
570 }
571
572 #[test]
573 fn roundtrip_empty_stream() {
574 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
575 let mut compressed = Vec::new();
576 c.finish(&mut compressed).unwrap();
577
578 let mut d = ZripDecompressor::new();
579 let mut plain = Vec::new();
580 d.decompress(&compressed, &mut plain).unwrap();
581 d.finish(&mut plain).unwrap();
582 assert!(plain.is_empty());
583 }
584
585 #[test]
586 fn empty_input_chunks_are_noops() {
587 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
588 let mut out = Vec::new();
589 c.compress(&[], &mut out).unwrap();
590 c.compress(b"hello", &mut out).unwrap();
591 c.finish(&mut out).unwrap();
592
593 let mut d = ZripDecompressor::new();
594 let mut plain = Vec::new();
595 d.decompress(&out, &mut plain).unwrap();
596 d.finish(&mut plain).unwrap();
597 assert_eq!(plain, b"hello");
598 }
599
600 #[test]
601 fn format_header_roundtrip() {
602 assert_eq!(CompressionFormat::parse_header("zrip"), Some(CompressionFormat::Zrip));
603 assert_eq!(CompressionFormat::parse_header("ZSTD"), Some(CompressionFormat::Zrip));
604 assert_eq!(CompressionFormat::parse_header("identity"), Some(CompressionFormat::None));
605 assert_eq!(CompressionFormat::parse_header(""), Some(CompressionFormat::None));
606 assert_eq!(CompressionFormat::parse_header("br"), None);
607 assert_eq!(CompressionFormat::Zrip.as_str(), "zrip");
608 assert_eq!(CompressionFormat::None.as_str(), "identity");
609 }
610
611 #[test]
612 fn passthrough_roundtrip() {
613 let mut c = compressor(CompressionFormat::None).unwrap();
614 let mut d = decompressor(CompressionFormat::None);
615 let mut compressed = Vec::new();
616 let mut plain = Vec::new();
617 c.compress(b"abc", &mut compressed).unwrap();
618 c.finish(&mut compressed).unwrap();
619 d.decompress(&compressed, &mut plain).unwrap();
620 d.finish(&mut plain).unwrap();
621 assert_eq!(plain, b"abc");
622 }
623
624 #[test]
625 fn truncated_stream_is_detected() {
626 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
627 let mut compressed = Vec::new();
628 c.compress(&vec![7u8; 5000], &mut compressed).unwrap();
629 c.finish(&mut compressed).unwrap();
630 compressed.truncate(compressed.len() - 1); let mut d = ZripDecompressor::new();
633 let mut plain = Vec::new();
634 d.decompress(&compressed, &mut plain).unwrap();
635 assert!(matches!(d.finish(&mut plain), Err(DecompressError::Truncated(_))));
636 }
637
638 #[test]
639 fn corrupt_stream_is_detected() {
640 let mut d = ZripDecompressor::new();
641 let mut plain = Vec::new();
642 let err = d.decompress(b"this is not a zstd frame at all", &mut plain);
643 assert!(err.is_err());
644 }
645
646 #[test]
647 fn per_call_output_budget_rejects_multi_frame_bomb() {
648 let mut compressed = Vec::new();
652 for _ in 0..64 {
653 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
654 c.compress(&vec![7u8; STREAM_BUF_SIZE], &mut compressed)
655 .unwrap();
656 c.finish(&mut compressed).unwrap();
657 }
658 let mut d = ZripDecompressor::with_max_output(MAX_FRAME_OUTPUT);
659 let mut plain = Vec::new();
660 let err = d.decompress(&compressed, &mut plain);
661 assert!(
662 matches!(err, Err(DecompressError::TooLarge { .. })),
663 "expected TooLarge, got {err:?}"
664 );
665 assert!(plain.len() <= MAX_FRAME_OUTPUT);
667 }
668
669 #[test]
670 fn generous_default_budget_allows_coalesced_frames() {
671 let mut compressed = Vec::new();
674 for _ in 0..8 {
675 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
676 c.compress(&vec![7u8; STREAM_BUF_SIZE], &mut compressed)
677 .unwrap();
678 c.finish(&mut compressed).unwrap();
679 }
680 let mut d = ZripDecompressor::new();
681 let mut plain = Vec::new();
682 d.decompress(&compressed, &mut plain).unwrap();
683 d.finish(&mut plain).unwrap();
684 assert_eq!(plain.len(), 8 * STREAM_BUF_SIZE);
685 }
686
687 #[test]
688 fn zstd_compat_interop() {
689 let data: Vec<u8> = (0..50_000u32).map(|i| (i % 31) as u8).collect();
692 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
693 let mut compressed = Vec::new();
694 c.compress(&data, &mut compressed).unwrap();
695 c.finish(&mut compressed).unwrap();
696 let decoded = zstd::stream::decode_all(&compressed[..]).unwrap();
697 assert_eq!(decoded, data);
698
699 let zstd_enc = zstd::stream::encode_all(&data[..], 1).unwrap();
700 let mut d = ZripDecompressor::new();
701 let mut plain = Vec::new();
702 d.decompress(&zstd_enc, &mut plain).unwrap();
703 d.finish(&mut plain).unwrap();
704 assert_eq!(plain, data);
705 }
706}