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 literals;
13pub(crate) mod ring_buffer;
14pub(crate) mod sequences;
15#[cfg(feature = "std")]
16pub mod streaming;
17
18#[cfg(all(
19    any(target_arch = "x86_64", target_arch = "aarch64"),
20    not(feature = "paranoid")
21))]
22pub(crate) mod simd_decode;
23
24#[cfg(feature = "alloc")]
25use alloc::boxed::Box;
26#[cfg(feature = "alloc")]
27use alloc::vec::Vec;
28
29use crate::exec::decode_execute_sequences;
30use crate::literals::decode_literals_ws;
31use crate::sequences::{SequenceDecodeTables, parse_sequence_count, parse_sequence_tables_ws};
32use zrip_core::block::{BlockType, parse_block_header};
33use zrip_core::error::DecompressError;
34use zrip_core::frame::MAX_WINDOW_SIZE;
35use zrip_core::frame::header::parse_frame_header;
36use zrip_core::huffman::HuffmanDecodeEntry;
37#[cfg(all(
38    any(target_arch = "x86_64", target_arch = "aarch64"),
39    not(feature = "paranoid")
40))]
41use zrip_core::simd::CpuTier;
42use zrip_core::xxhash::Xxh64State;
43
44pub(crate) struct BlockDecodeWorkspace {
45    pub literal_buf: Vec<u8>,
46    pub huf_table: Vec<HuffmanDecodeEntry>,
47    pub huf_table_log: u8,
48    pub huf_valid: bool,
49    pub huf_all_weights: Vec<u8>,
50    pub huf_rank_count: Vec<u32>,
51    pub huf_rank_start: Vec<u32>,
52    pub fse_dist: Vec<i16>,
53    pub fse_symbol_next: Vec<u16>,
54    pub fse_build_buf: Vec<zrip_core::fse::FseDecodeEntry>,
55}
56
57impl BlockDecodeWorkspace {
58    pub(crate) fn new() -> Self {
59        Self {
60            literal_buf: Vec::new(),
61            huf_table: Vec::new(),
62            huf_table_log: 0,
63            huf_valid: false,
64            huf_all_weights: Vec::new(),
65            huf_rank_count: Vec::new(),
66            huf_rank_start: Vec::new(),
67            fse_dist: Vec::new(),
68            fse_symbol_next: Vec::new(),
69            fse_build_buf: Vec::new(),
70        }
71    }
72}
73
74pub(crate) fn skip_skippable_frame(data: &[u8]) -> Option<usize> {
75    if data.len() < 8 {
76        return None;
77    }
78    let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
79    if (magic & 0xFFFF_FFF0) != 0x184D_2A50 {
80        return None;
81    }
82    let frame_size = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize;
83    let total = 8 + frame_size;
84    if total > data.len() {
85        return None;
86    }
87    Some(total)
88}
89
90pub fn decompress(input: &[u8]) -> Result<Vec<u8>, DecompressError> {
91    decompress_with_dict(input, None)
92}
93
94/// Decompress with an explicit output size limit.
95///
96/// Returns [`DecompressError::OutputTooSmall`] if the decompressed output would
97/// exceed `max_output_size` bytes. Use [`SAFE_DECOMPRESS_LIMIT`](zrip_core::SAFE_DECOMPRESS_LIMIT)
98/// when processing untrusted input to prevent memory exhaustion attacks.
99pub fn decompress_with_limit(
100    input: &[u8],
101    max_output_size: usize,
102) -> Result<Vec<u8>, DecompressError> {
103    let mut output = Vec::new();
104    let mut ws = Box::new(BlockDecodeWorkspace::new());
105    let mut offset = 0;
106    while offset < input.len() {
107        let remaining = &input[offset..];
108        if let Some(skip_len) = skip_skippable_frame(remaining) {
109            offset += skip_len;
110            continue;
111        }
112        let consumed = decompress_frame(remaining, &mut output, max_output_size, None, &mut ws)?;
113        offset += consumed;
114    }
115    Ok(output)
116}
117
118pub fn decompress_into(input: &[u8], output: &mut Vec<u8>) -> Result<usize, DecompressError> {
119    let max_output = zrip_core::DEFAULT_DECOMPRESS_LIMIT;
120    let mut ws = Box::new(BlockDecodeWorkspace::new());
121    let start = output.len();
122    let mut offset = 0;
123    while offset < input.len() {
124        let remaining = &input[offset..];
125        if let Some(skip_len) = skip_skippable_frame(remaining) {
126            offset += skip_len;
127            continue;
128        }
129        let consumed = decompress_frame(remaining, output, max_output, None, &mut ws)?;
130        offset += consumed;
131    }
132    Ok(output.len() - start)
133}
134
135pub fn decompress_with_dict(
136    input: &[u8],
137    dict: Option<&zrip_core::dict::Dictionary>,
138) -> Result<Vec<u8>, DecompressError> {
139    let max_output = zrip_core::DEFAULT_DECOMPRESS_LIMIT;
140    let mut output = Vec::new();
141    let mut ws = Box::new(BlockDecodeWorkspace::new());
142    let mut offset = 0;
143
144    while offset < input.len() {
145        let remaining = &input[offset..];
146        if let Some(skip_len) = skip_skippable_frame(remaining) {
147            offset += skip_len;
148            continue;
149        }
150        let consumed = decompress_frame(remaining, &mut output, max_output, dict, &mut ws)?;
151        offset += consumed;
152    }
153
154    Ok(output)
155}
156
157pub(crate) fn decompress_frame(
158    input: &[u8],
159    output: &mut Vec<u8>,
160    max_output: usize,
161    dict: Option<&zrip_core::dict::Dictionary>,
162    ws: &mut BlockDecodeWorkspace,
163) -> Result<usize, DecompressError> {
164    let header = parse_frame_header(input)?;
165
166    if header.window_size > MAX_WINDOW_SIZE {
167        return Err(DecompressError::WindowTooLarge {
168            requested: header.window_size,
169            max: MAX_WINDOW_SIZE,
170        });
171    }
172
173    if let Some(frame_dict_id) = header.dict_id {
174        match dict {
175            Some(d) if d.id() == frame_dict_id => {}
176            Some(d) => {
177                return Err(DecompressError::DictMismatch {
178                    expected: frame_dict_id,
179                    got: d.id(),
180                });
181            }
182            None => return Err(DecompressError::DictRequired),
183        }
184    }
185
186    if let Some(fcs) = header.frame_content_size {
187        if max_output < usize::MAX && fcs as usize > max_output {
188            return Err(DecompressError::OutputTooSmall);
189        }
190        let hint = (fcs as usize).min(MAX_WINDOW_SIZE as usize);
191        output.reserve(hint + 32);
192    }
193
194    let mut offset = header.header_size;
195    let output_start = output.len();
196
197    let dict_history: &[u8] = if let Some(d) = dict { d.content() } else { &[] };
198
199    let mut seq_tables = if let Some(d) = dict {
200        let mut st = SequenceDecodeTables::new_default();
201        if let Some((t, l)) = d.of_table() {
202            st.of_table = zrip_core::fse::promote_of_table(t);
203            st.of_accuracy = l;
204            st.of_set = true;
205        }
206        if let Some((t, l)) = d.ml_table() {
207            st.ml_table = zrip_core::fse::promote_ml_table(t);
208            st.ml_accuracy = l;
209            st.ml_set = true;
210        }
211        if let Some((t, l)) = d.ll_table() {
212            st.ll_table = zrip_core::fse::promote_ll_table(t);
213            st.ll_accuracy = l;
214            st.ll_set = true;
215        }
216        st
217    } else {
218        SequenceDecodeTables::new_default()
219    };
220    let mut rep_offsets: [u32; 3] = if let Some(d) = dict {
221        *d.rep_offsets()
222    } else {
223        [1, 4, 8]
224    };
225    ws.huf_valid = false;
226    if let Some(d) = dict {
227        if let Some((t, l)) = d.huf_table() {
228            ws.huf_table.clear();
229            ws.huf_table.extend_from_slice(t);
230            ws.huf_table_log = l;
231            ws.huf_valid = true;
232        }
233    }
234
235    let mut hasher = if header.content_checksum {
236        Some(Xxh64State::new(0))
237    } else {
238        None
239    };
240
241    loop {
242        if offset + 3 > input.len() {
243            return Err(DecompressError::InputExhausted);
244        }
245        let block_header = parse_block_header(&input[offset..])?;
246        offset += 3;
247
248        let block_size = block_header.block_size as usize;
249
250        if block_size > zrip_core::frame::MAX_BLOCK_SIZE {
251            match block_header.block_type {
252                BlockType::Raw | BlockType::Rle => {
253                    return Err(DecompressError::BlockTooLarge);
254                }
255                BlockType::Compressed => {}
256            }
257        }
258
259        match block_header.block_type {
260            BlockType::Raw => {
261                if offset + block_size > input.len() {
262                    return Err(DecompressError::InputExhausted);
263                }
264                if output.len() - output_start + block_size > max_output {
265                    return Err(DecompressError::OutputTooSmall);
266                }
267                output.extend_from_slice(&input[offset..offset + block_size]);
268                offset += block_size;
269            }
270            BlockType::Rle => {
271                if offset >= input.len() {
272                    return Err(DecompressError::InputExhausted);
273                }
274                if output.len() - output_start + block_size > max_output {
275                    return Err(DecompressError::OutputTooSmall);
276                }
277                let byte = input[offset];
278                output.resize(output.len() + block_size, byte);
279                offset += 1;
280            }
281            BlockType::Compressed => {
282                if offset + block_size > input.len() {
283                    return Err(DecompressError::InputExhausted);
284                }
285                let block_data = &input[offset..offset + block_size];
286                decode_compressed_block(
287                    block_data,
288                    output,
289                    output_start,
290                    max_output,
291                    &mut seq_tables,
292                    &mut rep_offsets,
293                    ws,
294                    dict_history,
295                )?;
296                offset += block_size;
297            }
298        }
299
300        if block_header.last_block {
301            break;
302        }
303    }
304
305    if let Some(ref mut hasher) = hasher {
306        hasher.update(&output[output_start..]);
307        let hash = hasher.finish();
308        let expected_checksum = (hash & 0xFFFF_FFFF) as u32;
309
310        if offset + 4 > input.len() {
311            return Err(DecompressError::InputExhausted);
312        }
313        let stored_checksum = u32::from_le_bytes([
314            input[offset],
315            input[offset + 1],
316            input[offset + 2],
317            input[offset + 3],
318        ]);
319        offset += 4;
320
321        if expected_checksum != stored_checksum {
322            return Err(DecompressError::ChecksumMismatch {
323                expected: stored_checksum,
324                got: expected_checksum,
325            });
326        }
327    }
328
329    if let Some(fcs) = header.frame_content_size {
330        if (output.len() - output_start) as u64 != fcs {
331            return Err(DecompressError::FrameSizeMismatch);
332        }
333    }
334
335    Ok(offset)
336}
337
338#[allow(clippy::too_many_arguments)]
339fn decode_compressed_block(
340    data: &[u8],
341    output: &mut Vec<u8>,
342    output_start: usize,
343    max_output: usize,
344    seq_tables: &mut SequenceDecodeTables,
345    rep_offsets: &mut [u32; 3],
346    ws: &mut BlockDecodeWorkspace,
347    dict_history: &[u8],
348) -> Result<(), DecompressError> {
349    let lit_consumed = decode_literals_ws(data, ws)?;
350
351    let remaining = &data[lit_consumed..];
352
353    if remaining.is_empty() {
354        if output.len() - output_start + ws.literal_buf.len() > max_output {
355            return Err(DecompressError::OutputTooSmall);
356        }
357        output.extend_from_slice(&ws.literal_buf);
358        return Ok(());
359    }
360
361    let (num_sequences, seq_count_size) = parse_sequence_count(remaining)?;
362
363    if num_sequences == 0 {
364        if output.len() - output_start + ws.literal_buf.len() > max_output {
365            return Err(DecompressError::OutputTooSmall);
366        }
367        output.extend_from_slice(&ws.literal_buf);
368        return Ok(());
369    }
370
371    let table_data = &remaining[seq_count_size..];
372    let tables_consumed = parse_sequence_tables_ws(table_data, seq_tables, ws)?;
373
374    let seq_data = &table_data[tables_consumed..];
375
376    let before = output.len();
377
378    #[cfg(all(target_arch = "x86_64", not(feature = "paranoid")))]
379    {
380        if zrip_core::simd::cpu_tier() >= CpuTier::Avx2 {
381            decode_execute_block_avx2(
382                seq_data,
383                num_sequences,
384                seq_tables,
385                rep_offsets,
386                &ws.literal_buf,
387                output,
388                dict_history,
389            )?;
390            if output.len() - before > zrip_core::frame::MAX_BLOCK_SIZE {
391                return Err(DecompressError::BlockTooLarge);
392            }
393            return Ok(());
394        }
395    }
396    #[cfg(all(target_arch = "aarch64", not(feature = "paranoid")))]
397    {
398        if zrip_core::simd::cpu_tier() >= CpuTier::Neon {
399            decode_execute_block_neon(
400                seq_data,
401                num_sequences,
402                seq_tables,
403                rep_offsets,
404                &ws.literal_buf,
405                output,
406                dict_history,
407            )?;
408            if output.len() - before > zrip_core::frame::MAX_BLOCK_SIZE {
409                return Err(DecompressError::BlockTooLarge);
410            }
411            return Ok(());
412        }
413    }
414
415    decode_execute_sequences(
416        seq_data,
417        num_sequences,
418        seq_tables,
419        rep_offsets,
420        &ws.literal_buf,
421        output,
422        dict_history,
423    )?;
424    if output.len() - before > zrip_core::frame::MAX_BLOCK_SIZE {
425        return Err(DecompressError::BlockTooLarge);
426    }
427
428    Ok(())
429}
430
431#[cfg(all(target_arch = "x86_64", not(feature = "paranoid")))]
432fn decode_execute_block_avx2(
433    seq_data: &[u8],
434    num_sequences: u32,
435    tables: &mut SequenceDecodeTables,
436    rep_offsets: &mut [u32; 3],
437    literals: &[u8],
438    output: &mut Vec<u8>,
439    history: &[u8],
440) -> Result<(), DecompressError> {
441    crate::simd_decode::x86_64::decode::decode_execute_avx2_safe(
442        seq_data,
443        num_sequences,
444        tables,
445        rep_offsets,
446        literals,
447        output,
448        history,
449    )
450}
451
452#[cfg(all(target_arch = "aarch64", not(feature = "paranoid")))]
453fn decode_execute_block_neon(
454    seq_data: &[u8],
455    num_sequences: u32,
456    tables: &mut SequenceDecodeTables,
457    rep_offsets: &mut [u32; 3],
458    literals: &[u8],
459    output: &mut Vec<u8>,
460    history: &[u8],
461) -> Result<(), DecompressError> {
462    crate::simd_decode::aarch64::decode::decode_execute_neon_safe(
463        seq_data,
464        num_sequences,
465        tables,
466        rep_offsets,
467        literals,
468        output,
469        history,
470    )
471}
zrip_decode/lib.rs

zrip_decode/
lib.rs
