Skip to main content

libfw_core/
compress.rs

1//! Streaming compression abstractions backed by [`zrip`] (zstd).
2//!
3//! # Design
4//!
5//! Both [`Compressor`] and [`Decompressor`] are *stream-oriented*: they
6//! accept bounded input chunks and produce bounded output chunks, so the
7//! heap footprint of processing a file does **not** grow with the file
8//! size. The [`STREAM_BUF_SIZE`](crate::STREAM_BUF_SIZE) sliding window
9//! keeps memory constant (~64 KiB typical, capped well below 4 MiB even
10//! for adversarial frames).
11//!
12//! # Wire format
13//!
14//! The compressed body is a concatenation of independent zstd *frames*:
15//! each [`Compressor::compress`] call emits one complete frame when the
16//! next call arrives (or on [`Compressor::finish`]).
17//!
18//! The decompressor buffers compressed bytes until a complete frame is
19//! present (frame boundaries are detected by a small built-in zstd frame
20//! parser), then decodes exactly that frame with
21//! [`zrip::decompress_with_limit`]. A frame may therefore span several
22//! `decompress()` inputs, and one input may contain several frames.
23//!
24//! # Constant memory
25//!
26//! - **Compressor**: buffers at most one input chunk internally plus the
27//!   zrip encoder workspace (~150 KiB).
28//! - **Decompressor**: holds at most one *compressed* frame
29//!   ([`MAX_PENDING_FRAME`]) and one *decompressed* frame
30//!   ([`MAX_FRAME_OUTPUT`]) at a time.
31//!
32//! The server should feed [`Compressor::compress`] in
33//! [`STREAM_BUF_SIZE`](crate::STREAM_BUF_SIZE) windows (64 KiB), keeping
34//! frames ~64 KiB and the client's transient buffers small.
35
36use std::io::Write;
37
38use crate::error::{CompressError, DecompressError};
39use crate::{CHUNK_SIZE, STREAM_BUF_SIZE};
40
41/// Compression formats understood by libfw.
42#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
43#[serde(rename_all = "lowercase")]
44pub enum CompressionFormat {
45    /// No compression; the body passes through verbatim.
46    #[serde(rename = "identity")]
47    None,
48    /// zstd via the [`zrip`] codec (levels -8..=4).
49    Zrip,
50}
51
52impl CompressionFormat {
53    /// Wire representation used in
54    /// [`HEADER_COMPRESS`](crate::HEADER_COMPRESS) / `Content-Encoding`.
55    pub fn as_str(self) -> &'static str {
56        match self {
57            CompressionFormat::None => "identity",
58            CompressionFormat::Zrip => "zrip",
59        }
60    }
61
62    /// Parses a header value (`identity`, `zrip`, …).
63    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
72/// Safety ceiling for the decompressed size of a single frame.
73///
74/// The protocol keeps frames ≤ [`CHUNK_SIZE`](crate::CHUNK_SIZE); this
75/// limit is a guard against buggy or hostile peers.
76pub const MAX_FRAME_OUTPUT: usize = CHUNK_SIZE as usize;
77
78/// Maximum compressed bytes buffered while waiting for a frame boundary.
79///
80/// The worst-case incompressible [`CHUNK_SIZE`](crate::CHUNK_SIZE) chunk
81/// compresses to ≈ chunk size + framing overhead, hence the slack.
82pub const MAX_PENDING_FRAME: usize = CHUNK_SIZE as usize + STREAM_BUF_SIZE;
83
84/// Default cap on the total bytes a single [`Decompressor::decompress`] call
85/// may append to the output buffer.
86///
87/// A hostile peer could otherwise send many *small* frames in one network
88/// chunk, each expanding to [`MAX_FRAME_OUTPUT`] — inflating memory by a
89/// large factor before the caller drains the buffer. The default (a handful
90/// of frames) keeps the transient peak bounded while comfortably allowing
91/// legitimate coalesced reads (e.g. a browser delivering several 64 KiB
92/// download frames at once).
93pub const MAX_OUTPUT_PER_CALL: usize = MAX_FRAME_OUTPUT.saturating_mul(8);
94
95/// Default zrip compression level (Fast strategy, good for network transfer).
96pub const ZRIP_DEFAULT_LEVEL: i32 = 1;
97
98/// Lowest zrip compression level (Fast strategy; least CPU, worst ratio).
99pub const ZRIP_MIN_LEVEL: i32 = -8;
100
101/// Highest zrip compression level (bounded memory, best ratio).
102pub const ZRIP_MAX_LEVEL: i32 = 4;
103
104/// Whether `level` is an addressable zrip compression level.
105pub fn is_valid_zrip_level(level: i32) -> bool {
106    (ZRIP_MIN_LEVEL..=ZRIP_MAX_LEVEL).contains(&level)
107}
108
109/// Clamp a requested zrip level into `min..=max`.
110///
111/// Servers must never trust a client-supplied level: out-of-range requests
112/// are clamped (and the actual level echoed back on the response header),
113/// keeping the wire format honest without hard-failing older clients.
114pub fn negotiate_level(req: Option<i32>, min: i32, max: i32, default: i32) -> i32 {
115    if min > max {
116        // Empty range (misconfigured advertisement): fall back to the
117        // default rather than panicking inside `i32::clamp`.
118        return default;
119    }
120    match req {
121        None => default,
122        Some(l) => l.clamp(min, max),
123    }
124}
125
126/// Streaming compressor. Feed it bounded input chunks; it appends the
127/// compressed bytes produced so far to the provided output buffer.
128pub trait Compressor: Send {
129    /// The format this compressor produces.
130    fn format(&self) -> CompressionFormat;
131
132    /// Compress `input`, appending compressed bytes to `out`.
133    ///
134    /// `input` should be kept ≤ [`CHUNK_SIZE`](crate::CHUNK_SIZE) to keep
135    /// the output frames small and memory constant. When the next chunk
136    /// arrives, the previous chunk's frame is finalized and emitted.
137    fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError>;
138
139    /// Finalize the stream, appending the last frame (if any) to `out`.
140    ///
141    /// Must be called exactly once; afterwards the compressor is spent.
142    fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError>;
143}
144
145/// Streaming decompressor. Feed it arbitrary compressed bytes; decoded
146/// output is appended to the provided buffer.
147pub trait Decompressor: Send {
148    /// The format this decompressor consumes.
149    fn format(&self) -> CompressionFormat;
150
151    /// Decompress `input`, appending decoded bytes to `out`.
152    ///
153    /// Input may split frames arbitrarily. `out` should be drained by the
154    /// caller after each call to keep memory bounded.
155    fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError>;
156
157    /// Signal the end of the stream.
158    ///
159    /// Flushes any complete trailing frames and verifies the stream is
160    /// well-formed (no truncated frames). Call exactly once.
161    fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError>;
162}
163
164/// Construct a compressor for `format`.
165pub fn compressor(format: CompressionFormat) -> Result<Box<dyn Compressor>, CompressError> {
166    compressor_with_level(format, ZRIP_DEFAULT_LEVEL)
167}
168
169/// Construct a compressor for `format` at an explicit zrip level.
170///
171/// The level is only meaningful for [`CompressionFormat::Zrip`]; for
172/// [`CompressionFormat::None`] the body passes through verbatim regardless
173/// of `level`. Callers should pre-validate/clamp `level` with
174/// [`negotiate_level`] before constructing (constructors still error on
175/// out-of-range levels from zrip itself).
176pub 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
186/// Construct a decompressor for `format`.
187pub fn decompressor(format: CompressionFormat) -> Box<dyn Decompressor> {
188    decompressor_with_limit(format, MAX_OUTPUT_PER_CALL)
189}
190
191/// Construct a decompressor for `format` with an explicit per-call output
192/// budget (bytes a single [`Decompressor::decompress`] call may append).
193///
194/// Use a tight budget (e.g. [`MAX_FRAME_OUTPUT`]) on the server where the
195/// peer is potentially hostile; use the default generous budget on the
196/// client so coalesced multi-frame reads are never rejected.
197pub 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
207// ---------------------------------------------------------------------------
208// Passthrough (identity)
209// ---------------------------------------------------------------------------
210
211struct 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
245// ---------------------------------------------------------------------------
246// zrip (zstd) implementation
247// ---------------------------------------------------------------------------
248
249/// zstd compressor emitting one independent frame per input chunk.
250pub struct ZripCompressor {
251    encoder: Option<zrip::FrameEncoder<Vec<u8>>>,
252    /// Whether data has been written since the last frame boundary.
253    dirty: bool,
254}
255
256impl ZripCompressor {
257    /// Create a compressor at `level` (-8..=4; 0 = library default).
258    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        // Finalize the previous chunk's frame and hand its compressed bytes back.
280        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
301/// zstd decompressor reassembling frames across arbitrary chunk boundaries.
302pub struct ZripDecompressor {
303    /// Compressed bytes not yet turned into complete frames.
304    pending: Vec<u8>,
305    finished: bool,
306    /// Hard cap on the total bytes appended to `out` during a single
307    /// [`Decompressor::decompress`] / [`Decompressor::finish`] call.
308    max_output_per_call: usize,
309}
310
311impl ZripDecompressor {
312    /// Create a decompressor with the default per-call output budget
313    /// ([`MAX_OUTPUT_PER_CALL`]).
314    pub fn new() -> Self {
315        ZripDecompressor::with_max_output(MAX_OUTPUT_PER_CALL)
316    }
317
318    /// Create a decompressor with an explicit per-call output budget.
319    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    /// Decode as many complete frames as `pending` holds, appending at most
328    /// `max_add` bytes to `out` across the whole call.
329    fn drain_frames(&mut self, out: &mut Vec<u8>, max_add: usize) -> Result<(), DecompressError> {
330        let mut added = 0usize;
331        // Consume via a cursor and compact `pending` once at the end: a
332        // per-frame `drain(..len)` would memmove the whole remaining buffer
333        // for every frame (O(n²) across a many-frame stream).
334        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                    // Enforce the per-call budget *before* appending so the
354                    // memory spike of a multi-frame bomb never materializes.
355                    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                    // Keep the one compaction cheap: only memmove when the
364                    // consumed prefix is non-trivially large.
365                    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
426// ---------------------------------------------------------------------------
427// Minimal zstd frame-boundary parser
428// ---------------------------------------------------------------------------
429
430const ZSTD_MAGIC: u32 = 0xFD2F_B528;
431const SKIPPABLE_MASK: u32 = 0xFFFF_FFF0;
432const SKIPPABLE_MAGIC: u32 = 0x184D_2A50;
433
434enum FrameBoundary {
435    /// A complete frame occupying `len` bytes, declaring `content_size`
436    /// output bytes (`None` when the frame header does not state it).
437    Complete { len: usize, content_size: Option<u64> },
438    /// Need more bytes to know the boundary.
439    Incomplete,
440    /// Corrupt / unsupported.
441    Invalid,
442}
443
444/// Determine the byte length of the next zstd frame in `buf` (if present).
445fn 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
459/// Skippable frames: `magic(4) | skip_size u32(4) | payload(skip_size)`.
460fn 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
476/// Walk the frame header + blocks of a standard zstd frame.
477///
478/// Layout: `magic(4) descriptor(1) [window(1)] [dict_id(0/1/2/4)]
479/// [content_size(0/1/2/4/8)] block* checksum(0/4)`.
480fn zstd_frame_boundary(buf: &[u8]) -> FrameBoundary {
481    if buf.len() < 5 {
482        return FrameBoundary::Incomplete;
483    }
484    let descriptor = buf[4];
485    // Reserved bits (3..=4) must be zero.
486    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; // window descriptor
497    }
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    // Walk blocks until the last-block flag.
530    let mut off = hdr_len;
531    loop {
532        if buf.len() < off + 3 {
533            return FrameBoundary::Incomplete;
534        }
535        // Standard zstd block header (as read by the reference decoder):
536        // bit 0 = last_block, bits 1..=2 = block_type, bits 3..=23 = size.
537        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            // Reserved block type.
543            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        // Decode with arbitrarily split inputs, including mid-frame splits.
584        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; // 1..=5
593        }
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        // Server-style: feed in 64 KiB sliding windows.
607        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        // Feed sizes cycling 1KiB..16KiB until the whole input is consumed.
616        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        // Keep the parser defensive: malformed / hostile frame headers must
644        // surface as Invalid instead of crashing the WASM guest with an
645        // internal `unreachable!`.
646        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); // chop one byte
699
700        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        // A hostile peer delivers many small frames in ONE decompress() call;
717        // the cumulative decoded output must trip the per-call budget instead
718        // of ballooning memory.
719        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        // The buffer must not have been inflated past the budget.
734        assert!(plain.len() <= MAX_FRAME_OUTPUT);
735    }
736
737    #[test]
738    fn generous_default_budget_allows_coalesced_frames() {
739        // Several small frames in one call are fine under the default budget
740        // (what a browser may hand the client's decompressor).
741        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        // Our frames are standard zstd: the reference `zstd` codec can
758        // decode them, and we can decode its output.
759        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}