Skip to main content

zstd_rs/decode/
mod.rs

1//! Decompression: frames, blocks, and the public `Decompressor`.
2mod execute;
3mod literals;
4mod seqtable;
5mod sequences;
6mod state;
7pub mod tables;
8
9use crate::dict::Dictionary;
10use crate::error::{Error, Result};
11use crate::frame::{parse_block_header, BlockType, FrameFormat, FrameHeader, BLOCK_MAX, SKIPPABLE_BASE, SKIPPABLE_MASK, WINDOW_LOG_MAX};
12use crate::huf::dtable::HufDTable;
13use crate::xxh64::xxh64;
14use alloc::vec::Vec;
15use execute::Exec;
16use seqtable::{Kind, SeqTable};
17use state::{DecState, Predefined};
18
19/// A dictionary prepared for decoding (tables built once, shared by any number of frames).
20pub struct DecoderDictionary {
21    pub(crate) dict: Dictionary,
22    pub(crate) huf: Option<HufDTable>,
23    pub(crate) seq: Option<[SeqTable; 3]>,
24}
25
26impl DecoderDictionary {
27    /// Builds the decoding tables for `dict` once.
28    pub fn new(dict: Dictionary) -> Result<DecoderDictionary> {
29        let (huf, seq) = match &dict.tables {
30            Some(t) => {
31                let mut s = [SeqTable::default(), SeqTable::default(), SeqTable::default()];
32                s[0].build(&t.ll, Kind::LiteralLength)?;
33                s[1].build(&t.of, Kind::Offset)?;
34                s[2].build(&t.ml, Kind::MatchLength)?;
35                (Some(HufDTable::new(&t.huf)), Some(s))
36            }
37            None => (None, None),
38        };
39        Ok(DecoderDictionary { dict, huf, seq })
40    }
41
42    /// The underlying dictionary.
43    pub fn dictionary(&self) -> &Dictionary {
44        &self.dict
45    }
46}
47
48/// Reusable decompression context. Holds no configuration: every call states its limit.
49pub struct Decompressor {
50    st: DecState,
51    pre: Predefined,
52}
53
54impl Decompressor {
55    /// Creates a decompression context (its predefined tables are built here).
56    pub fn new() -> Decompressor {
57        let pre = match Predefined::new() {
58            Ok(p) => p,
59            Err(_) => unreachable!("predefined distributions are valid constants"),
60        };
61        Decompressor { st: DecState::new(), pre }
62    }
63
64    /// Decodes every frame in `src` (skipping skippable frames), appending to `out`.
65    /// Fails with `OutputLimit` before `out` would grow past `max_output` bytes in total.
66    /// Returns the number of bytes appended. Like the reference, empty input is zero frames.
67    pub fn decompress(&mut self, src: &[u8], dict: Option<&DecoderDictionary>, max_output: usize, out: &mut Vec<u8>) -> Result<usize> {
68        self.decompress_format(FrameFormat::Standard, src, dict, max_output, out)
69    }
70
71    /// Like [`Decompressor::decompress`] for frames in `format` (magicless frames
72    /// are concatenated without magic numbers and cannot include skippable frames).
73    pub fn decompress_format(
74        &mut self,
75        format: FrameFormat,
76        src: &[u8],
77        dict: Option<&DecoderDictionary>,
78        max_output: usize,
79        out: &mut Vec<u8>,
80    ) -> Result<usize> {
81        let before = out.len();
82        let limit = before.checked_add(max_output).ok_or(Error::Parameter("output limit overflows"))?;
83        let mut at = 0;
84        while at < src.len() {
85            at += self.frame(format, &src[at..], dict, limit, out)?;
86        }
87        Ok(out.len() - before)
88    }
89
90    /// Decodes one frame (or skips one skippable frame) at the start of `src`,
91    /// appending to `out` without letting it exceed `limit` bytes. Returns bytes consumed.
92    pub fn decompress_frame(&mut self, src: &[u8], dict: Option<&DecoderDictionary>, limit: usize, out: &mut Vec<u8>) -> Result<usize> {
93        self.frame(FrameFormat::Standard, src, dict, limit, out)
94    }
95
96    fn frame(
97        &mut self,
98        format: FrameFormat,
99        src: &[u8],
100        dict: Option<&DecoderDictionary>,
101        limit: usize,
102        out: &mut Vec<u8>,
103    ) -> Result<usize> {
104        if format == FrameFormat::Magicless {
105            let h = FrameHeader::parse_format(src, format)?;
106            return self.frame_with(src, &h, dict, limit, out);
107        }
108        let magic = u32::from_le_bytes(src.get(..4).ok_or(Error::Truncated)?.try_into().map_err(|_| Error::Truncated)?);
109        if magic & SKIPPABLE_MASK == SKIPPABLE_BASE {
110            let len = u32::from_le_bytes(src.get(4..8).ok_or(Error::Truncated)?.try_into().map_err(|_| Error::Truncated)?);
111            let end = 8usize.checked_add(len as usize).ok_or(Error::Truncated)?;
112            return if end <= src.len() { Ok(end) } else { Err(Error::Truncated) };
113        }
114        let h = FrameHeader::parse(src)?;
115        self.frame_with(src, &h, dict, limit, out)
116    }
117
118    /// Decodes the frame whose header `h` was parsed from the start of `src`.
119    fn frame_with(
120        &mut self,
121        src: &[u8],
122        h: &FrameHeader,
123        dict: Option<&DecoderDictionary>,
124        limit: usize,
125        out: &mut Vec<u8>,
126    ) -> Result<usize> {
127        if h.window_size > 1u64 << WINDOW_LOG_MAX {
128            return Err(Error::WindowTooLarge);
129        }
130        let dict = match (h.dict_id, dict) {
131            (0, d) => d,
132            (id, Some(d)) if d.dict.id == id => Some(d),
133            (_, Some(_)) => return Err(Error::WrongDictionary),
134            (_, None) => return Err(Error::MissingDictionary),
135        };
136        if let Some(cs) = h.content_size {
137            let room = limit.saturating_sub(out.len()) as u64;
138            if cs > room {
139                return Err(Error::OutputLimit);
140            }
141            out.reserve(cs as usize);
142        }
143        self.st.reset(dict);
144        let frame_start = out.len();
145        if let Some(cs) = h.content_size {
146            out.resize(frame_start + cs as usize + execute::SLACK, 0);
147        }
148        let mut pos = frame_start;
149        let body = self.frame_blocks(src, h, dict, limit, out, &mut pos);
150        out.truncate(pos);
151        let mut at = body?;
152        let produced = (pos - frame_start) as u64;
153        if h.content_size.is_some_and(|cs| cs != produced) {
154            return Err(Error::ContentSize);
155        }
156        if h.checksum {
157            let want = src.get(at..at + 4).ok_or(Error::Truncated)?;
158            let got = (xxh64(&out[frame_start..], 0) as u32).to_le_bytes();
159            if want != got {
160                return Err(Error::Checksum);
161            }
162            at += 4;
163        }
164        Ok(at)
165    }
166
167    /// Decodes the blocks of a frame into `out` at `*pos` (advanced as output is
168    /// produced, also on error); returns the offset just past the last block.
169    fn frame_blocks(
170        &mut self,
171        src: &[u8],
172        h: &FrameHeader,
173        dict: Option<&DecoderDictionary>,
174        limit: usize,
175        out: &mut Vec<u8>,
176        pos: &mut usize,
177    ) -> Result<usize> {
178        let frame_start = *pos;
179        let block_max = (h.window_size.min(BLOCK_MAX as u64)) as usize;
180        let mut at = h.header_len;
181        loop {
182            let (last, ty, size) = parse_block_header(src.get(at..).ok_or(Error::Truncated)?)?;
183            at += 3;
184            if size > block_max {
185                return Err(Error::Corrupt("block larger than the block size limit"));
186            }
187            match ty {
188                BlockType::Rle => {
189                    let b = *src.get(at).ok_or(Error::Truncated)?;
190                    make_room(out, *pos, size, limit)?;
191                    out[*pos..*pos + size].fill(b);
192                    *pos += size;
193                    at += 1;
194                }
195                BlockType::Raw => {
196                    let body = src.get(at..at + size).ok_or(Error::Truncated)?;
197                    make_room(out, *pos, size, limit)?;
198                    out[*pos..*pos + size].copy_from_slice(body);
199                    *pos += size;
200                    at += size;
201                }
202                BlockType::Compressed => {
203                    let body = src.get(at..at + size).ok_or(Error::Truncated)?;
204                    *pos = self.block(body, dict, &Bounds { frame_start, limit, block_max }, *pos, out)?;
205                    at += size;
206                }
207            }
208            if last {
209                return Ok(at);
210            }
211        }
212    }
213
214    /// Decodes one compressed block at `pos`; returns the new logical end.
215    fn block(&mut self, body: &[u8], dict: Option<&DecoderDictionary>, b: &Bounds, pos: usize, out: &mut Vec<u8>) -> Result<usize> {
216        let (lits, used) = literals::decode(body, &mut self.st, dict, b.block_max)?;
217        let history: &[u8] = dict.map(|d| d.dict.content.as_slice()).unwrap_or(&[]);
218        let mut ex = Exec::new(out, pos, b.frame_start, history, b.limit, b.block_max);
219        let r = sequences::decode(body, used, lits, &mut self.st, &self.pre, dict, &mut ex);
220        let end = ex.pos;
221        r.map(|_| end)
222    }
223}
224
225impl Default for Decompressor {
226    fn default() -> Self {
227        Self::new()
228    }
229}
230
231/// Checks the output limit and makes `out[pos..pos + n]` writable.
232fn make_room(out: &mut Vec<u8>, pos: usize, n: usize, limit: usize) -> Result<()> {
233    let end = pos.checked_add(n).filter(|&e| e <= limit).ok_or(Error::OutputLimit)?;
234    if out.len() < end {
235        out.resize(end, 0);
236    }
237    Ok(())
238}
239
240/// Per-frame bounds every block is checked against.
241struct Bounds {
242    /// Index in the output where the frame starts.
243    frame_start: usize,
244    /// Caller's cap on the output length.
245    limit: usize,
246    /// Block_Maximum_Size: min(window, 128 KiB).
247    block_max: usize,
248}