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