zrip_decode/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![cfg_attr(feature = "nightly", feature(optimize_attribute))]
3#![cfg_attr(feature = "paranoid", forbid(unsafe_code))]
4
5#[cfg(feature = "alloc")]
6extern crate alloc;
7
8pub(crate) mod block_decoder;
9#[cfg(feature = "std")]
10pub mod context;
11pub(crate) mod exec;
12pub(crate) mod fast_vec;
13pub(crate) mod literals;
14pub(crate) mod ring_buffer;
15pub(crate) mod seq_table;
16pub(crate) mod sequences;
17#[cfg(feature = "std")]
18pub mod streaming;
19
20#[cfg(feature = "alloc")]
21use alloc::boxed::Box;
22#[cfg(feature = "alloc")]
23use alloc::vec::Vec;
24
25use crate::exec::{decode_execute_sequences, decode_execute_single_sequence};
26use crate::literals::decode_literals_ws;
27use crate::sequences::{SequenceDecodeTables, parse_sequence_count, parse_sequence_tables_ws};
28use zrip_core::block::{BlockType, parse_block_header};
29use zrip_core::error::DecompressError;
30use zrip_core::frame::MAX_WINDOW_SIZE;
31#[cfg(feature = "std")]
32use zrip_core::frame::header::parse_frame_header_after_magic;
33use zrip_core::frame::header::{FrameHeader, parse_frame_header};
34use zrip_core::huffman::HuffmanDecodeEntry;
35use zrip_core::xxhash::Xxh64State;
36
37#[allow(clippy::struct_excessive_bools)]
38pub(crate) struct BlockDecodeWorkspace {
39    pub literal_buf: Vec<u8>,
40    pub huf_table: Vec<HuffmanDecodeEntry>,
41    pub huf_table_log: u8,
42    pub huf_valid: bool,
43    pub huf_all_weights: Vec<u8>,
44    pub huf_rank_count: Vec<u32>,
45    pub huf_rank_start: Vec<u32>,
46    pub huf_weights: Vec<u8>,
47    pub huf_last_weights: Vec<u8>,
48    pub huf_last_weights_valid: bool,
49    pub huf_last_header: Vec<u8>,
50    pub huf_last_header_valid: bool,
51    pub fse_dist: Vec<i16>,
52    pub fse_symbol_next: Vec<u16>,
53    pub fse_build_buf: Vec<zrip_core::fse::FseDecodeEntry>,
54    pub seq_tables: Option<Box<SequenceDecodeTables>>,
55    pub seq_table_header: Vec<u8>,
56    pub seq_table_cache: Option<Box<SequenceDecodeTables>>,
57    pub seq_table_cache_tables_current: bool,
58    pub cached_dict_tables: Option<Box<SequenceDecodeTables>>,
59    pub cached_dict_rep: [u32; 3],
60    pub cached_dict_huf: Option<(Vec<HuffmanDecodeEntry>, u8)>,
61}
62
63impl BlockDecodeWorkspace {
64    pub(crate) fn new() -> Self {
65        Self {
66            literal_buf: Vec::new(),
67            huf_table: Vec::new(),
68            huf_table_log: 0,
69            huf_valid: false,
70            huf_all_weights: Vec::new(),
71            huf_rank_count: Vec::new(),
72            huf_rank_start: Vec::new(),
73            huf_weights: Vec::new(),
74            huf_last_weights: Vec::new(),
75            huf_last_weights_valid: false,
76            huf_last_header: Vec::new(),
77            huf_last_header_valid: false,
78            fse_dist: Vec::new(),
79            fse_symbol_next: Vec::new(),
80            fse_build_buf: Vec::new(),
81            seq_tables: Some(Box::new(SequenceDecodeTables::new_default())),
82            seq_table_header: Vec::new(),
83            seq_table_cache: None,
84            seq_table_cache_tables_current: false,
85            cached_dict_tables: None,
86            cached_dict_rep: [1, 4, 8],
87            cached_dict_huf: None,
88        }
89    }
90
91    pub(crate) fn reset_huffman_state(&mut self) {
92        self.huf_valid = false;
93    }
94
95    #[cfg(feature = "std")]
96    pub(crate) fn cache_dict(&mut self, dict: &zrip_core::dict::Dictionary) {
97        let mut st = SequenceDecodeTables::new_default();
98        if let Some((t, l)) = dict.of_table() {
99            st.of_table = crate::seq_table::SeqTable::promote_of(t);
100            st.of_accuracy = l;
101            st.of_set = true;
102        }
103        if let Some((t, l)) = dict.ml_table() {
104            st.ml_table = crate::seq_table::SeqTable::promote_ml(t);
105            st.ml_accuracy = l;
106            st.ml_set = true;
107        }
108        if let Some((t, l)) = dict.ll_table() {
109            st.ll_table = crate::seq_table::SeqTable::promote_ll(t);
110            st.ll_accuracy = l;
111            st.ll_set = true;
112        }
113        self.cached_dict_tables = Some(Box::new(st));
114        self.cached_dict_rep = *dict.rep_offsets();
115        if let Some((t, l)) = dict.huf_table() {
116            self.cached_dict_huf = Some((t.to_vec(), l));
117        }
118    }
119}
120
121pub(crate) fn skip_skippable_frame(data: &[u8]) -> Option<usize> {
122    if data.len() < 8 {
123        return None;
124    }
125    let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
126    if (magic & 0xFFFF_FFF0) != 0x184D_2A50 {
127        return None;
128    }
129    let frame_size = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize;
130    let total = 8 + frame_size;
131    if total > data.len() {
132        return None;
133    }
134    Some(total)
135}
136
137pub fn decompress(input: &[u8]) -> Result<Vec<u8>, DecompressError> {
138    decompress_with_dict(input, None)
139}
140
141/// Decompress with an explicit output size limit.
142///
143/// Returns [`DecompressError::OutputTooSmall`] if the decompressed output would
144/// exceed `max_output_size` bytes. Use [`SAFE_DECOMPRESS_LIMIT`](zrip_core::SAFE_DECOMPRESS_LIMIT)
145/// when processing untrusted input to prevent memory exhaustion attacks.
146pub fn decompress_with_limit(
147    input: &[u8],
148    max_output_size: usize,
149) -> Result<Vec<u8>, DecompressError> {
150    let mut output = Vec::new();
151    let mut ws = Box::new(BlockDecodeWorkspace::new());
152    let mut offset = 0;
153    while offset < input.len() {
154        let remaining = &input[offset..];
155        if let Some(skip_len) = skip_skippable_frame(remaining) {
156            offset += skip_len;
157            continue;
158        }
159        let consumed = decompress_frame(remaining, &mut output, max_output_size, None, &mut ws)?;
160        offset += consumed;
161    }
162    Ok(output)
163}
164
165pub fn decompress_into(input: &[u8], output: &mut Vec<u8>) -> Result<usize, DecompressError> {
166    let max_output = zrip_core::DEFAULT_DECOMPRESS_LIMIT;
167    let mut ws = Box::new(BlockDecodeWorkspace::new());
168    let start = output.len();
169    let mut offset = 0;
170    while offset < input.len() {
171        let remaining = &input[offset..];
172        if let Some(skip_len) = skip_skippable_frame(remaining) {
173            offset += skip_len;
174            continue;
175        }
176        let consumed = decompress_frame(remaining, output, max_output, None, &mut ws)?;
177        offset += consumed;
178    }
179    Ok(output.len() - start)
180}
181
182pub fn decompress_with_dict(
183    input: &[u8],
184    dict: Option<&zrip_core::dict::Dictionary>,
185) -> Result<Vec<u8>, DecompressError> {
186    let max_output = zrip_core::DEFAULT_DECOMPRESS_LIMIT;
187    let mut output = Vec::new();
188    let mut ws = Box::new(BlockDecodeWorkspace::new());
189    let mut offset = 0;
190
191    while offset < input.len() {
192        let remaining = &input[offset..];
193        if let Some(skip_len) = skip_skippable_frame(remaining) {
194            offset += skip_len;
195            continue;
196        }
197        let consumed = decompress_frame(remaining, &mut output, max_output, dict, &mut ws)?;
198        offset += consumed;
199    }
200
201    Ok(output)
202}
203
204pub(crate) fn decompress_frame(
205    input: &[u8],
206    output: &mut Vec<u8>,
207    max_output: usize,
208    dict: Option<&zrip_core::dict::Dictionary>,
209    ws: &mut BlockDecodeWorkspace,
210) -> Result<usize, DecompressError> {
211    let header = parse_frame_header(input)?;
212    decompress_frame_with_header(input, output, max_output, dict, ws, header)
213}
214
215#[cfg(feature = "std")]
216pub(crate) fn decompress_frame_after_magic(
217    input: &[u8],
218    output: &mut Vec<u8>,
219    max_output: usize,
220    dict: Option<&zrip_core::dict::Dictionary>,
221    ws: &mut BlockDecodeWorkspace,
222) -> Result<usize, DecompressError> {
223    let header = parse_frame_header_after_magic(input, 0)?;
224    decompress_frame_with_header(input, output, max_output, dict, ws, header)
225}
226
227fn decompress_frame_with_header(
228    input: &[u8],
229    output: &mut Vec<u8>,
230    max_output: usize,
231    dict: Option<&zrip_core::dict::Dictionary>,
232    ws: &mut BlockDecodeWorkspace,
233    header: FrameHeader,
234) -> Result<usize, DecompressError> {
235    if header.window_size > MAX_WINDOW_SIZE && !header.single_segment {
236        return Err(DecompressError::WindowTooLarge {
237            requested: header.window_size,
238            max: MAX_WINDOW_SIZE,
239        });
240    }
241
242    if let Some(frame_dict_id) = header.dict_id {
243        match dict {
244            Some(d) if d.id() == frame_dict_id => {}
245            Some(d) => {
246                return Err(DecompressError::DictMismatch {
247                    expected: frame_dict_id,
248                    got: d.id(),
249                });
250            }
251            None => return Err(DecompressError::DictRequired),
252        }
253    }
254
255    if let Some(fcs) = header.frame_content_size {
256        if max_output < usize::MAX && fcs as usize > max_output {
257            return Err(DecompressError::OutputTooSmall);
258        }
259        let hint = (fcs as usize).min(MAX_WINDOW_SIZE as usize);
260        output.reserve(hint + 32);
261    }
262
263    let mut offset = header.header_size;
264    let output_start = output.len();
265
266    let dict_history: &[u8] = if let Some(d) = dict { d.content() } else { &[] };
267
268    let mut seq_tables = None;
269    let mut rep_offsets = [1u32, 4, 8];
270    ws.reset_huffman_state();
271    if let Some((ref t, l)) = ws.cached_dict_huf {
272        ws.huf_table.clear();
273        ws.huf_table.extend_from_slice(t);
274        ws.huf_table_log = l;
275        ws.huf_valid = true;
276        ws.huf_last_weights_valid = false;
277        ws.huf_last_header_valid = false;
278    } else if let Some(d) = dict
279        && let Some((t, l)) = d.huf_table()
280    {
281        ws.huf_table.clear();
282        ws.huf_table.extend_from_slice(t);
283        ws.huf_table_log = l;
284        ws.huf_valid = true;
285        ws.huf_last_weights_valid = false;
286        ws.huf_last_header_valid = false;
287    }
288
289    let mut hasher = if header.content_checksum {
290        Some(Xxh64State::new(0))
291    } else {
292        None
293    };
294
295    loop {
296        if offset + 3 > input.len() {
297            return Err(DecompressError::InputExhausted);
298        }
299        let block_header = parse_block_header(&input[offset..])?;
300        offset += 3;
301
302        let block_size = block_header.block_size as usize;
303
304        if block_size > zrip_core::frame::MAX_BLOCK_SIZE {
305            match block_header.block_type {
306                BlockType::Raw | BlockType::Rle => {
307                    return Err(DecompressError::BlockTooLarge);
308                }
309                BlockType::Compressed => {}
310            }
311        }
312
313        let block_output_start = output.len();
314        match block_header.block_type {
315            BlockType::Raw => {
316                if offset + block_size > input.len() {
317                    return Err(DecompressError::InputExhausted);
318                }
319                if output.len() - output_start + block_size > max_output {
320                    return Err(DecompressError::OutputTooSmall);
321                }
322                output.extend_from_slice(&input[offset..offset + block_size]);
323                offset += block_size;
324            }
325            BlockType::Rle => {
326                if offset >= input.len() {
327                    return Err(DecompressError::InputExhausted);
328                }
329                if output.len() - output_start + block_size > max_output {
330                    return Err(DecompressError::OutputTooSmall);
331                }
332                let byte = input[offset];
333                output.resize(output.len() + block_size, byte);
334                offset += 1;
335            }
336            BlockType::Compressed => {
337                if offset + block_size > input.len() {
338                    return Err(DecompressError::InputExhausted);
339                }
340                if seq_tables.is_none() {
341                    let mut initial_tables = ws
342                        .seq_tables
343                        .take()
344                        .unwrap_or_else(|| Box::new(SequenceDecodeTables::new_default()));
345                    let initial_rep_offsets =
346                        initial_sequence_state(initial_tables.as_mut(), ws, dict);
347                    seq_tables = Some(initial_tables);
348                    rep_offsets = initial_rep_offsets;
349                }
350                let block_data = &input[offset..offset + block_size];
351                decode_compressed_block(
352                    block_data,
353                    output,
354                    output_start,
355                    max_output,
356                    seq_tables
357                        .as_deref_mut()
358                        .expect("sequence tables are initialized before compressed blocks"),
359                    &mut rep_offsets,
360                    ws,
361                    dict_history,
362                )?;
363                offset += block_size;
364            }
365        }
366        if let Some(ref mut hasher) = hasher {
367            hasher.update(&output[block_output_start..]);
368        }
369
370        if block_header.last_block {
371            break;
372        }
373    }
374
375    if let Some(tables) = seq_tables.take() {
376        ws.seq_tables = Some(tables);
377    }
378
379    if let Some(ref mut hasher) = hasher {
380        let hash = hasher.finish();
381        let expected_checksum = (hash & 0xFFFF_FFFF) as u32;
382
383        if offset + 4 > input.len() {
384            return Err(DecompressError::InputExhausted);
385        }
386        let stored_checksum = u32::from_le_bytes([
387            input[offset],
388            input[offset + 1],
389            input[offset + 2],
390            input[offset + 3],
391        ]);
392        offset += 4;
393
394        if expected_checksum != stored_checksum {
395            return Err(DecompressError::ChecksumMismatch {
396                expected: stored_checksum,
397                got: expected_checksum,
398            });
399        }
400    }
401
402    if let Some(fcs) = header.frame_content_size
403        && (output.len() - output_start) as u64 != fcs
404    {
405        return Err(DecompressError::FrameSizeMismatch);
406    }
407
408    Ok(offset)
409}
410
411fn initial_sequence_state(
412    tables: &mut SequenceDecodeTables,
413    ws: &mut BlockDecodeWorkspace,
414    dict: Option<&zrip_core::dict::Dictionary>,
415) -> [u32; 3] {
416    if let Some(ref cached) = ws.cached_dict_tables {
417        *tables = (**cached).clone();
418        ws.seq_table_cache_tables_current = false;
419        ws.cached_dict_rep
420    } else if let Some(d) = dict {
421        tables.reset_default();
422        ws.seq_table_cache_tables_current = false;
423        if let Some((t, l)) = d.of_table() {
424            tables.of_table = crate::seq_table::SeqTable::promote_of(t);
425            tables.of_accuracy = l;
426            tables.of_kind = crate::sequences::SequenceTableKind::Other;
427            tables.of_set = true;
428        }
429        if let Some((t, l)) = d.ml_table() {
430            tables.ml_table = crate::seq_table::SeqTable::promote_ml(t);
431            tables.ml_accuracy = l;
432            tables.ml_kind = crate::sequences::SequenceTableKind::Other;
433            tables.ml_set = true;
434        }
435        if let Some((t, l)) = d.ll_table() {
436            tables.ll_table = crate::seq_table::SeqTable::promote_ll(t);
437            tables.ll_accuracy = l;
438            tables.ll_kind = crate::sequences::SequenceTableKind::Other;
439            tables.ll_set = true;
440        }
441        *d.rep_offsets()
442    } else {
443        tables.clear_repeat_flags();
444        [1u32, 4, 8]
445    }
446}
447
448#[allow(clippy::too_many_arguments)]
449fn decode_compressed_block(
450    data: &[u8],
451    output: &mut Vec<u8>,
452    output_start: usize,
453    max_output: usize,
454    seq_tables: &mut SequenceDecodeTables,
455    rep_offsets: &mut [u32; 3],
456    ws: &mut BlockDecodeWorkspace,
457    dict_history: &[u8],
458) -> Result<(), DecompressError> {
459    let lit_consumed = decode_literals_ws(data, ws)?;
460
461    let remaining = &data[lit_consumed..];
462
463    if remaining.is_empty() {
464        if output.len() - output_start + ws.literal_buf.len() > max_output {
465            return Err(DecompressError::OutputTooSmall);
466        }
467        output.extend_from_slice(&ws.literal_buf);
468        return Ok(());
469    }
470
471    let (num_sequences, seq_count_size) = parse_sequence_count(remaining)?;
472
473    if num_sequences == 0 {
474        if output.len() - output_start + ws.literal_buf.len() > max_output {
475            return Err(DecompressError::OutputTooSmall);
476        }
477        output.extend_from_slice(&ws.literal_buf);
478        return Ok(());
479    }
480
481    let table_data = &remaining[seq_count_size..];
482    let tables_consumed = parse_sequence_tables_ws(table_data, seq_tables, ws)?;
483
484    let seq_data = &table_data[tables_consumed..];
485
486    let before = output.len();
487
488    let result = decode_sequences_dispatch(
489        seq_data,
490        num_sequences,
491        seq_tables,
492        rep_offsets,
493        &ws.literal_buf,
494        output,
495        dict_history,
496    );
497    result?;
498    if output.len() - before > zrip_core::frame::MAX_BLOCK_SIZE {
499        return Err(DecompressError::BlockTooLarge);
500    }
501
502    Ok(())
503}
504
505#[inline(always)]
506pub(crate) fn decode_sequences_dispatch(
507    seq_data: &[u8],
508    num_sequences: u32,
509    seq_tables: &mut SequenceDecodeTables,
510    rep_offsets: &mut [u32; 3],
511    literals: &[u8],
512    output: &mut Vec<u8>,
513    history: &[u8],
514) -> Result<(), DecompressError> {
515    if num_sequences == 1 {
516        if history.is_empty() {
517            return decode_execute_single_sequence::<false>(
518                seq_data,
519                seq_tables,
520                rep_offsets,
521                literals,
522                output,
523                history,
524            );
525        }
526        return decode_execute_single_sequence::<true>(
527            seq_data,
528            seq_tables,
529            rep_offsets,
530            literals,
531            output,
532            history,
533        );
534    }
535
536    #[cfg(all(feature = "std", feature = "simd"))]
537    {
538        use std::sync::OnceLock;
539        static LEVEL: OnceLock<fearless_simd::Level> = OnceLock::new();
540        let level = *LEVEL.get_or_init(fearless_simd::Level::new);
541        return fearless_simd::dispatch!(level, _simd => {
542            if history.is_empty() {
543                decode_execute_sequences::<false>(
544                    seq_data,
545                    num_sequences,
546                    seq_tables,
547                    rep_offsets,
548                    literals,
549                    output,
550                    history,
551                )
552            } else {
553                decode_execute_sequences::<true>(
554                    seq_data,
555                    num_sequences,
556                    seq_tables,
557                    rep_offsets,
558                    literals,
559                    output,
560                    history,
561                )
562            }
563        });
564    }
565
566    #[allow(unreachable_code)]
567    if history.is_empty() {
568        decode_execute_sequences::<false>(
569            seq_data,
570            num_sequences,
571            seq_tables,
572            rep_offsets,
573            literals,
574            output,
575            history,
576        )
577    } else {
578        decode_execute_sequences::<true>(
579            seq_data,
580            num_sequences,
581            seq_tables,
582            rep_offsets,
583            literals,
584            output,
585            history,
586        )
587    }
588}
589
590#[cfg(test)]
591mod tests {
592    use super::*;
593    use alloc::vec::Vec;
594
595    fn push_block_header(out: &mut Vec<u8>, last: bool, block_type: u32, block_size: usize) {
596        let raw = ((block_size as u32) << 3) | (block_type << 1) | u32::from(last);
597        out.push(raw as u8);
598        out.push((raw >> 8) as u8);
599        out.push((raw >> 16) as u8);
600    }
601
602    #[test]
603    fn decompresses_frame_after_magic() {
604        let mut frame = Vec::new();
605        frame.push(0x20);
606        frame.push(5);
607        push_block_header(&mut frame, true, 0, 5);
608        frame.extend_from_slice(b"hello");
609
610        let mut output = Vec::new();
611        let mut ws = BlockDecodeWorkspace::new();
612        let consumed =
613            decompress_frame_after_magic(&frame, &mut output, usize::MAX, None, &mut ws).unwrap();
614        assert_eq!(consumed, frame.len());
615        assert_eq!(output, b"hello");
616    }
617
618    #[test]
619    fn frame_reset_keeps_explicit_huffman_cache() {
620        let mut ws = BlockDecodeWorkspace::new();
621        ws.huf_valid = true;
622        ws.huf_last_weights_valid = true;
623        ws.huf_last_weights.extend_from_slice(&[1, 2, 3]);
624
625        ws.reset_huffman_state();
626
627        assert!(!ws.huf_valid);
628        assert!(ws.huf_last_weights_valid);
629        assert_eq!(ws.huf_last_weights, [1, 2, 3]);
630    }
631}
632
633#[cfg(all(test, miri, not(feature = "paranoid")))]
634mod ub_tests {
635    use super::*;
636    use alloc::vec::Vec;
637    use zrip_core::bitstream::writer::BitWriter;
638    use zrip_core::frame::{MAX_BLOCK_SIZE, ZSTD_MAGIC};
639
640    fn push_block_header(out: &mut Vec<u8>, last: bool, block_type: u32, block_size: usize) {
641        let raw = ((block_size as u32) << 3) | (block_type << 1) | u32::from(last);
642        out.push(raw as u8);
643        out.push((raw >> 8) as u8);
644        out.push((raw >> 16) as u8);
645    }
646
647    fn frame_with_oversized_compressed_block_output() -> Vec<u8> {
648        let mut frame = Vec::new();
649        frame.extend_from_slice(&ZSTD_MAGIC.to_le_bytes());
650        frame.push(0x00);
651        frame.push(0x00);
652
653        push_block_header(&mut frame, false, 0, 1);
654        frame.push(b'A');
655
656        let mut block = Vec::new();
657        let trailing_literals = 65usize;
658        block.push(0x04 | (((trailing_literals & 0x0f) as u8) << 4));
659        block.push((trailing_literals >> 4) as u8);
660        block.extend(core::iter::repeat_n(b'B', trailing_literals));
661
662        block.push(1);
663        block.push(0x54);
664        block.extend_from_slice(&[0, 2, 52]);
665
666        let mut seq_bits = BitWriter::new();
667        let ml_extra = MAX_BLOCK_SIZE as u32 - 65_539;
668        seq_bits.write_bits(ml_extra, 16);
669        seq_bits.write_bits(0, 2);
670        seq_bits.close_reverse_stream();
671        block.extend_from_slice(&seq_bits.into_bytes());
672
673        push_block_header(&mut frame, true, 2, block.len());
674        frame.extend_from_slice(&block);
675        frame
676    }
677
678    #[test]
679    fn compressed_block_trailing_literals_overrun_wildcopy_headroom() {
680        // Issue: sequence execution reserves one block plus 64 bytes of wild-copy
681        // headroom and checks each sequence against that limit, but the final
682        // trailing literals are appended afterward without a matching capacity or
683        // block-size check. This safe frame seeds one prior byte, emits a 128 KiB
684        // match at offset 1, then appends 65 literals; miri reports the resulting
685        // out-of-bounds write in fast_extend_from_slice.
686        let frame = frame_with_oversized_compressed_block_output();
687        let _ = decompress(&frame);
688    }
689}
zrip_decode/lib.rs

zrip_decode/
lib.rs
