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, serde::Serialize, serde::Deserialize)]
43#[serde(rename_all = "lowercase")]
44pub enum CompressionFormat {
45 #[serde(rename = "identity")]
47 None,
48 Zrip,
50}
51
52impl CompressionFormat {
53 pub fn as_str(self) -> &'static str {
56 match self {
57 CompressionFormat::None => "identity",
58 CompressionFormat::Zrip => "zrip",
59 }
60 }
61
62 pub fn parse_header(s: &str) -> Option<Self> {
64 match s.trim().to_ascii_lowercase().as_str() {
65 "" | "identity" | "none" => Some(CompressionFormat::None),
66 "zrip" | "zstd" => Some(CompressionFormat::Zrip),
67 _ => None,
68 }
69 }
70}
71
72pub const MAX_FRAME_OUTPUT: usize = CHUNK_SIZE as usize;
77
78pub const MAX_PENDING_FRAME: usize = CHUNK_SIZE as usize + STREAM_BUF_SIZE;
83
84pub const MAX_OUTPUT_PER_CALL: usize = MAX_FRAME_OUTPUT.saturating_mul(8);
94
95pub const ZRIP_DEFAULT_LEVEL: i32 = 1;
97
98pub const ZRIP_MIN_LEVEL: i32 = -8;
100
101pub const ZRIP_MAX_LEVEL: i32 = 4;
103
104pub fn is_valid_zrip_level(level: i32) -> bool {
106 (ZRIP_MIN_LEVEL..=ZRIP_MAX_LEVEL).contains(&level)
107}
108
109pub fn negotiate_level(req: Option<i32>, min: i32, max: i32, default: i32) -> i32 {
115 if min > max {
116 return default;
119 }
120 match req {
121 None => default,
122 Some(l) => l.clamp(min, max),
123 }
124}
125
126pub trait Compressor: Send {
129 fn format(&self) -> CompressionFormat;
131
132 fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError>;
138
139 fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError>;
143}
144
145pub trait Decompressor: Send {
148 fn format(&self) -> CompressionFormat;
150
151 fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError>;
156
157 fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError>;
162}
163
164pub fn compressor(format: CompressionFormat) -> Result<Box<dyn Compressor>, CompressError> {
166 compressor_with_level(format, ZRIP_DEFAULT_LEVEL)
167}
168
169pub fn compressor_with_level(
177 format: CompressionFormat,
178 level: i32,
179) -> Result<Box<dyn Compressor>, CompressError> {
180 match format {
181 CompressionFormat::None => Ok(Box::new(PassthroughCompressor)),
182 CompressionFormat::Zrip => Ok(Box::new(ZripCompressor::new(level)?)),
183 }
184}
185
186pub fn decompressor(format: CompressionFormat) -> Box<dyn Decompressor> {
188 decompressor_with_limit(format, MAX_OUTPUT_PER_CALL)
189}
190
191pub fn decompressor_with_limit(
198 format: CompressionFormat,
199 max_output_per_call: usize,
200) -> Box<dyn Decompressor> {
201 match format {
202 CompressionFormat::None => Box::new(PassthroughDecompressor),
203 CompressionFormat::Zrip => Box::new(ZripDecompressor::with_max_output(max_output_per_call)),
204 }
205}
206
207struct PassthroughCompressor;
212
213impl Compressor for PassthroughCompressor {
214 fn format(&self) -> CompressionFormat {
215 CompressionFormat::None
216 }
217
218 fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError> {
219 out.extend_from_slice(input);
220 Ok(())
221 }
222
223 fn finish(&mut self, _out: &mut Vec<u8>) -> Result<(), CompressError> {
224 Ok(())
225 }
226}
227
228struct PassthroughDecompressor;
229
230impl Decompressor for PassthroughDecompressor {
231 fn format(&self) -> CompressionFormat {
232 CompressionFormat::None
233 }
234
235 fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError> {
236 out.extend_from_slice(input);
237 Ok(())
238 }
239
240 fn finish(&mut self, _out: &mut Vec<u8>) -> Result<(), DecompressError> {
241 Ok(())
242 }
243}
244
245pub struct ZripCompressor {
251 encoder: Option<zrip::FrameEncoder<Vec<u8>>>,
252 dirty: bool,
254}
255
256impl ZripCompressor {
257 pub fn new(level: i32) -> Result<Self, CompressError> {
259 Ok(ZripCompressor {
260 encoder: Some(zrip::FrameEncoder::new(Vec::new(), level)?),
261 dirty: false,
262 })
263 }
264}
265
266impl Compressor for ZripCompressor {
267 fn format(&self) -> CompressionFormat {
268 CompressionFormat::Zrip
269 }
270
271 fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError> {
272 if input.is_empty() {
273 return Ok(());
274 }
275 let encoder = self
276 .encoder
277 .as_mut()
278 .ok_or_else(|| std::io::Error::other("compressor already finished"))?;
279 if self.dirty {
281 let finished = encoder.reset(Vec::new())?;
282 out.extend_from_slice(&finished);
283 }
284 encoder.write_all(input)?;
285 self.dirty = true;
286 Ok(())
287 }
288
289 fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError> {
290 let encoder = self
291 .encoder
292 .take()
293 .ok_or_else(|| std::io::Error::other("compressor already finished"))?;
294 let tail = encoder.finish()?;
295 out.extend_from_slice(&tail);
296 self.dirty = false;
297 Ok(())
298 }
299}
300
301pub struct ZripDecompressor {
303 pending: Vec<u8>,
305 finished: bool,
306 max_output_per_call: usize,
309}
310
311impl ZripDecompressor {
312 pub fn new() -> Self {
315 ZripDecompressor::with_max_output(MAX_OUTPUT_PER_CALL)
316 }
317
318 pub fn with_max_output(max_output_per_call: usize) -> Self {
320 ZripDecompressor {
321 pending: Vec::with_capacity(STREAM_BUF_SIZE),
322 finished: false,
323 max_output_per_call,
324 }
325 }
326
327 fn drain_frames(&mut self, out: &mut Vec<u8>, max_add: usize) -> Result<(), DecompressError> {
330 let mut added = 0usize;
331 let mut consumed = 0usize;
335 loop {
336 let boundary = frame_boundary(&self.pending[consumed..]);
337 match boundary {
338 FrameBoundary::Complete { len, content_size } => {
339 if content_size.is_some_and(|cs| cs > MAX_FRAME_OUTPUT as u64) {
340 return Err(DecompressError::TooLarge {
341 limit: MAX_FRAME_OUTPUT,
342 });
343 }
344 let decoded = {
345 let frame = &self.pending[consumed..consumed + len];
346 zrip::decompress_with_limit(frame, MAX_FRAME_OUTPUT).map_err(|e| {
347 DecompressError::Io(std::io::Error::new(
348 std::io::ErrorKind::InvalidData,
349 e,
350 ))
351 })?
352 };
353 if added.saturating_add(decoded.len()) > max_add {
356 return Err(DecompressError::TooLarge { limit: max_add });
357 }
358 out.extend_from_slice(&decoded);
359 added = added.saturating_add(decoded.len());
360 consumed += len;
361 }
362 FrameBoundary::Incomplete => {
363 if consumed > 0 {
366 self.pending.drain(..consumed);
367 }
368 return Ok(());
369 }
370 FrameBoundary::Invalid => {
371 return Err(DecompressError::Io(std::io::Error::new(
372 std::io::ErrorKind::InvalidData,
373 "invalid zstd frame data",
374 )))
375 }
376 }
377 }
378 }
379}
380
381impl Default for ZripDecompressor {
382 fn default() -> Self {
383 ZripDecompressor::new()
384 }
385}
386
387impl Decompressor for ZripDecompressor {
388 fn format(&self) -> CompressionFormat {
389 CompressionFormat::Zrip
390 }
391
392 fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError> {
393 if self.finished {
394 return Err(DecompressError::Io(std::io::Error::other(
395 "decompressor already finished",
396 )));
397 }
398 if !input.is_empty() {
399 self.pending.extend_from_slice(input);
400 }
401 self.drain_frames(out, self.max_output_per_call)?;
402 if self.pending.len() > MAX_PENDING_FRAME {
403 return Err(DecompressError::TooLarge {
404 limit: MAX_PENDING_FRAME,
405 });
406 }
407 Ok(())
408 }
409
410 fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError> {
411 if self.finished {
412 return Ok(());
413 }
414 self.finished = true;
415 self.drain_frames(out, self.max_output_per_call)?;
416 if !self.pending.is_empty() {
417 return Err(DecompressError::Truncated(std::io::Error::new(
418 std::io::ErrorKind::UnexpectedEof,
419 format!("{} trailing compressed bytes", self.pending.len()),
420 )));
421 }
422 Ok(())
423 }
424}
425
426const ZSTD_MAGIC: u32 = 0xFD2F_B528;
431const SKIPPABLE_MASK: u32 = 0xFFFF_FFF0;
432const SKIPPABLE_MAGIC: u32 = 0x184D_2A50;
433
434enum FrameBoundary {
435 Complete { len: usize, content_size: Option<u64> },
438 Incomplete,
440 Invalid,
442}
443
444fn frame_boundary(buf: &[u8]) -> FrameBoundary {
446 if buf.len() < 4 {
447 return FrameBoundary::Incomplete;
448 }
449 let magic = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]);
450 if magic == ZSTD_MAGIC {
451 zstd_frame_boundary(buf)
452 } else if (magic & SKIPPABLE_MASK) == SKIPPABLE_MAGIC {
453 skippable_frame_boundary(buf)
454 } else {
455 FrameBoundary::Invalid
456 }
457}
458
459fn skippable_frame_boundary(buf: &[u8]) -> FrameBoundary {
461 if buf.len() < 8 {
462 return FrameBoundary::Incomplete;
463 }
464 let skip = u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]) as usize;
465 let total = 8usize.saturating_add(skip);
466 if buf.len() >= total {
467 FrameBoundary::Complete {
468 len: total,
469 content_size: Some(0),
470 }
471 } else {
472 FrameBoundary::Incomplete
473 }
474}
475
476fn zstd_frame_boundary(buf: &[u8]) -> FrameBoundary {
481 if buf.len() < 5 {
482 return FrameBoundary::Incomplete;
483 }
484 let descriptor = buf[4];
485 if descriptor & 0x18 != 0 {
487 return FrameBoundary::Invalid;
488 }
489 let single_segment = descriptor & 0x20 != 0;
490 let checksum = descriptor & 0x04 != 0;
491 let dict_id_flag = descriptor & 0x03;
492 let fcs_flag = (descriptor >> 6) & 0x03;
493
494 let mut hdr_len = 5usize;
495 if !single_segment {
496 hdr_len += 1; }
498 let dict_id_size = match dict_id_flag {
499 0 => 0,
500 1 => 1,
501 2 => 2,
502 3 => 4,
503 _ => return FrameBoundary::Invalid,
504 };
505 hdr_len += dict_id_size;
506 let fcs_size: usize = match fcs_flag {
507 0 if single_segment => 1,
508 0 => 0,
509 1 => 2,
510 2 => 4,
511 3 => 8,
512 _ => return FrameBoundary::Invalid,
513 };
514 hdr_len += fcs_size;
515
516 if buf.len() < hdr_len {
517 return FrameBoundary::Incomplete;
518 }
519 let content_size = if fcs_size > 0 {
520 let mut v = 0u64;
521 for (i, &b) in buf[5..5 + fcs_size].iter().enumerate() {
522 v |= (b as u64) << (8 * i);
523 }
524 Some(v)
525 } else {
526 None
527 };
528
529 let mut off = hdr_len;
531 loop {
532 if buf.len() < off + 3 {
533 return FrameBoundary::Incomplete;
534 }
535 let block_header = u32::from_le_bytes([buf[off], buf[off + 1], buf[off + 2], 0]);
538 let last = block_header & 0x01 != 0;
539 let block_type = (block_header >> 1) & 0x03;
540 let block_size = (block_header >> 3) as usize;
541 if block_type == 3 {
542 return FrameBoundary::Invalid;
544 }
545 off += 3 + block_size;
546 if off > MAX_PENDING_FRAME {
547 return FrameBoundary::Invalid;
548 }
549 if last {
550 break;
551 }
552 }
553 if checksum {
554 off += 4;
555 }
556 if buf.len() < off {
557 return FrameBoundary::Incomplete;
558 }
559 FrameBoundary::Complete {
560 len: off,
561 content_size,
562 }
563}
564
565#[cfg(test)]
566mod tests {
567 use super::*;
568
569 fn roundtrip_chunks(data: &[u8], feed: &[usize]) {
570 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
571 let mut compressed = Vec::new();
572 let mut off = 0;
573 for &n in feed {
574 let end = (off + n).min(data.len());
575 if end > off {
576 c.compress(&data[off..end], &mut compressed).unwrap();
577 }
578 off = end;
579 }
580 c.finish(&mut compressed).unwrap();
581 assert!(!compressed.is_empty());
582
583 let mut d = ZripDecompressor::new();
585 let mut plain = Vec::new();
586 let mut step = 1;
587 let mut i = 0;
588 while i < compressed.len() {
589 let end = (i + step).min(compressed.len());
590 d.decompress(&compressed[i..end], &mut plain).unwrap();
591 i = end;
592 step = step % 5 + 1; }
594 d.finish(&mut plain).unwrap();
595 assert_eq!(plain, data, "roundtrip mismatch with feed {feed:?}");
596 }
597
598 #[test]
599 fn roundtrip_single_chunk() {
600 let data: Vec<u8> = (0..100_000u32).map(|i| (i % 251) as u8).collect();
601 roundtrip_chunks(&data, &[data.len()]);
602 }
603
604 #[test]
605 fn roundtrip_multi_chunk_64k_windows() {
606 let data: Vec<u8> = (0..300_000u32).map(|i| (i / 7) as u8).collect();
608 let feed: Vec<usize> = std::iter::repeat(STREAM_BUF_SIZE).take(5).collect();
609 roundtrip_chunks(&data, &feed);
610 }
611
612 #[test]
613 fn roundtrip_highly_compressible() {
614 let data = b"libfw streaming compression test. ".repeat(10_000);
615 let mut feed = Vec::new();
617 let mut consumed = 0;
618 for (_, &size) in [1024usize, 2048, 4096, 8192, 16384].iter().cycle().enumerate() {
619 feed.push(size);
620 consumed += size;
621 if consumed >= data.len() {
622 break;
623 }
624 }
625 roundtrip_chunks(&data, &feed);
626 }
627
628 #[test]
629 fn roundtrip_empty_stream() {
630 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
631 let mut compressed = Vec::new();
632 c.finish(&mut compressed).unwrap();
633
634 let mut d = ZripDecompressor::new();
635 let mut plain = Vec::new();
636 d.decompress(&compressed, &mut plain).unwrap();
637 d.finish(&mut plain).unwrap();
638 assert!(plain.is_empty());
639 }
640
641 #[test]
642 fn malformed_zstd_headers_are_rejected_without_panicking() {
643 let bad = [0xFD, 0x2F, 0xB5, 0x28, 0xFF];
647 assert!(matches!(frame_boundary(&bad), FrameBoundary::Invalid));
648
649 let bad2 = [0xFD, 0x2F, 0xB5, 0x28, 0x00, 0x00, 0x00, 0x00, 0x00];
650 assert!(matches!(frame_boundary(&bad2), FrameBoundary::Invalid));
651 }
652
653 #[test]
654 fn empty_input_chunks_are_noops() {
655 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
656 let mut out = Vec::new();
657 c.compress(&[], &mut out).unwrap();
658 c.compress(b"hello", &mut out).unwrap();
659 c.finish(&mut out).unwrap();
660
661 let mut d = ZripDecompressor::new();
662 let mut plain = Vec::new();
663 d.decompress(&out, &mut plain).unwrap();
664 d.finish(&mut plain).unwrap();
665 assert_eq!(plain, b"hello");
666 }
667
668 #[test]
669 fn format_header_roundtrip() {
670 assert_eq!(CompressionFormat::parse_header("zrip"), Some(CompressionFormat::Zrip));
671 assert_eq!(CompressionFormat::parse_header("ZSTD"), Some(CompressionFormat::Zrip));
672 assert_eq!(CompressionFormat::parse_header("identity"), Some(CompressionFormat::None));
673 assert_eq!(CompressionFormat::parse_header(""), Some(CompressionFormat::None));
674 assert_eq!(CompressionFormat::parse_header("br"), None);
675 assert_eq!(CompressionFormat::Zrip.as_str(), "zrip");
676 assert_eq!(CompressionFormat::None.as_str(), "identity");
677 }
678
679 #[test]
680 fn passthrough_roundtrip() {
681 let mut c = compressor(CompressionFormat::None).unwrap();
682 let mut d = decompressor(CompressionFormat::None);
683 let mut compressed = Vec::new();
684 let mut plain = Vec::new();
685 c.compress(b"abc", &mut compressed).unwrap();
686 c.finish(&mut compressed).unwrap();
687 d.decompress(&compressed, &mut plain).unwrap();
688 d.finish(&mut plain).unwrap();
689 assert_eq!(plain, b"abc");
690 }
691
692 #[test]
693 fn truncated_stream_is_detected() {
694 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
695 let mut compressed = Vec::new();
696 c.compress(&vec![7u8; 5000], &mut compressed).unwrap();
697 c.finish(&mut compressed).unwrap();
698 compressed.truncate(compressed.len() - 1); let mut d = ZripDecompressor::new();
701 let mut plain = Vec::new();
702 d.decompress(&compressed, &mut plain).unwrap();
703 assert!(matches!(d.finish(&mut plain), Err(DecompressError::Truncated(_))));
704 }
705
706 #[test]
707 fn corrupt_stream_is_detected() {
708 let mut d = ZripDecompressor::new();
709 let mut plain = Vec::new();
710 let err = d.decompress(b"this is not a zstd frame at all", &mut plain);
711 assert!(err.is_err());
712 }
713
714 #[test]
715 fn per_call_output_budget_rejects_multi_frame_bomb() {
716 let mut compressed = Vec::new();
720 for _ in 0..64 {
721 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
722 c.compress(&vec![7u8; STREAM_BUF_SIZE], &mut compressed)
723 .unwrap();
724 c.finish(&mut compressed).unwrap();
725 }
726 let mut d = ZripDecompressor::with_max_output(MAX_FRAME_OUTPUT);
727 let mut plain = Vec::new();
728 let err = d.decompress(&compressed, &mut plain);
729 assert!(
730 matches!(err, Err(DecompressError::TooLarge { .. })),
731 "expected TooLarge, got {err:?}"
732 );
733 assert!(plain.len() <= MAX_FRAME_OUTPUT);
735 }
736
737 #[test]
738 fn generous_default_budget_allows_coalesced_frames() {
739 let mut compressed = Vec::new();
742 for _ in 0..8 {
743 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
744 c.compress(&vec![7u8; STREAM_BUF_SIZE], &mut compressed)
745 .unwrap();
746 c.finish(&mut compressed).unwrap();
747 }
748 let mut d = ZripDecompressor::new();
749 let mut plain = Vec::new();
750 d.decompress(&compressed, &mut plain).unwrap();
751 d.finish(&mut plain).unwrap();
752 assert_eq!(plain.len(), 8 * STREAM_BUF_SIZE);
753 }
754
755 #[test]
756 fn zstd_compat_interop() {
757 let data: Vec<u8> = (0..50_000u32).map(|i| (i % 31) as u8).collect();
760 let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
761 let mut compressed = Vec::new();
762 c.compress(&data, &mut compressed).unwrap();
763 c.finish(&mut compressed).unwrap();
764 let decoded = zstd::stream::decode_all(&compressed[..]).unwrap();
765 assert_eq!(decoded, data);
766
767 let zstd_enc = zstd::stream::encode_all(&data[..], 1).unwrap();
768 let mut d = ZripDecompressor::new();
769 let mut plain = Vec::new();
770 d.decompress(&zstd_enc, &mut plain).unwrap();
771 d.finish(&mut plain).unwrap();
772 assert_eq!(plain, data);
773 }
774}