autoit/decompress/
mod.rs

1//! AutoIt compressed payload decompression.
2//!
3//! Dispatches on the 8-byte wrapper header (4-byte magic plus a big-endian
4//! decompressed-size field) to the matching codec: the `EA04`/`EA05`/`EA06` LZ
5//! scheme handled here, or the `JB00`/`JB01` adaptive-Huffman + LZSS scheme in
6//! the `jb` submodule.
7
8mod bitstream;
9mod jb;
10
11use crate::{Encoding, Error};
12use bitstream::BitStream;
13
14const HEADER_LEN: usize = 8;
15const MAGIC_EA04: &[u8; 4] = b"EA04";
16const MAGIC_EA05: &[u8; 4] = b"EA05";
17const MAGIC_EA06: &[u8; 4] = b"EA06";
18const MAGIC_JB00: &[u8; 4] = b"JB00";
19const MAGIC_JB01: &[u8; 4] = b"JB01";
20
21/// Decompression limits.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct Limits {
24    /// Maximum decompressed output size in bytes.
25    pub max_output_size: usize,
26}
27
28impl Default for Limits {
29    /// Returns limits with a 64 MiB maximum decompressed output size.
30    ///
31    /// # Returns
32    ///
33    /// A [`Limits`] whose `max_output_size` is 64 MiB.
34    fn default() -> Self {
35        Self {
36            max_output_size: 64 * 1024 * 1024,
37        }
38    }
39}
40
41/// Decompresses an AutoIt compressed blob.
42///
43/// Parses the 8-byte header (4-byte magic and a big-endian decompressed size)
44/// and dispatches to the LZ codec for `EA04`/`EA05`/`EA06` or the JB codec for
45/// `JB00`/`JB01`.
46///
47/// # Arguments
48///
49/// * `data` - The full compressed blob including the 8-byte wrapper header.
50/// * `limits` - Caps applied during decompression, notably the maximum output size.
51///
52/// # Returns
53///
54/// The fully decompressed bytes on success.
55///
56/// # Errors
57///
58/// Returns [`Error::truncated`] if the header or its size field is missing or
59/// the payload slice is absent. Returns [`Error::limit_exceeded`] if the
60/// advertised size does not fit in a `usize` or exceeds `limits.max_output_size`.
61/// Returns [`Error::compression_error`] for an unrecognized magic. Errors from
62/// the selected codec are propagated.
63pub fn decompress(data: &[u8], limits: Limits) -> Result<Vec<u8>, Error> {
64    let header = data.get(0..HEADER_LEN).ok_or_else(Error::truncated)?;
65    let magic = header.get(0..4).ok_or_else(Error::truncated)?;
66    let size_bytes: [u8; 4] = header
67        .get(4..8)
68        .ok_or_else(Error::truncated)?
69        .try_into()
70        .map_err(|_err| Error::truncated())?;
71    let output_size =
72        usize::try_from(u32::from_be_bytes(size_bytes)).map_err(|_err| Error::limit_exceeded())?;
73    if output_size > limits.max_output_size {
74        return Err(Error::limit_exceeded());
75    }
76    let payload = data.get(HEADER_LEN..).ok_or_else(Error::truncated)?;
77
78    if magic == MAGIC_EA04 {
79        decompress_lz(payload, output_size, Encoding::Ea04)
80    } else if magic == MAGIC_EA05 {
81        decompress_lz(payload, output_size, Encoding::Ea05)
82    } else if magic == MAGIC_EA06 {
83        decompress_lz(payload, output_size, Encoding::Ea06)
84    } else if magic == MAGIC_JB00 || magic == MAGIC_JB01 {
85        jb::decompress_jb01(payload, output_size)
86    } else {
87        Err(Error::compression_error())
88    }
89}
90
91/// Decompresses an `EA04`/`EA05`/`EA06` LZ payload into `output_size` bytes.
92///
93/// Each token starts with a single control bit: when it equals the encoding's
94/// literal symbol the next 8 bits are an output byte, otherwise a 15-bit
95/// back-reference offset followed by a staged match length (see
96/// [`read_match_len`]) selects bytes to copy from earlier output. The literal
97/// control bit is `0` for `EA04`/`EA05` and `1` for `EA06`.
98///
99/// # Arguments
100///
101/// * `data` - The compressed payload after the 8-byte wrapper header.
102/// * `output_size` - The exact decompressed length to produce.
103/// * `encoding` - Which EA variant the payload uses, selecting the literal control bit.
104///
105/// # Returns
106///
107/// The decompressed bytes of length `output_size`.
108///
109/// # Errors
110///
111/// Returns [`Error::unsupported_encoding`] if `encoding` is [`Encoding::Jb01`].
112/// Returns [`Error::compression_error`] if a decoded byte, offset, or match
113/// length cannot be represented or a match is invalid. Propagates
114/// [`Error::truncated`] from the bit stream when input is exhausted.
115fn decompress_lz(data: &[u8], output_size: usize, encoding: Encoding) -> Result<Vec<u8>, Error> {
116    let literal_symbol = match encoding {
117        Encoding::Ea04 | Encoding::Ea05 => 0,
118        Encoding::Ea06 => 1,
119        Encoding::Jb01 => return Err(Error::unsupported_encoding()),
120    };
121    let mut bits = BitStream::new(data);
122    let mut output = Vec::with_capacity(output_size);
123    while output.len() < output_size {
124        let control = bits.read_bits(1)?;
125        if control == literal_symbol {
126            let byte =
127                u8::try_from(bits.read_bits(8)?).map_err(|_err| Error::compression_error())?;
128            output.push(byte);
129        } else {
130            let offset =
131                usize::try_from(bits.read_bits(15)?).map_err(|_err| Error::compression_error())?;
132            let match_len = read_match_len(&mut bits)?;
133            copy_match(&mut output, offset, match_len, output_size)?;
134        }
135    }
136    Ok(output)
137}
138
139/// Reads a staged, variable-width match length from the bit stream.
140///
141/// The length is encoded in successive stages, each with a base value and a
142/// fixed-width field: if the field's value is below the stage's "more" sentinel
143/// the length is `base + value`; if it equals the sentinel decoding advances to
144/// the next stage. The final (base 299) stage repeats, accumulating the
145/// 8-bit sentinel each iteration until a sub-sentinel field terminates it.
146///
147/// # Arguments
148///
149/// * `bits` - The bit stream positioned at the start of the length code.
150///
151/// # Returns
152///
153/// The decoded match length.
154///
155/// # Errors
156///
157/// Returns [`Error::compression_error`] if an accumulated length overflows or a
158/// field cannot be converted. Propagates [`Error::truncated`] when the bit
159/// stream is exhausted.
160fn read_match_len(bits: &mut BitStream<'_>) -> Result<usize, Error> {
161    let stages = [
162        (3usize, 2usize, 0b11u32),
163        (6, 3, 0b111),
164        (13, 5, 0b1_1111),
165        (44, 8, 0xff),
166        (299, 8, 0xff),
167    ];
168    for (base, width, more) in stages {
169        let add = bits.read_bits(width)?;
170        if add != more {
171            return usize::try_from(add)
172                .ok()
173                .and_then(|value| base.checked_add(value))
174                .ok_or_else(Error::compression_error);
175        }
176        if base == 299 {
177            let mut length = base;
178            loop {
179                length = length
180                    .checked_add(usize::try_from(more).map_err(|_err| Error::compression_error())?)
181                    .ok_or_else(Error::compression_error)?;
182                let extra = bits.read_bits(width)?;
183                if extra != more {
184                    return length
185                        .checked_add(
186                            usize::try_from(extra).map_err(|_err| Error::compression_error())?,
187                        )
188                        .ok_or_else(Error::compression_error);
189                }
190            }
191        }
192    }
193    Err(Error::compression_error())
194}
195
196/// Copies a back-reference match onto the end of the output buffer.
197///
198/// Bytes are copied one at a time from `offset` positions behind the current
199/// end, so overlapping matches (offset smaller than the length) repeat the
200/// preceding window as they extend.
201///
202/// # Arguments
203///
204/// * `output` - The decompressed buffer being appended to.
205/// * `offset` - Distance behind the current end to copy from; must be in `1..=output.len()`.
206/// * `match_len` - Number of bytes to copy.
207/// * `output_size` - The final output length, used to reject overruns.
208///
209/// # Returns
210///
211/// The unit value on success, with `match_len` bytes appended to `output`.
212///
213/// # Errors
214///
215/// Returns [`Error::compression_error`] if `offset` is zero or larger than the
216/// current output, if appending would overflow or exceed `output_size`, or if a
217/// source index cannot be reached.
218fn copy_match(
219    output: &mut Vec<u8>,
220    offset: usize,
221    match_len: usize,
222    output_size: usize,
223) -> Result<(), Error> {
224    if offset == 0 || offset > output.len() {
225        return Err(Error::compression_error());
226    }
227    let end = output
228        .len()
229        .checked_add(match_len)
230        .ok_or_else(Error::compression_error)?;
231    if end > output_size {
232        return Err(Error::compression_error());
233    }
234    for _ in 0..match_len {
235        let source = output
236            .len()
237            .checked_sub(offset)
238            .ok_or_else(Error::compression_error)?;
239        let byte = *output.get(source).ok_or_else(Error::compression_error)?;
240        output.push(byte);
241    }
242    Ok(())
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248
249    #[test]
250    fn decompresses_ea06_literal_only() -> Result<(), String> {
251        let mut blob = Vec::from(*MAGIC_EA06);
252        blob.extend_from_slice(&3u32.to_be_bytes());
253        blob.extend_from_slice(
254            pack_bits(&[
255                1, 0, 1, 0, 0, 0, 0, 0, 1, // A
256                1, 0, 1, 0, 0, 0, 0, 1, 0, // B
257                1, 0, 1, 0, 0, 0, 0, 1, 1, // C
258            ])?
259            .as_slice(),
260        );
261
262        let decompressed =
263            decompress(blob.as_slice(), Limits::default()).map_err(|err| err.to_string())?;
264        check_eq(decompressed, b"ABC".to_vec(), "decompressed")
265    }
266
267    #[test]
268    fn decompresses_ea05_literal_only() -> Result<(), String> {
269        let mut blob = Vec::from(*MAGIC_EA05);
270        blob.extend_from_slice(&2u32.to_be_bytes());
271        blob.extend_from_slice(
272            pack_bits(&[
273                0, 0, 1, 0, 0, 1, 0, 0, 0, // H
274                0, 0, 1, 0, 0, 1, 0, 0, 1, // I
275            ])?
276            .as_slice(),
277        );
278
279        let decompressed =
280            decompress(blob.as_slice(), Limits::default()).map_err(|err| err.to_string())?;
281        check_eq(decompressed, b"HI".to_vec(), "decompressed")
282    }
283
284    #[test]
285    fn decompresses_overlapping_match() -> Result<(), String> {
286        let mut blob = Vec::from(*MAGIC_EA06);
287        blob.extend_from_slice(&5u32.to_be_bytes());
288        blob.extend_from_slice(
289            pack_bits(&[
290                1, 0, 1, 0, 0, 0, 0, 0, 1, // A
291                1, 0, 1, 0, 0, 0, 0, 1, 0, // B
292                0, // match
293                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, // offset 2
294                0, 0, // length 3
295            ])?
296            .as_slice(),
297        );
298
299        let decompressed =
300            decompress(blob.as_slice(), Limits::default()).map_err(|err| err.to_string())?;
301        check_eq(decompressed, b"ABABA".to_vec(), "decompressed")
302    }
303
304    #[test]
305    fn rejects_invalid_match_offset() -> Result<(), String> {
306        let mut blob = Vec::from(*MAGIC_EA06);
307        blob.extend_from_slice(&3u32.to_be_bytes());
308        blob.extend_from_slice(
309            pack_bits(&[
310                0, // match
311                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // offset 0
312                0, 0, // length 3
313            ])?
314            .as_slice(),
315        );
316
317        let Err(err) = decompress(blob.as_slice(), Limits::default()) else {
318            return Err("unexpected decompression success".to_string());
319        };
320        check_eq(
321            err.recognition_failure(),
322            Some(crate::RecognitionFailure::CompressionError),
323            "error",
324        )
325    }
326
327    #[test]
328    fn rejects_truncated_bitstream() -> Result<(), String> {
329        let mut blob = Vec::from(*MAGIC_EA06);
330        blob.extend_from_slice(&1u32.to_be_bytes());
331        blob.push(0b1000_0000);
332
333        let Err(err) = decompress(blob.as_slice(), Limits::default()) else {
334            return Err("unexpected decompression success".to_string());
335        };
336        check_eq(
337            err.recognition_failure(),
338            Some(crate::RecognitionFailure::Truncated),
339            "error",
340        )
341    }
342
343    #[test]
344    fn rejects_advertised_size_above_limit() -> Result<(), String> {
345        let mut blob = Vec::from(*MAGIC_EA06);
346        blob.extend_from_slice(&2u32.to_be_bytes());
347
348        let Err(err) = decompress(blob.as_slice(), Limits { max_output_size: 1 }) else {
349            return Err("unexpected decompression success".to_string());
350        };
351        check_eq(
352            err.recognition_failure(),
353            Some(crate::RecognitionFailure::LimitExceeded),
354            "error",
355        )
356    }
357
358    fn pack_bits(bits: &[u8]) -> Result<Vec<u8>, String> {
359        let mut out = Vec::new();
360        let mut cursor = 0usize;
361        while cursor < bits.len() {
362            let mut byte = 0u8;
363            for bit_index in 0..8usize {
364                let source_index = cursor
365                    .checked_add(bit_index)
366                    .ok_or_else(|| "bit offset overflow".to_string())?;
367                let bit = bits
368                    .get(source_index)
369                    .copied()
370                    .map_or(0, core::convert::identity);
371                byte = (byte << 1) | bit;
372            }
373            out.push(byte);
374            cursor = cursor
375                .checked_add(8)
376                .ok_or_else(|| "bit offset overflow".to_string())?;
377        }
378        Ok(out)
379    }
380
381    fn check_eq<T>(actual: T, expected: T, context: &str) -> Result<(), String>
382    where
383        T: core::fmt::Debug + PartialEq,
384    {
385        if actual == expected {
386            Ok(())
387        } else {
388            Err(format!("{context}: got {actual:?}, expected {expected:?}"))
389        }
390    }
391}
autoit/decompress/mod.rs

autoit/decompress/
mod.rs
