Skip to main content

zstd_complete/decoding/
dictionary.rs

1use alloc::vec::Vec;
2use core::convert::TryInto;
3
4use crate::decoding::errors::DictionaryDecodeError;
5use crate::decoding::scratch::FSEScratch;
6use crate::decoding::scratch::HuffmanScratch;
7
8/// Zstandard includes support for "raw content" dictionaries, that store bytes optionally used
9/// during sequence execution.
10///
11/// <https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#dictionary-format>
12pub struct Dictionary {
13    /// A 4 byte value used by decoders to check if they can use
14    /// the correct dictionary. This value must not be zero.
15    pub id: u32,
16    /// A dictionary can contain an entropy table, either FSE or
17    /// Huffman.
18    pub fse: FSEScratch,
19    /// A dictionary can contain an entropy table, either FSE or
20    /// Huffman.
21    pub huf: HuffmanScratch,
22    /// The content of a dictionary acts as a "past" in front of data
23    /// to compress or decompress,
24    /// so it can be referenced in sequence commands.
25    /// As long as the amount of data decoded from this frame is less than or
26    /// equal to Window_Size, sequence commands may specify offsets longer than
27    /// the total length of decoded output so far to reference back to the
28    /// dictionary, even parts of the dictionary with offsets larger than Window_Size.
29    /// After the total output has surpassed Window_Size however,
30    /// this is no longer allowed and the dictionary is no longer accessible
31    pub dict_content: Vec<u8>,
32    /// The 3 most recent offsets are stored so that they can be used
33    /// during sequence execution, see
34    /// <https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#repeat-offsets>
35    /// for more.
36    pub offset_hist: [u32; 3],
37}
38
39/// This 4 byte (little endian) magic number refers to the start of a dictionary
40pub const MAGIC_NUM: [u8; 4] = [0x37, 0xA4, 0x30, 0xEC];
41
42impl Dictionary {
43    /// Parses the dictionary from `raw` and set the tables
44    /// it returns the dict_id for checking with the frame's `dict_id``
45    pub fn decode_dict(raw: &[u8]) -> Result<Dictionary, DictionaryDecodeError> {
46        if raw.len() < 8 {
47            return Err(DictionaryDecodeError::NotEnoughBytes);
48        }
49
50        let mut new_dict = Dictionary {
51            id: 0,
52            fse: FSEScratch::new(),
53            huf: HuffmanScratch::new(),
54            dict_content: Vec::new(),
55            offset_hist: [2, 4, 8],
56        };
57
58        let magic_num: [u8; 4] = raw[..4].try_into().expect("optimized away");
59        if magic_num != MAGIC_NUM {
60            return Err(DictionaryDecodeError::BadMagicNum { got: magic_num });
61        }
62
63        let dict_id = raw[4..8].try_into().expect("optimized away");
64        let dict_id = u32::from_le_bytes(dict_id);
65        new_dict.id = dict_id;
66
67        let raw_tables = &raw[8..];
68
69        let huf_size = new_dict.huf.table.build_decoder(raw_tables)?;
70
71        if raw_tables.len() < huf_size as usize {
72            return Err(DictionaryDecodeError::NotEnoughBytes);
73        }
74        let raw_tables = &raw_tables[huf_size as usize..];
75
76        let of_size = new_dict.fse.offsets.build_decoder(
77            raw_tables,
78            crate::decoding::sequence_section_decoder::OF_MAX_LOG,
79        )?;
80
81        if raw_tables.len() < of_size {
82            return Err(DictionaryDecodeError::NotEnoughBytes);
83        }
84        let raw_tables = &raw_tables[of_size..];
85
86        let ml_size = new_dict.fse.match_lengths.build_decoder(
87            raw_tables,
88            crate::decoding::sequence_section_decoder::ML_MAX_LOG,
89        )?;
90
91        if raw_tables.len() < ml_size {
92            return Err(DictionaryDecodeError::NotEnoughBytes);
93        }
94        let raw_tables = &raw_tables[ml_size..];
95
96        let ll_size = new_dict.fse.literal_lengths.build_decoder(
97            raw_tables,
98            crate::decoding::sequence_section_decoder::LL_MAX_LOG,
99        )?;
100
101        if raw_tables.len() < ll_size {
102            return Err(DictionaryDecodeError::NotEnoughBytes);
103        }
104        let raw_tables = &raw_tables[ll_size..];
105
106        if raw_tables.len() < 12 {
107            return Err(DictionaryDecodeError::NotEnoughBytes);
108        }
109        let offset1 = raw_tables[0..4].try_into().expect("optimized away");
110        let offset1 = u32::from_le_bytes(offset1);
111
112        let offset2 = raw_tables[4..8].try_into().expect("optimized away");
113        let offset2 = u32::from_le_bytes(offset2);
114
115        let offset3 = raw_tables[8..12].try_into().expect("optimized away");
116        let offset3 = u32::from_le_bytes(offset3);
117
118        new_dict.offset_hist[0] = offset1;
119        new_dict.offset_hist[1] = offset2;
120        new_dict.offset_hist[2] = offset3;
121
122        let raw_content = &raw_tables[12..];
123        new_dict.dict_content.extend(raw_content);
124
125        Ok(new_dict)
126    }
127}
128
129#[test]
130fn truncated_dictionary() {
131    use alloc::vec;
132
133    // First case: Valid dictionary magic number, but missing the 4-byte dictionary ID.
134    let raw = [0x37, 0xA4, 0x30, 0xEC];
135    let _ = Dictionary::decode_dict(&raw);
136
137    // Second case: Valid dictionary magic number, non-zero dictionary ID, table bytes known to parse successfully. But fewer than 12 bytes remain for the 3 offset-history u32 values..
138    let mut raw = vec![0u8; 8];
139    raw[0] = 0x37;
140    raw[1] = 0xA4;
141    raw[2] = 0x30;
142    raw[3] = 0xEC;
143    raw[4] = 0x01;
144    raw[5] = 0x21;
145    raw[6] = 0x23;
146    raw[7] = 0x47;
147
148    let raw_tables = [
149        54, 16, 192, 155, 4, 0, 207, 59, 239, 121, 158, 116, 220, 93, 114, 229, 110, 41, 249, 95,
150        165, 255, 83, 202, 254, 68, 74, 159, 63, 161, 100, 151, 137, 21, 184, 183, 189, 100, 235,
151        209, 251, 174, 91, 75, 91, 185, 19, 39, 75, 146, 98, 177, 249, 14, 4, 35, 0, 0, 0, 40, 40,
152        20, 10, 12, 204, 37, 196, 1, 173, 122, 0, 4, 0, 128, 1, 2, 2, 25, 32, 27, 27, 22, 24, 26,
153        18, 12, 12, 15, 16, 11, 69, 37, 225, 48, 20, 12, 6, 2, 161, 80, 40, 20, 44, 137, 145, 204,
154        46, 0, 0, 0, 0, 0, 116, 253, 16, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
155    ];
156
157    raw.extend_from_slice(&raw_tables);
158
159    // Fewer than 12 bytes remain for the 3 offset-history u32 values.
160    raw.extend_from_slice(&[3, 0, 0, 0, 10, 0, 0, 0, 0xEF, 0xCD, 0xAB]);
161
162    let _ = Dictionary::decode_dict(&raw);
163}