zrip_decode/
lib.rs

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

zrip_decode/
lib.rs
