zstd_zero/
lib.rs

1#![doc = include_str!("../README.md")]
2#![no_std]
3#![forbid(unsafe_code)]
4
5#[cfg(test)]
6extern crate std;
7
8mod bitstream;
9mod error;
10mod fse;
11mod huffman;
12mod xxhash;
13
14use bitstream::BackwardBits;
15pub use error::DecodeError;
16use fse::Table as FseTable;
17use huffman::Table as HuffmanTable;
18use xxhash::XxHash64;
19
20pub use fse::Entry as FseEntry;
21pub use huffman::Entry as HuffmanEntry;
22/// FSE slots required for three sequence tables and Huffman weight scratch.
23pub const FSE_ENTRIES: usize = 4 * 512;
24/// Number of scratch counters shared by all FSE table constructors.
25/// Contents may be arbitrary on entry and are not retained across tables.
26pub const FSE_SCRATCH_LEN: usize = fse::MAX_SYMBOLS;
27/// Huffman slots required for literal decoding.
28pub const HUFFMAN_ENTRIES: usize = 2048;
29
30pub const MAX_BLOCK_SIZE: usize = 128 * 1024;
31pub const MAX_FRAME_HEADER_SIZE: usize = 18;
32
33const ZSTD_MAGIC: u32 = 0xfd2f_b528;
34const SKIPPABLE_MAGIC_MIN: u32 = 0x184d_2a50;
35const SKIPPABLE_MAGIC_MAX: u32 = 0x184d_2a5f;
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
38pub struct FrameHeader {
39    pub window_size: u64,
40    pub content_size: Option<u64>,
41    pub dictionary_id: u32,
42    pub has_checksum: bool,
43}
44
45#[derive(Clone, Copy, Debug, Eq, PartialEq)]
46pub enum StreamHeader {
47    Zstandard(FrameHeader),
48    Skippable { magic: u32, size: u32 },
49}
50
51#[derive(Clone, Copy, Debug, Eq, PartialEq)]
52pub enum FrameKind {
53    Zstandard,
54    Skippable { magic: u32, size: u32 },
55}
56
57#[derive(Clone, Copy, Debug, Eq, PartialEq)]
58pub enum HeaderStatus {
59    NeedMore { minimum: usize },
60    Complete { header: StreamHeader, size: usize },
61}
62
63#[derive(Debug, Eq, PartialEq)]
64pub enum DecodeStep<'a> {
65    NeedInput {
66        consumed: usize,
67    },
68    FrameStarted {
69        consumed: usize,
70        header: StreamHeader,
71    },
72    Output {
73        consumed: usize,
74        bytes: &'a [u8],
75    },
76    FrameFinished {
77        consumed: usize,
78        kind: FrameKind,
79    },
80}
81
82impl DecodeStep<'_> {
83    pub const fn consumed(&self) -> usize {
84        match self {
85            Self::NeedInput { consumed }
86            | Self::FrameStarted { consumed, .. }
87            | Self::Output { consumed, .. }
88            | Self::FrameFinished { consumed, .. } => *consumed,
89        }
90    }
91}
92
93/// A failure while driving a [`Decoder`] through its callback API.
94#[derive(Debug, Eq, PartialEq)]
95pub enum StreamError<E> {
96    Decode(DecodeError),
97    Output(E),
98    DecoderStalled,
99}
100
101impl<E: core::fmt::Display> core::fmt::Display for StreamError<E> {
102    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
103        match self {
104            Self::Decode(error) => write!(formatter, "Zstandard decode failed: {error}"),
105            Self::Output(error) => write!(formatter, "Zstandard output failed: {error}"),
106            Self::DecoderStalled => {
107                formatter.write_str("Zstandard decoder stopped making progress")
108            }
109        }
110    }
111}
112
113enum InternalStep {
114    NeedInput {
115        consumed: usize,
116    },
117    FrameStarted {
118        consumed: usize,
119        header: StreamHeader,
120    },
121    Output {
122        consumed: usize,
123        start: usize,
124        length: usize,
125    },
126    FrameFinished {
127        consumed: usize,
128        kind: FrameKind,
129    },
130}
131
132/// Initialized caller-owned storage, retained for the decoder's lifetime.
133pub struct DecoderBuffers<'a> {
134    /// At least [`FSE_SCRATCH_LEN`] elements for probabilities and next states.
135    pub fse_scratch: &'a mut [i16],
136    /// At least [`FSE_ENTRIES`] entries, including Huffman weight scratch.
137    pub fse: &'a mut [FseEntry],
138    /// At least [`HUFFMAN_ENTRIES`] entries.
139    pub huffman: &'a mut [HuffmanEntry],
140    pub history: &'a mut [u8],
141    pub block: &'a mut [u8],
142    pub literals: &'a mut [u8],
143}
144
145/// Tuning knobs for a [`Decoder`].
146///
147/// [`Default`] enables every validation check; prefer it unless you have a
148/// specific reason not to.
149#[derive(Clone, Copy, Debug, Eq, PartialEq)]
150pub struct DecoderOptions {
151    /// Reject Huffman literal streams that do not end exactly on their final bit.
152    ///
153    /// Enabled by default. Disabling it makes the decoder tolerate a truncated
154    /// literal bitstream by padding the remaining symbols instead of failing,
155    /// which is only useful for best-effort salvage of damaged data.
156    ///
157    /// # Warning
158    ///
159    /// With this disabled, a corrupt frame can decode "successfully" to bytes
160    /// that differ from what other zstd implementations produce, with no error
161    /// reported. Never disable it for content-addressed or otherwise
162    /// security-sensitive input.
163    pub strict_literal_bitstream: bool,
164}
165
166impl Default for DecoderOptions {
167    fn default() -> Self {
168        Self {
169            strict_literal_bitstream: true,
170        }
171    }
172}
173
174#[derive(Clone, Copy)]
175enum State {
176    Header,
177    Skippable {
178        magic: u32,
179        size: u32,
180        remaining: usize,
181    },
182    BlockHeader,
183    BlockPayload {
184        header: BlockHeader,
185        filled: usize,
186    },
187    Checksum {
188        filled: usize,
189    },
190    FrameDone(FrameKind),
191}
192
193#[derive(Clone, Copy)]
194struct BlockHeader {
195    last: bool,
196    kind: BlockKind,
197    size: usize,
198}
199
200#[derive(Clone, Copy, Eq, PartialEq)]
201enum BlockKind {
202    Raw,
203    Rle,
204    Compressed,
205}
206
207pub struct Decoder<'a> {
208    fse_scratch: &'a mut [i16],
209    options: DecoderOptions,
210    history: &'a mut [u8],
211    block: &'a mut [u8],
212    literals: &'a mut [u8],
213    state: State,
214    poisoned: bool,
215    completed_frames: u64,
216    header_buffer: [u8; MAX_FRAME_HEADER_SIZE],
217    header_len: usize,
218    block_header_buffer: [u8; 3],
219    block_header_len: usize,
220    checksum_buffer: [u8; 4],
221    current_header: Option<FrameHeader>,
222    block_limit: usize,
223    history_position: usize,
224    frame_output: u64,
225    pending_start: usize,
226    pending_len: usize,
227    offsets: [u32; 3],
228    checksum: XxHash64,
229    huffman: HuffmanTable<'a>,
230    literal_lengths: FseTable<'a>,
231    offsets_table: FseTable<'a>,
232    match_lengths: FseTable<'a>,
233}
234
235impl<'a> Decoder<'a> {
236    /// Create a decoder with every validation check enabled.
237    ///
238    /// Returns [`DecodeError::InvalidEntropyTable`] if entropy storage is too short.
239    pub fn new(buffers: DecoderBuffers<'a>) -> Result<Self, DecodeError> {
240        Self::with_options(buffers, DecoderOptions::default())
241    }
242
243    /// Create a decoder with explicit [`DecoderOptions`].
244    ///
245    /// Returns [`DecodeError::InvalidEntropyTable`] if entropy storage is too short.
246    pub fn with_options(
247        buffers: DecoderBuffers<'a>,
248        options: DecoderOptions,
249    ) -> Result<Self, DecodeError> {
250        if buffers.fse.len() < FSE_ENTRIES
251            || buffers.huffman.len() < HUFFMAN_ENTRIES
252            || buffers.fse_scratch.len() < FSE_SCRATCH_LEN
253        {
254            return Err(DecodeError::InvalidEntropyTable);
255        }
256        let (scratch, rest) = buffers.fse[..FSE_ENTRIES].split_at_mut(512);
257        let (literal, rest) = rest.split_at_mut(512);
258        let (offsets, matches) = rest.split_at_mut(512);
259        Ok(Self {
260            fse_scratch: &mut buffers.fse_scratch[..FSE_SCRATCH_LEN],
261            options,
262            history: buffers.history,
263            block: buffers.block,
264            literals: buffers.literals,
265            state: State::Header,
266            poisoned: false,
267            completed_frames: 0,
268            header_buffer: [0; MAX_FRAME_HEADER_SIZE],
269            header_len: 0,
270            block_header_buffer: [0; 3],
271            block_header_len: 0,
272            checksum_buffer: [0; 4],
273            current_header: None,
274            block_limit: 0,
275            history_position: 0,
276            frame_output: 0,
277            pending_start: 0,
278            pending_len: 0,
279            offsets: [1, 4, 8],
280            checksum: XxHash64::new(),
281            huffman: HuffmanTable::new(
282                &mut buffers.huffman[..HUFFMAN_ENTRIES],
283                FseTable::new(scratch),
284            ),
285            literal_lengths: FseTable::new(literal),
286            offsets_table: FseTable::new(offsets),
287            match_lengths: FseTable::new(matches),
288        })
289    }
290
291    pub fn decode<'decoder>(
292        &'decoder mut self,
293        input: &[u8],
294    ) -> Result<DecodeStep<'decoder>, DecodeError> {
295        if self.poisoned {
296            return Err(DecodeError::DecoderPoisoned);
297        }
298        let step = match self.decode_inner(input) {
299            Ok(step) => step,
300            Err(error) => {
301                self.poisoned = true;
302                return Err(error);
303            }
304        };
305        Ok(match step {
306            InternalStep::NeedInput { consumed } => DecodeStep::NeedInput { consumed },
307            InternalStep::FrameStarted { consumed, header } => {
308                DecodeStep::FrameStarted { consumed, header }
309            }
310            InternalStep::Output {
311                consumed,
312                start,
313                length,
314            } => DecodeStep::Output {
315                consumed,
316                bytes: &self.history[start..start + length],
317            },
318            InternalStep::FrameFinished { consumed, kind } => {
319                DecodeStep::FrameFinished { consumed, kind }
320            }
321        })
322    }
323
324    /// Consumes an input fragment and sends decoded bytes to `output`.
325    ///
326    /// This is the convenient callback-driven counterpart to [`decode`](Self::decode).
327    /// Frame events are handled internally while output errors remain distinct
328    /// from compressed-stream errors.
329    pub fn push<E>(
330        &mut self,
331        mut input: &[u8],
332        mut output: impl FnMut(&[u8]) -> Result<(), E>,
333    ) -> Result<(), StreamError<E>> {
334        let mut idle_steps = 0;
335        while !input.is_empty() {
336            let step = self.decode(input).map_err(StreamError::Decode)?;
337            let consumed = step.consumed();
338            if consumed > input.len() {
339                return Err(StreamError::DecoderStalled);
340            }
341            let needs_input = matches!(step, DecodeStep::NeedInput { .. });
342            let produced = match step {
343                DecodeStep::Output { bytes, .. } if !bytes.is_empty() => {
344                    output(bytes).map_err(StreamError::Output)?;
345                    true
346                }
347                _ => false,
348            };
349            input = &input[consumed..];
350            if needs_input && !input.is_empty() {
351                return Err(StreamError::DecoderStalled);
352            }
353            if consumed != 0 || produced {
354                idle_steps = 0;
355            } else {
356                idle_steps += 1;
357                if idle_steps > 16 {
358                    return Err(StreamError::DecoderStalled);
359                }
360            }
361        }
362        Ok(())
363    }
364
365    /// Drains buffered output and validates the end of the stream.
366    pub fn finish_with<E>(
367        &mut self,
368        mut output: impl FnMut(&[u8]) -> Result<(), E>,
369    ) -> Result<(), StreamError<E>> {
370        let mut idle_steps = 0;
371        loop {
372            match self.decode(&[]).map_err(StreamError::Decode)? {
373                DecodeStep::NeedInput { .. } => break,
374                DecodeStep::Output { bytes, .. } if !bytes.is_empty() => {
375                    output(bytes).map_err(StreamError::Output)?;
376                    idle_steps = 0;
377                }
378                _ => {
379                    idle_steps += 1;
380                    if idle_steps > 16 {
381                        return Err(StreamError::DecoderStalled);
382                    }
383                }
384            }
385        }
386        self.finish().map_err(StreamError::Decode)
387    }
388
389    pub fn finish(&self) -> Result<(), DecodeError> {
390        if self.poisoned {
391            return Err(DecodeError::DecoderPoisoned);
392        }
393        if self.pending_len != 0 || self.header_len != 0 || self.completed_frames == 0 {
394            return Err(DecodeError::UnexpectedEof);
395        }
396        match self.state {
397            State::Header | State::FrameDone(_) => Ok(()),
398            _ => Err(DecodeError::UnexpectedEof),
399        }
400    }
401
402    pub fn reset(&mut self) {
403        self.state = State::Header;
404        self.poisoned = false;
405        self.completed_frames = 0;
406        self.header_len = 0;
407        self.block_header_len = 0;
408        self.current_header = None;
409        self.reset_frame();
410    }
411
412    fn reset_frame(&mut self) {
413        self.frame_output = 0;
414        self.pending_len = 0;
415        self.offsets = [1, 4, 8];
416        self.checksum = XxHash64::new();
417        self.huffman.reset();
418        self.literal_lengths.reset();
419        self.offsets_table.reset();
420        self.match_lengths.reset();
421    }
422
423    fn decode_inner(&mut self, input: &[u8]) -> Result<InternalStep, DecodeError> {
424        if self.pending_len != 0 {
425            let start = self.pending_start;
426            let length = core::cmp::min(self.pending_len, self.history.len() - start);
427            self.pending_start = (start + length) % self.history.len();
428            self.pending_len -= length;
429            return Ok(InternalStep::Output {
430                consumed: 0,
431                start,
432                length,
433            });
434        }
435
436        let mut consumed = 0usize;
437        loop {
438            match self.state {
439                State::Header => match inspect_frame(&self.header_buffer[..self.header_len])? {
440                    HeaderStatus::NeedMore { minimum } => {
441                        if consumed == input.len() {
442                            return Ok(InternalStep::NeedInput { consumed });
443                        }
444                        let amount =
445                            core::cmp::min(minimum - self.header_len, input.len() - consumed);
446                        self.header_buffer[self.header_len..self.header_len + amount]
447                            .copy_from_slice(&input[consumed..consumed + amount]);
448                        self.header_len += amount;
449                        consumed += amount;
450                    }
451                    HeaderStatus::Complete { header, .. } => {
452                        self.header_len = 0;
453                        match header {
454                            StreamHeader::Zstandard(frame) => {
455                                self.start_frame(frame)?;
456                                self.state = State::BlockHeader;
457                            }
458                            StreamHeader::Skippable { magic, size } => {
459                                self.state = State::Skippable {
460                                    magic,
461                                    size,
462                                    remaining: size as usize,
463                                };
464                            }
465                        }
466                        return Ok(InternalStep::FrameStarted { consumed, header });
467                    }
468                },
469                State::Skippable {
470                    magic,
471                    size,
472                    remaining,
473                } => {
474                    let amount = core::cmp::min(remaining, input.len() - consumed);
475                    consumed += amount;
476                    let remaining = remaining - amount;
477                    if remaining == 0 {
478                        self.state = State::FrameDone(FrameKind::Skippable { magic, size });
479                    } else {
480                        self.state = State::Skippable {
481                            magic,
482                            size,
483                            remaining,
484                        };
485                        return Ok(InternalStep::NeedInput { consumed });
486                    }
487                }
488                State::BlockHeader => {
489                    let amount = core::cmp::min(3 - self.block_header_len, input.len() - consumed);
490                    self.block_header_buffer[self.block_header_len..self.block_header_len + amount]
491                        .copy_from_slice(&input[consumed..consumed + amount]);
492                    self.block_header_len += amount;
493                    consumed += amount;
494                    if self.block_header_len != 3 {
495                        return Ok(InternalStep::NeedInput { consumed });
496                    }
497                    self.block_header_len = 0;
498                    let header = parse_block_header(self.block_header_buffer, self.block_limit)?;
499                    let payload_size = if header.kind == BlockKind::Rle {
500                        1
501                    } else {
502                        header.size
503                    };
504                    if payload_size > self.block.len() {
505                        return Err(DecodeError::BlockScratchTooSmall {
506                            required: payload_size,
507                            provided: self.block.len(),
508                        });
509                    }
510                    self.state = State::BlockPayload { header, filled: 0 };
511                }
512                State::BlockPayload { header, filled } => {
513                    let payload_size = if header.kind == BlockKind::Rle {
514                        1
515                    } else {
516                        header.size
517                    };
518                    let amount = core::cmp::min(payload_size - filled, input.len() - consumed);
519                    self.block[filled..filled + amount]
520                        .copy_from_slice(&input[consumed..consumed + amount]);
521                    consumed += amount;
522                    let filled = filled + amount;
523                    if filled != payload_size {
524                        self.state = State::BlockPayload { header, filled };
525                        return Ok(InternalStep::NeedInput { consumed });
526                    }
527                    self.process_block(header)?;
528                    self.state = if header.last {
529                        if self.current_header()?.has_checksum {
530                            self.checksum_buffer = [0; 4];
531                            State::Checksum { filled: 0 }
532                        } else {
533                            self.validate_frame_end()?;
534                            State::FrameDone(FrameKind::Zstandard)
535                        }
536                    } else {
537                        State::BlockHeader
538                    };
539                    if self.pending_len != 0 {
540                        let start = self.pending_start;
541                        let length = core::cmp::min(self.pending_len, self.history.len() - start);
542                        self.pending_start = (start + length) % self.history.len();
543                        self.pending_len -= length;
544                        return Ok(InternalStep::Output {
545                            consumed,
546                            start,
547                            length,
548                        });
549                    }
550                }
551                State::Checksum { filled } => {
552                    let amount = core::cmp::min(4 - filled, input.len() - consumed);
553                    self.checksum_buffer[filled..filled + amount]
554                        .copy_from_slice(&input[consumed..consumed + amount]);
555                    consumed += amount;
556                    let filled = filled + amount;
557                    if filled != 4 {
558                        self.state = State::Checksum { filled };
559                        return Ok(InternalStep::NeedInput { consumed });
560                    }
561                    let expected = u32::from_le_bytes(self.checksum_buffer);
562                    let actual = self.checksum.digest() as u32;
563                    if actual != expected {
564                        return Err(DecodeError::ChecksumMismatch { expected, actual });
565                    }
566                    self.validate_frame_end()?;
567                    self.state = State::FrameDone(FrameKind::Zstandard);
568                }
569                State::FrameDone(kind) => {
570                    self.completed_frames = self
571                        .completed_frames
572                        .checked_add(1)
573                        .ok_or(DecodeError::ArithmeticOverflow)?;
574                    self.current_header = None;
575                    self.state = State::Header;
576                    return Ok(InternalStep::FrameFinished { consumed, kind });
577                }
578            }
579        }
580    }
581
582    fn start_frame(&mut self, header: FrameHeader) -> Result<(), DecodeError> {
583        if header.dictionary_id != 0 {
584            return Err(DecodeError::UnsupportedDictionary {
585                id: header.dictionary_id,
586            });
587        }
588        let window =
589            usize::try_from(header.window_size).map_err(|_| DecodeError::WindowTooLarge)?;
590        if window > self.history.len() {
591            return Err(DecodeError::HistoryTooSmall {
592                required: window,
593                provided: self.history.len(),
594            });
595        }
596        self.current_header = Some(header);
597        self.block_limit = core::cmp::min(window, MAX_BLOCK_SIZE);
598        self.reset_frame();
599        Ok(())
600    }
601
602    fn current_header(&self) -> Result<FrameHeader, DecodeError> {
603        self.current_header.ok_or(DecodeError::InvalidFrameHeader)
604    }
605
606    fn process_block(&mut self, header: BlockHeader) -> Result<(), DecodeError> {
607        let start = self.history_position;
608        let before = self.frame_output;
609        match header.kind {
610            BlockKind::Raw => {
611                for index in 0..header.size {
612                    self.write_byte(self.block[index])?;
613                }
614            }
615            BlockKind::Rle => {
616                let byte = self.block[0];
617                for _ in 0..header.size {
618                    self.write_byte(byte)?;
619                }
620            }
621            BlockKind::Compressed => self.decode_compressed(header.size)?,
622        }
623        let produced = usize::try_from(self.frame_output - before)
624            .map_err(|_| DecodeError::ArithmeticOverflow)?;
625        if produced > self.block_limit {
626            return Err(DecodeError::InvalidBlock);
627        }
628        self.record_output(start, produced);
629        Ok(())
630    }
631
632    fn decode_compressed(&mut self, size: usize) -> Result<(), DecodeError> {
633        let (literal_count, consumed) = decode_literals(
634            &self.block[..size],
635            self.literals,
636            &mut self.huffman,
637            self.fse_scratch,
638            self.block_limit,
639            self.options.strict_literal_bitstream,
640        )?;
641        let sequence_input = &self.block[consumed..size];
642        let (sequence_count, mut position) = parse_sequence_count(sequence_input)?;
643        if sequence_count == 0 {
644            if position != sequence_input.len() {
645                return Err(DecodeError::InvalidBlock);
646            }
647            let window_size = self.current_header()?.window_size;
648            let mut writer = HistoryWriter {
649                history: self.history,
650                position: &mut self.history_position,
651                frame_output: &mut self.frame_output,
652                window_size,
653            };
654            for byte in &self.literals[..literal_count] {
655                writer.write(*byte)?;
656            }
657            return Ok(());
658        }
659        let modes = *sequence_input
660            .get(position)
661            .ok_or(DecodeError::InvalidBlock)?;
662        position += 1;
663        if modes & 3 != 0 {
664            return Err(DecodeError::InvalidBlock);
665        }
666        position += build_sequence_table(
667            &mut self.literal_lengths,
668            self.fse_scratch,
669            modes >> 6,
670            &sequence_input[position..],
671            &LL_DEFAULT,
672            6,
673            35,
674            9,
675        )?;
676        position += build_sequence_table(
677            &mut self.offsets_table,
678            self.fse_scratch,
679            (modes >> 4) & 3,
680            &sequence_input[position..],
681            &OF_DEFAULT,
682            5,
683            31,
684            8,
685        )?;
686        position += build_sequence_table(
687            &mut self.match_lengths,
688            self.fse_scratch,
689            (modes >> 2) & 3,
690            &sequence_input[position..],
691            &ML_DEFAULT,
692            6,
693            52,
694            9,
695        )?;
696        let mut bits = BackwardBits::new(
697            sequence_input
698                .get(position..)
699                .ok_or(DecodeError::InvalidBitstream)?,
700        )?;
701        let mut ll_state = bits.read(self.literal_lengths.log())?;
702        let mut of_state = bits.read(self.offsets_table.log())?;
703        let mut ml_state = bits.read(self.match_lengths.log())?;
704        let mut literal_position = 0usize;
705        let window_size = self.current_header()?.window_size;
706        let literal_lengths = &self.literal_lengths;
707        let offsets_table = &self.offsets_table;
708        let match_lengths = &self.match_lengths;
709        let literals = &self.literals[..literal_count];
710        let repeated_offsets = &mut self.offsets;
711        let mut writer = HistoryWriter {
712            history: self.history,
713            position: &mut self.history_position,
714            frame_output: &mut self.frame_output,
715            window_size,
716        };
717
718        for sequence in 0..sequence_count {
719            let ll_code = literal_lengths.symbol(ll_state)? as usize;
720            let of_code = offsets_table.symbol(of_state)?;
721            let ml_code = match_lengths.symbol(ml_state)? as usize;
722            if ll_code >= LL_BASE.len() || ml_code >= ML_BASE.len() || of_code > 31 {
723                return Err(DecodeError::InvalidEntropyTable);
724            }
725            let raw_offset = (1u32 << of_code)
726                .checked_add(bits.read(of_code)?)
727                .ok_or(DecodeError::ArithmeticOverflow)?;
728            let match_length = ML_BASE[ml_code]
729                .checked_add(bits.read(ML_BITS[ml_code])?)
730                .ok_or(DecodeError::ArithmeticOverflow)?;
731            let literal_length = LL_BASE[ll_code]
732                .checked_add(bits.read(LL_BITS[ll_code])?)
733                .ok_or(DecodeError::ArithmeticOverflow)?;
734            let literal_end = literal_position
735                .checked_add(literal_length as usize)
736                .ok_or(DecodeError::ArithmeticOverflow)?;
737            if literal_end > literal_count {
738                return Err(DecodeError::InvalidBlock);
739            }
740            for byte in &literals[literal_position..literal_end] {
741                writer.write(*byte)?;
742            }
743            literal_position = literal_end;
744            let offset = resolve_offset(raw_offset, literal_length, repeated_offsets)?;
745            writer.copy_match(offset as usize, match_length as usize)?;
746
747            if sequence + 1 != sequence_count {
748                literal_lengths.update(&mut ll_state, &mut bits)?;
749                match_lengths.update(&mut ml_state, &mut bits)?;
750                offsets_table.update(&mut of_state, &mut bits)?;
751            }
752        }
753        if bits.remaining() != 0 {
754            return Err(DecodeError::InvalidBitstream);
755        }
756        for byte in &literals[literal_position..literal_count] {
757            writer.write(*byte)?;
758        }
759        Ok(())
760    }
761
762    fn write_byte(&mut self, byte: u8) -> Result<(), DecodeError> {
763        write_history_byte(
764            self.history,
765            &mut self.history_position,
766            &mut self.frame_output,
767            byte,
768        )
769    }
770
771    fn record_output(&mut self, start: usize, length: usize) {
772        if length == 0 {
773            return;
774        }
775        let first = core::cmp::min(length, self.history.len() - start);
776        self.checksum.update(&self.history[start..start + first]);
777        if first != length {
778            self.checksum.update(&self.history[..length - first]);
779        }
780        self.pending_start = start;
781        self.pending_len = length;
782    }
783
784    fn validate_frame_end(&self) -> Result<(), DecodeError> {
785        if let Some(expected) = self.current_header()?.content_size {
786            if self.frame_output != expected {
787                return Err(DecodeError::ContentSizeMismatch {
788                    expected,
789                    actual: self.frame_output,
790                });
791            }
792        }
793        Ok(())
794    }
795}
796
797struct HistoryWriter<'a> {
798    history: &'a mut [u8],
799    position: &'a mut usize,
800    frame_output: &'a mut u64,
801    window_size: u64,
802}
803
804impl HistoryWriter<'_> {
805    fn write(&mut self, byte: u8) -> Result<(), DecodeError> {
806        write_history_byte(self.history, self.position, self.frame_output, byte)
807    }
808
809    fn copy_match(&mut self, offset: usize, length: usize) -> Result<(), DecodeError> {
810        let available = core::cmp::min(*self.frame_output, self.window_size);
811        if offset == 0 || offset as u64 > available {
812            return Err(DecodeError::InvalidOffset);
813        }
814        for _ in 0..length {
815            let source = (*self.position + self.history.len() - offset) % self.history.len();
816            self.write(self.history[source])?;
817        }
818        Ok(())
819    }
820}
821
822fn write_history_byte(
823    history: &mut [u8],
824    position: &mut usize,
825    frame_output: &mut u64,
826    byte: u8,
827) -> Result<(), DecodeError> {
828    if history.is_empty() {
829        return Err(DecodeError::HistoryTooSmall {
830            required: 1,
831            provided: 0,
832        });
833    }
834    history[*position] = byte;
835    *position += 1;
836    if *position == history.len() {
837        *position = 0;
838    }
839    *frame_output = frame_output
840        .checked_add(1)
841        .ok_or(DecodeError::ArithmeticOverflow)?;
842    Ok(())
843}
844
845pub fn inspect_frame(input: &[u8]) -> Result<HeaderStatus, DecodeError> {
846    if input.len() < 4 {
847        return Ok(HeaderStatus::NeedMore { minimum: 4 });
848    }
849    let magic = read_u32(input);
850    if (SKIPPABLE_MAGIC_MIN..=SKIPPABLE_MAGIC_MAX).contains(&magic) {
851        if input.len() < 8 {
852            return Ok(HeaderStatus::NeedMore { minimum: 8 });
853        }
854        return Ok(HeaderStatus::Complete {
855            header: StreamHeader::Skippable {
856                magic,
857                size: read_u32(&input[4..]),
858            },
859            size: 8,
860        });
861    }
862    if magic != ZSTD_MAGIC {
863        return Err(DecodeError::InvalidMagic);
864    }
865    if input.len() < 5 {
866        return Ok(HeaderStatus::NeedMore { minimum: 5 });
867    }
868    let descriptor = input[4];
869    if descriptor & 0x08 != 0 {
870        return Err(DecodeError::InvalidFrameHeader);
871    }
872    let single_segment = descriptor & 0x20 != 0;
873    let dictionary_size = [0usize, 1, 2, 4][(descriptor & 3) as usize];
874    let content_flag = descriptor >> 6;
875    let content_size_bytes = match content_flag {
876        0 if single_segment => 1,
877        0 => 0,
878        1 => 2,
879        2 => 4,
880        _ => 8,
881    };
882    let size = 5 + usize::from(!single_segment) + dictionary_size + content_size_bytes;
883    if input.len() < size {
884        return Ok(HeaderStatus::NeedMore { minimum: size });
885    }
886    let mut position = 5usize;
887    let encoded_window = if single_segment {
888        None
889    } else {
890        let value = input[position];
891        position += 1;
892        let exponent = (value >> 3) as u32;
893        let base = 1u64 << (10 + exponent);
894        Some(base + (base / 8) * (value as u64 & 7))
895    };
896    let dictionary_id = read_variable(&input[position..position + dictionary_size]);
897    position += dictionary_size;
898    let content_size = if content_size_bytes == 0 {
899        None
900    } else {
901        let mut value = read_variable_u64(&input[position..position + content_size_bytes]);
902        if content_size_bytes == 2 {
903            value += 256;
904        }
905        Some(value)
906    };
907    let window_size = if single_segment {
908        content_size.ok_or(DecodeError::InvalidFrameHeader)?
909    } else {
910        encoded_window.ok_or(DecodeError::InvalidFrameHeader)?
911    };
912    Ok(HeaderStatus::Complete {
913        header: StreamHeader::Zstandard(FrameHeader {
914            window_size,
915            content_size,
916            dictionary_id,
917            has_checksum: descriptor & 4 != 0,
918        }),
919        size,
920    })
921}
922
923fn parse_block_header(bytes: [u8; 3], limit: usize) -> Result<BlockHeader, DecodeError> {
924    let value = bytes[0] as u32 | (bytes[1] as u32) << 8 | (bytes[2] as u32) << 16;
925    let kind = match (value >> 1) & 3 {
926        0 => BlockKind::Raw,
927        1 => BlockKind::Rle,
928        2 => BlockKind::Compressed,
929        _ => return Err(DecodeError::InvalidBlock),
930    };
931    let size = (value >> 3) as usize;
932    if size > limit {
933        return Err(DecodeError::InvalidBlock);
934    }
935    Ok(BlockHeader {
936        last: value & 1 != 0,
937        kind,
938        size,
939    })
940}
941
942fn decode_literals(
943    input: &[u8],
944    output: &mut [u8],
945    table: &mut HuffmanTable<'_>,
946    scratch: &mut [i16],
947    block_limit: usize,
948    strict: bool,
949) -> Result<(usize, usize), DecodeError> {
950    let first = *input.first().ok_or(DecodeError::InvalidBlock)?;
951    let kind = first & 3;
952    let format = (first >> 2) & 3;
953    if kind <= 1 {
954        let (regenerated, header_size): (usize, usize) = match format {
955            0 | 2 => ((first >> 3) as usize, 1),
956            1 => {
957                let second = *input.get(1).ok_or(DecodeError::InvalidBlock)?;
958                (((first >> 4) as usize) | ((second as usize) << 4), 2)
959            }
960            _ => {
961                let second = *input.get(1).ok_or(DecodeError::InvalidBlock)?;
962                let third = *input.get(2).ok_or(DecodeError::InvalidBlock)?;
963                (
964                    ((first >> 4) as usize) | ((second as usize) << 4) | ((third as usize) << 12),
965                    3,
966                )
967            }
968        };
969        ensure_literal_space(regenerated, output.len(), block_limit)?;
970        return if kind == 0 {
971            let end = header_size
972                .checked_add(regenerated)
973                .ok_or(DecodeError::ArithmeticOverflow)?;
974            let source = input
975                .get(header_size..end)
976                .ok_or(DecodeError::InvalidBlock)?;
977            output[..regenerated].copy_from_slice(source);
978            Ok((regenerated, end))
979        } else {
980            let value = *input.get(header_size).ok_or(DecodeError::InvalidBlock)?;
981            output[..regenerated].fill(value);
982            Ok((regenerated, header_size + 1))
983        };
984    }
985    let (regenerated, compressed, streams, header_size): (usize, usize, usize, usize) = match format
986    {
987        0 | 1 => {
988            if input.len() < 3 {
989                return Err(DecodeError::InvalidBlock);
990            }
991            let combined = input[0] as u32 | (input[1] as u32) << 8 | (input[2] as u32) << 16;
992            (
993                ((combined >> 4) & 0x3ff) as usize,
994                ((combined >> 14) & 0x3ff) as usize,
995                if format == 0 { 1 } else { 4 },
996                3,
997            )
998        }
999        2 => {
1000            if input.len() < 4 {
1001                return Err(DecodeError::InvalidBlock);
1002            }
1003            let combined = read_u32(input);
1004            (
1005                ((combined >> 4) & 0x3fff) as usize,
1006                ((combined >> 18) & 0x3fff) as usize,
1007                4,
1008                4,
1009            )
1010        }
1011        _ => {
1012            if input.len() < 5 {
1013                return Err(DecodeError::InvalidBlock);
1014            }
1015            let combined = read_variable_u64(&input[..5]);
1016            (
1017                ((combined >> 4) & 0x3ffff) as usize,
1018                ((combined >> 22) & 0x3ffff) as usize,
1019                4,
1020                5,
1021            )
1022        }
1023    };
1024    ensure_literal_space(regenerated, output.len(), block_limit)?;
1025    let end = header_size
1026        .checked_add(compressed)
1027        .ok_or(DecodeError::ArithmeticOverflow)?;
1028    let mut encoded = input
1029        .get(header_size..end)
1030        .ok_or(DecodeError::InvalidBlock)?;
1031    if kind == 2 {
1032        let table_size = table.read_description(encoded, scratch)?;
1033        encoded = &encoded[table_size..];
1034    } else if !table.is_valid() {
1035        return Err(DecodeError::InvalidEntropyTable);
1036    }
1037    if streams == 1 {
1038        table.decode(encoded, &mut output[..regenerated], strict)?;
1039    } else {
1040        table.decode_four(encoded, &mut output[..regenerated], strict)?;
1041    }
1042    Ok((regenerated, end))
1043}
1044
1045fn ensure_literal_space(
1046    required: usize,
1047    provided: usize,
1048    block_limit: usize,
1049) -> Result<(), DecodeError> {
1050    if required > block_limit {
1051        return Err(DecodeError::InvalidBlock);
1052    }
1053    if required > provided {
1054        return Err(DecodeError::LiteralScratchTooSmall { required, provided });
1055    }
1056    Ok(())
1057}
1058
1059fn parse_sequence_count(input: &[u8]) -> Result<(usize, usize), DecodeError> {
1060    let first = *input.first().ok_or(DecodeError::InvalidBlock)? as usize;
1061    match first {
1062        0 => Ok((0, 1)),
1063        1..=127 => Ok((first, 1)),
1064        128..=254 => Ok((
1065            ((first - 128) << 8) + *input.get(1).ok_or(DecodeError::InvalidBlock)? as usize,
1066            2,
1067        )),
1068        _ => Ok((
1069            0x7f00
1070                + u16::from_le_bytes([
1071                    *input.get(1).ok_or(DecodeError::InvalidBlock)?,
1072                    *input.get(2).ok_or(DecodeError::InvalidBlock)?,
1073                ]) as usize,
1074            3,
1075        )),
1076    }
1077}
1078
1079#[allow(clippy::too_many_arguments)]
1080fn build_sequence_table(
1081    table: &mut FseTable<'_>,
1082    scratch: &mut [i16],
1083    mode: u8,
1084    input: &[u8],
1085    predefined: &[i16],
1086    predefined_log: u8,
1087    max_symbol: usize,
1088    max_log: u8,
1089) -> Result<usize, DecodeError> {
1090    match mode {
1091        0 => {
1092            table.build(predefined, predefined_log, scratch)?;
1093            Ok(0)
1094        }
1095        1 => {
1096            table.build_rle(*input.first().ok_or(DecodeError::InvalidEntropyTable)?);
1097            Ok(1)
1098        }
1099        2 => table.read_description(input, max_symbol, max_log, scratch),
1100        3 if table.is_valid() => Ok(0),
1101        _ => Err(DecodeError::InvalidEntropyTable),
1102    }
1103}
1104
1105fn resolve_offset(
1106    encoded: u32,
1107    literal_length: u32,
1108    repeated: &mut [u32; 3],
1109) -> Result<u32, DecodeError> {
1110    let value = if encoded > 3 {
1111        let value = encoded - 3;
1112        repeated[2] = repeated[1];
1113        repeated[1] = repeated[0];
1114        repeated[0] = value;
1115        value
1116    } else if literal_length != 0 {
1117        match encoded {
1118            1 => repeated[0],
1119            2 => {
1120                repeated.swap(0, 1);
1121                repeated[0]
1122            }
1123            3 => {
1124                let value = repeated[2];
1125                repeated[2] = repeated[1];
1126                repeated[1] = repeated[0];
1127                repeated[0] = value;
1128                value
1129            }
1130            _ => return Err(DecodeError::InvalidOffset),
1131        }
1132    } else {
1133        match encoded {
1134            1 => {
1135                repeated.swap(0, 1);
1136                repeated[0]
1137            }
1138            2 => {
1139                let value = repeated[2];
1140                repeated[2] = repeated[1];
1141                repeated[1] = repeated[0];
1142                repeated[0] = value;
1143                value
1144            }
1145            3 => {
1146                let value = repeated[0]
1147                    .checked_sub(1)
1148                    .filter(|value| *value != 0)
1149                    .ok_or(DecodeError::InvalidOffset)?;
1150                repeated[2] = repeated[1];
1151                repeated[1] = repeated[0];
1152                repeated[0] = value;
1153                value
1154            }
1155            _ => return Err(DecodeError::InvalidOffset),
1156        }
1157    };
1158    Ok(value)
1159}
1160
1161fn read_u32(input: &[u8]) -> u32 {
1162    u32::from_le_bytes([input[0], input[1], input[2], input[3]])
1163}
1164
1165fn read_variable(input: &[u8]) -> u32 {
1166    read_variable_u64(input) as u32
1167}
1168
1169fn read_variable_u64(input: &[u8]) -> u64 {
1170    let mut value = 0u64;
1171    for (index, byte) in input.iter().enumerate() {
1172        value |= (*byte as u64) << (index * 8);
1173    }
1174    value
1175}
1176
1177const LL_BASE: [u32; 36] = [
1178    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 18, 20, 22, 24, 28, 32, 40, 48, 64,
1179    128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536,
1180];
1181const LL_BITS: [u8; 36] = [
1182    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 3, 3, 4, 6, 7, 8, 9, 10, 11,
1183    12, 13, 14, 15, 16,
1184];
1185const ML_BASE: [u32; 53] = [
1186    3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27,
1187    28, 29, 30, 31, 32, 33, 34, 35, 37, 39, 41, 43, 47, 51, 59, 67, 83, 99, 131, 259, 515, 1027,
1188    2051, 4099, 8195, 16387, 32771, 65539,
1189];
1190const ML_BITS: [u8; 53] = [
1191    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
1192    1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
1193];
1194const LL_DEFAULT: [i16; 36] = [
1195    4, 3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, 1, 1, 1, 1, 1,
1196    -1, -1, -1, -1,
1197];
1198const ML_DEFAULT: [i16; 53] = [
1199    1, 4, 3, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1200    1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1, -1, -1, -1,
1201];
1202const OF_DEFAULT: [i16; 29] = [
1203    1, 1, 1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1, -1,
1204];
zstd_zero/lib.rs

zstd_zero/
lib.rs
