steamkit/util/
lzma_decoder.rs

1//! A small, allocation-free LZMA decoder for raw streams whose uncompressed size
2//! is known up front. The output buffer doubles as the dictionary, so the whole
3//! thing is decoded in one call.
4
5use crate::error::SteamError;
6
7const STATE_COUNT: usize = 12;
8const LITERAL_STATE_COUNT: usize = 7;
9const POS_STATE_BITS_MAX: usize = 4;
10const LITERAL_CODER_SIZE: usize = 0x300;
11
12const END_POS_MODEL_INDEX: u32 = 14;
13const FULL_DISTANCE_COUNT: u32 = 1 << (END_POS_MODEL_INDEX >> 1);
14const ALIGN_BITS: u32 = 4;
15
16const LENGTH_CHOICE: usize = 0;
17const LENGTH_CHOICE2: usize = 1;
18const LENGTH_LOW: usize = 2;
19const LENGTH_MID: usize = LENGTH_LOW + (8 << POS_STATE_BITS_MAX);
20const LENGTH_HIGH: usize = LENGTH_MID + (8 << POS_STATE_BITS_MAX);
21const LENGTH_CODER_SIZE: usize = LENGTH_HIGH + 256;
22
23const IS_MATCH: usize = 0;
24const IS_REP: usize = IS_MATCH + (STATE_COUNT << POS_STATE_BITS_MAX);
25const IS_REP_G0: usize = IS_REP + STATE_COUNT;
26const IS_REP_G1: usize = IS_REP_G0 + STATE_COUNT;
27const IS_REP_G2: usize = IS_REP_G1 + STATE_COUNT;
28const IS_REP0_LONG: usize = IS_REP_G2 + STATE_COUNT;
29const POS_SLOT: usize = IS_REP0_LONG + (STATE_COUNT << POS_STATE_BITS_MAX);
30const SPEC_POS: usize = POS_SLOT + (4 << 6);
31const ALIGN: usize = SPEC_POS + FULL_DISTANCE_COUNT as usize - END_POS_MODEL_INDEX as usize;
32const MATCH_LENGTH: usize = ALIGN + (1 << ALIGN_BITS) as usize;
33const REP_LENGTH: usize = MATCH_LENGTH + LENGTH_CODER_SIZE;
34const LITERAL: usize = REP_LENGTH + LENGTH_CODER_SIZE;
35
36/// Decodes raw LZMA streams whose uncompressed size is known up front.
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
38pub struct LzmaDecoder;
39
40impl LzmaDecoder {
41    /// Size of the properties that precede an LZMA stream: one byte for
42    /// lc/lp/pb and four for the dictionary size.
43    pub const PROPERTIES_SIZE: usize = 5;
44
45    /// Decodes a raw LZMA stream into `output`, after an optional preset
46    /// dictionary at its start.
47    ///
48    /// * `properties` - the 5 property bytes.
49    /// * `input` - the compressed stream.
50    /// * `output` - starts with the preset dictionary; the rest receives the
51    ///   uncompressed bytes.
52    /// * `preset_dictionary_length` - how many bytes at the start of `output`
53    ///   matches can refer to before anything is decoded.
54    pub fn decode(
55        properties: &[u8],
56        input: &[u8],
57        output: &mut [u8],
58        preset_dictionary_length: usize,
59    ) -> Result<(), SteamError> {
60        if preset_dictionary_length > output.len() {
61            return Err(SteamError::InvalidArgument(
62                "the preset dictionary does not fit in the output buffer".to_string(),
63            ));
64        }
65
66        if properties.len() < Self::PROPERTIES_SIZE {
67            return Err(SteamError::Decompression(
68                "LZMA properties are too short".to_string(),
69            ));
70        }
71
72        let properties_byte = properties[0];
73
74        if properties_byte >= 9 * 5 * 5 {
75            return Err(SteamError::Decompression(
76                "Invalid LZMA properties".to_string(),
77            ));
78        }
79
80        let lc = (properties_byte % 9) as usize;
81        let lp = (properties_byte / 9 % 5) as usize;
82        let pb = (properties_byte / (9 * 5)) as usize;
83
84        let probability_count = LITERAL + (LITERAL_CODER_SIZE << (lc + lp));
85        let mut probabilities = vec![0u16; probability_count];
86
87        decode_stream(
88            &mut probabilities,
89            input,
90            output,
91            preset_dictionary_length,
92            lc,
93            lp,
94            pb,
95        )
96    }
97}
98
99fn get_state_after_literal(state: usize) -> usize {
100    if state < 4 {
101        return 0;
102    }
103    if state < 10 {
104        return state - 3;
105    }
106    state - 6
107}
108
109fn decode_stream(
110    probabilities: &mut [u16],
111    input: &[u8],
112    output: &mut [u8],
113    position: usize,
114    lc: usize,
115    lp: usize,
116    pb: usize,
117) -> Result<(), SteamError> {
118    for probability in probabilities.iter_mut() {
119        *probability = RangeDecoder::INITIAL_PROBABILITY;
120    }
121
122    let mut range_decoder = RangeDecoder::new(input)?;
123
124    let pos_state_mask = (1usize << pb) - 1;
125    let literal_pos_mask = (1usize << lp) - 1;
126    let literal_shift = 8 - lc;
127
128    let mut state = 0usize;
129    let mut rep0 = 0u32;
130    let mut rep1 = 0u32;
131    let mut rep2 = 0u32;
132    let mut rep3 = 0u32;
133    let mut position = position;
134
135    while position < output.len() {
136        let pos_state = position & pos_state_mask;
137        let state_index = (state << POS_STATE_BITS_MAX) + pos_state;
138
139        if range_decoder.decode_bit(&mut probabilities[IS_MATCH + state_index])? == 0 {
140            let previous_byte = if position > 0 { output[position - 1] } else { 0 };
141            let literal_state = ((position & literal_pos_mask) << lc)
142                + ((previous_byte as u32) >> literal_shift) as usize;
143            let start = LITERAL + LITERAL_CODER_SIZE * literal_state;
144
145            if state < LITERAL_STATE_COUNT {
146                let byte = range_decoder
147                    .decode_literal(&mut probabilities[start..start + LITERAL_CODER_SIZE])?;
148                output[position] = byte;
149            } else {
150                let match_index = position
151                    .checked_sub(rep0 as usize + 1)
152                    .ok_or_else(|| {
153                        SteamError::Decompression(
154                            "LZMA match distance is outside of the decoded data".to_string(),
155                        )
156                    })?;
157                let match_byte = *output.get(match_index).ok_or_else(|| {
158                    SteamError::Decompression(
159                        "LZMA match distance is outside of the decoded data".to_string(),
160                    )
161                })?;
162                let byte = range_decoder.decode_matched_literal(
163                    &mut probabilities[start..start + LITERAL_CODER_SIZE],
164                    match_byte,
165                )?;
166                output[position] = byte;
167            }
168
169            position += 1;
170            state = get_state_after_literal(state);
171            continue;
172        }
173
174        let length;
175
176        if range_decoder.decode_bit(&mut probabilities[IS_REP + state])? == 0 {
177            rep3 = rep2;
178            rep2 = rep1;
179            rep1 = rep0;
180
181            length = range_decoder.decode_length(
182                &mut probabilities[MATCH_LENGTH..MATCH_LENGTH + LENGTH_CODER_SIZE],
183                pos_state,
184            )?;
185            state = if state < LITERAL_STATE_COUNT { 7 } else { 10 };
186            rep0 = range_decoder.decode_distance(&mut *probabilities, length)?;
187
188            if rep0 == u32::MAX {
189                return Err(SteamError::Decompression(
190                    "LZMA stream ended before the expected uncompressed size".to_string(),
191                ));
192            }
193
194            if rep0 as usize >= position {
195                return Err(SteamError::Decompression(
196                    "LZMA match distance is outside of the decoded data".to_string(),
197                ));
198            }
199        } else {
200            if position == 0 {
201                return Err(SteamError::Decompression(
202                    "LZMA stream starts with a repeated match".to_string(),
203                ));
204            }
205
206            if range_decoder.decode_bit(&mut probabilities[IS_REP_G0 + state])? == 0 {
207                if range_decoder.decode_bit(&mut probabilities[IS_REP0_LONG + state_index])? == 0 {
208                    let source = position
209                        .checked_sub(rep0 as usize + 1)
210                        .ok_or_else(|| {
211                            SteamError::Decompression(
212                                "LZMA match distance is outside of the decoded data".to_string(),
213                            )
214                        })?;
215                    let byte = *output.get(source).ok_or_else(|| {
216                        SteamError::Decompression(
217                            "LZMA match distance is outside of the decoded data".to_string(),
218                        )
219                    })?;
220                    output[position] = byte;
221                    position += 1;
222                    state = if state < LITERAL_STATE_COUNT { 9 } else { 11 };
223                    continue;
224                }
225            } else {
226                let distance;
227
228                if range_decoder.decode_bit(&mut probabilities[IS_REP_G1 + state])? == 0 {
229                    distance = rep1;
230                } else {
231                    if range_decoder.decode_bit(&mut probabilities[IS_REP_G2 + state])? == 0 {
232                        distance = rep2;
233                    } else {
234                        distance = rep3;
235                        rep3 = rep2;
236                    }
237
238                    rep2 = rep1;
239                }
240
241                rep1 = rep0;
242                rep0 = distance;
243            }
244
245            length = range_decoder.decode_length(
246                &mut probabilities[REP_LENGTH..REP_LENGTH + LENGTH_CODER_SIZE],
247                pos_state,
248            )?;
249            state = if state < LITERAL_STATE_COUNT { 8 } else { 11 };
250        }
251
252        let match_length = ((length as usize) + 2).min(output.len() - position);
253        let match_distance = rep0 as usize + 1;
254
255        if match_distance >= match_length {
256            let source = position - match_distance;
257            output.copy_within(source..source + match_length, position);
258        } else if match_distance == 1 {
259            let byte = output[position - 1];
260            for slot in output[position..position + match_length].iter_mut() {
261                *slot = byte;
262            }
263        } else {
264            for offset in 0..match_length {
265                output[position + offset] = output[position + offset - match_distance];
266            }
267        }
268
269        position += match_length;
270    }
271
272    Ok(())
273}
274
275struct RangeDecoder<'a> {
276    input: &'a [u8],
277    position: usize,
278    range: u32,
279    code: u32,
280}
281
282impl<'a> RangeDecoder<'a> {
283    const INITIAL_PROBABILITY: u16 = 1 << (RangeDecoder::BIT_MODEL_TOTAL_BITS - 1);
284
285    const INITIAL_BYTES: usize = 5;
286    const BIT_MODEL_TOTAL_BITS: u32 = 11;
287    const MOVE_BITS: u32 = 5;
288    const TOP_VALUE: u32 = 1 << 24;
289    const PROBABILITY_ZERO_OFFSET: u32 = (1u32 << RangeDecoder::MOVE_BITS)
290        .wrapping_sub(1)
291        .wrapping_sub(1u32 << RangeDecoder::BIT_MODEL_TOTAL_BITS);
292
293    fn new(input: &'a [u8]) -> Result<Self, SteamError> {
294        if input.len() < Self::INITIAL_BYTES {
295            return Err(SteamError::Decompression(
296                "LZMA stream is truncated".to_string(),
297            ));
298        }
299
300        Ok(Self {
301            input,
302            position: Self::INITIAL_BYTES,
303            range: u32::MAX,
304            code: u32::from_be_bytes([input[1], input[2], input[3], input[4]]),
305        })
306    }
307
308    fn decode_bit(&mut self, probability: &mut u16) -> Result<u32, SteamError> {
309        let prob = *probability as u32;
310        let bound = (self.range >> Self::BIT_MODEL_TOTAL_BITS) * prob;
311        let bit;
312
313        if self.code < bound {
314            self.range = bound;
315            *probability =
316                (prob + (((1u32 << Self::BIT_MODEL_TOTAL_BITS) - prob) >> Self::MOVE_BITS)) as u16;
317            bit = 0;
318        } else {
319            self.range -= bound;
320            self.code -= bound;
321            *probability = (prob - (prob >> Self::MOVE_BITS)) as u16;
322            bit = 1;
323        }
324
325        if self.range < Self::TOP_VALUE {
326            self.normalize()?;
327        }
328
329        Ok(bit)
330    }
331
332    /// Same as [`Self::decode_bit`], but selects the results with masks
333    /// instead of branching.
334    fn decode_tree_bit(&mut self, probability: &mut u16) -> Result<u32, SteamError> {
335        let prob = *probability as u32;
336        let bound = (self.range >> Self::BIT_MODEL_TOTAL_BITS) * prob;
337
338        let zero_mask = 0u32.wrapping_sub(
339            (((self.code as u64).wrapping_sub(bound as u64)) >> 63) as u32,
340        );
341        let one_mask = !zero_mask;
342
343        self.range = (bound & zero_mask) | (self.range.wrapping_sub(bound) & one_mask);
344        self.code = self.code.wrapping_sub(bound & one_mask);
345
346        let adjusted = prob.wrapping_add(zero_mask & Self::PROBABILITY_ZERO_OFFSET);
347        *probability = prob.wrapping_sub(adjusted >> Self::MOVE_BITS) as u16;
348
349        if self.range < Self::TOP_VALUE {
350            self.normalize()?;
351        }
352
353        Ok(one_mask & 1)
354    }
355
356    fn decode_literal(&mut self, probabilities: &mut [u16]) -> Result<u8, SteamError> {
357        let mut symbol = 1u32;
358
359        while symbol < 0x100 {
360            let bit = self.decode_tree_bit(&mut probabilities[symbol as usize])?;
361            symbol = (symbol << 1) | bit;
362        }
363
364        Ok(symbol as u8)
365    }
366
367    fn decode_matched_literal(
368        &mut self,
369        probabilities: &mut [u16],
370        match_byte: u8,
371    ) -> Result<u8, SteamError> {
372        let mut symbol = 1u32;
373        let mut match_byte = match_byte as u32;
374        let mut offset = 0x100u32;
375
376        while symbol < 0x100 {
377            match_byte <<= 1;
378
379            let previous_offset = offset;
380            offset &= match_byte;
381
382            let index = (previous_offset + offset + symbol) as usize;
383            let bit = self.decode_tree_bit(&mut probabilities[index])?;
384            symbol = (symbol << 1) | bit;
385
386            // Leaves 0x100 only when the decoded bit equals the match bit.
387            // `bit - 1` has to wrap to 0xFFFFFFFF when the bit is zero.
388            offset ^= previous_offset & bit.wrapping_sub(1);
389        }
390
391        Ok(symbol as u8)
392    }
393
394    /// Decodes a match length, minus the minimum match length of 2.
395    fn decode_length(&mut self, probabilities: &mut [u16], pos_state: usize) -> Result<u32, SteamError> {
396        if self.decode_bit(&mut probabilities[LENGTH_CHOICE])? == 0 {
397            return self.decode_bit_tree(&mut probabilities[LENGTH_LOW + (pos_state << 3)..], 3);
398        }
399
400        if self.decode_bit(&mut probabilities[LENGTH_CHOICE2])? == 0 {
401            return Ok(8
402                + self.decode_bit_tree(&mut probabilities[LENGTH_MID + (pos_state << 3)..], 3)?);
403        }
404
405        Ok(16 + self.decode_bit_tree(&mut probabilities[LENGTH_HIGH..], 8)?)
406    }
407
408    /// Decodes a match distance, minus one. Returns `u32::MAX` for the end
409    /// marker.
410    fn decode_distance(&mut self, probabilities: &mut [u16], length: u32) -> Result<u32, SteamError> {
411        let length_state = length.min(3);
412        let pos_slot = self.decode_bit_tree(
413            &mut probabilities[POS_SLOT + ((length_state as usize) << 6)..],
414            6,
415        )?;
416
417        if pos_slot < 4 {
418            return Ok(pos_slot);
419        }
420
421        let direct_bits = (pos_slot >> 1) - 1;
422        let mut distance = (2 | (pos_slot & 1)) << direct_bits;
423
424        if pos_slot < END_POS_MODEL_INDEX {
425            let start = SPEC_POS + distance as usize - pos_slot as usize - 1;
426            return Ok(distance
427                + self.decode_reverse_bit_tree(&mut probabilities[start..], direct_bits as usize)?);
428        }
429
430        distance += self.decode_direct_bits((direct_bits - ALIGN_BITS) as usize)? << ALIGN_BITS;
431        Ok(distance
432            + self.decode_reverse_bit_tree(&mut probabilities[ALIGN..], ALIGN_BITS as usize)?)
433    }
434
435    fn decode_bit_tree(&mut self, probabilities: &mut [u16], bit_count: usize) -> Result<u32, SteamError> {
436        let mut symbol = 1u32;
437
438        for _ in 0..bit_count {
439            let bit = self.decode_tree_bit(&mut probabilities[symbol as usize])?;
440            symbol = (symbol << 1) | bit;
441        }
442
443        Ok(symbol - (1u32 << bit_count))
444    }
445
446    fn decode_reverse_bit_tree(
447        &mut self,
448        probabilities: &mut [u16],
449        bit_count: usize,
450    ) -> Result<u32, SteamError> {
451        let mut index = 1u32;
452        let mut symbol = 0u32;
453
454        for i in 0..bit_count {
455            let bit = self.decode_tree_bit(&mut probabilities[index as usize])?;
456            index = (index << 1) | bit;
457            symbol |= bit << i;
458        }
459
460        Ok(symbol)
461    }
462
463    fn decode_direct_bits(&mut self, bit_count: usize) -> Result<u32, SteamError> {
464        let mut result = 0u32;
465
466        for _ in 0..bit_count {
467            self.range >>= 1;
468
469            let mask = 0u32.wrapping_sub((self.code.wrapping_sub(self.range)) >> 31);
470            self.code = self.code.wrapping_sub(self.range & !mask);
471            result = (result << 1) | (!mask & 1);
472
473            if self.range < Self::TOP_VALUE {
474                self.normalize()?;
475            }
476        }
477
478        Ok(result)
479    }
480
481    fn normalize(&mut self) -> Result<(), SteamError> {
482        if self.position >= self.input.len() {
483            return Err(SteamError::Decompression(
484                "LZMA stream is truncated".to_string(),
485            ));
486        }
487
488        self.range <<= 8;
489        self.code = (self.code << 8) | self.input[self.position] as u32;
490        self.position += 1;
491        Ok(())
492    }
493}
steamkit/util/lzma_decoder.rs

steamkit/util/
lzma_decoder.rs
