arc_mac/
shrinkwrap.rs

1// SPDX-License-Identifier: LGPL-2.1-or-later
2// Derived from Aaru.Compression.Native and XADMaster; see THIRD_PARTY_NOTICES.md.
3//! ShrinkWrap 3 **"SIT2"** decompressor: the StuffIt codec behind Disk Copy 6 / NDIF
4//! chunk type `0xF0` (images made by Aladdin ShrinkWrap 3, creator `Wrap`).
5//!
6//! It is StuffIt method 14 ("Installer") cut down: one block per chunk, 276 literal /
7//! length symbols, 16 distance symbols, no sliding-window wrap. Ported from Aaru's
8//! `stuffit/shrinkwrap.c` (reverse-engineered from ShrinkWrap 3.5.1's `Dcmp` 240), whose
9//! tree reader is XADMaster's `SIT14_ReadTree`.
10//!
11//! Chunk layout: a 16-byte header (`u32` compressed length, `u32` LE uncompressed length
12//! twice, `0xFFFFFFFF`), then the LSB-first bitstream: the literal/length tree, the
13//! distance tree (each byte-aligned after), then the codes.
14
15use anyhow::{bail, Result};
16
17use crate::stuffit_lzh::BitReader;
18
19const LITLEN_SYMBOLS: usize = 276;
20const DIST_SYMBOLS: usize = 16;
21const HEADER_LEN: usize = 16;
22/// Nested code-length trees never go deeper than one level in real streams.
23const MAX_TREE_DEPTH: u32 = 4;
24
25/// Uncompressed length a ShrinkWrap chunk declares in its header, if it has one.
26pub fn declared_len(chunk: &[u8]) -> Option<usize> {
27    let b = chunk.get(4..8)?;
28    Some(u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize)
29}
30
31/// Decompress one ShrinkWrap SIT2 chunk (header included), refusing outputs over `max_out`.
32pub fn decompress(chunk: &[u8], max_out: usize) -> Result<Vec<u8>> {
33    if chunk.len() < HEADER_LEN {
34        bail!("ShrinkWrap: chunk shorter than its 16-byte header");
35    }
36    let want = declared_len(chunk).unwrap_or(0);
37    if want > max_out {
38        bail!("ShrinkWrap: chunk declares {want} bytes, more than the {max_out} it may hold");
39    }
40    let mut len_base = [0u32; 20];
41    let mut len_extra = [0u32; 20];
42    let mut k = 0u32;
43    for i in 0..20 {
44        len_extra[i] = if i >= 2 { (i as u32 - 2) >> 1 } else { 0 };
45        len_base[i] = k + 4;
46        k += 1 << len_extra[i];
47    }
48    let mut dist_base = [0u32; 16];
49    let mut k = 1u32;
50    for (i, base) in dist_base.iter_mut().enumerate() {
51        *base = k;
52        k += 1 << i;
53    }
54
55    let mut br = BitReader::new(&chunk[HEADER_LEN..]);
56    let litlen = read_tree(&mut br, LITLEN_SYMBOLS, 0)?;
57    let dist = read_tree(&mut br, DIST_SYMBOLS, 0)?;
58
59    let mut out = Vec::with_capacity(want);
60    while out.len() < want {
61        let sym = walk(&mut br, litlen.as_slice(), LITLEN_SYMBOLS)?;
62        if sym < 256 {
63            out.push(sym as u8);
64            continue;
65        }
66        let li = sym - 256;
67        if li >= 20 {
68            bail!("ShrinkWrap: length symbol {li} out of range");
69        }
70        let mut len = len_base[li] + br.next_bits(len_extra[li]);
71        let di = walk(&mut br, dist.as_slice(), DIST_SYMBOLS)?;
72        let distance = dist_base[di] as usize + br.next_bits(di as u32) as usize;
73        if distance == 0 || distance > out.len() {
74            bail!(
75                "ShrinkWrap: back-reference {distance} before the start of the output ({})",
76                out.len()
77            );
78        }
79        while len > 0 && out.len() < want {
80            let b = out[out.len() - distance];
81            out.push(b);
82            len -= 1;
83        }
84    }
85    Ok(out)
86}
87
88/// Follow one codeword through a tree built by [`read_tree`]; returns the symbol.
89fn walk(br: &mut BitReader, tree: &[u16], symbols: usize) -> Result<usize> {
90    let leaf = symbols * 2;
91    let mut node = 0usize;
92    for _ in 0..=32 {
93        node = match tree.get(node + br.next_bit() as usize) {
94            Some(&n) if n != 0 => n as usize,
95            _ => bail!("ShrinkWrap: invalid Huffman code in the stream"),
96        };
97        if node >= leaf {
98            return Ok(node - leaf);
99        }
100    }
101    bail!("ShrinkWrap: Huffman code longer than 32 bits")
102}
103
104/// Read a code-length table (method 14's `SIT14_ReadTree`) and build its decode tree.
105fn read_tree(br: &mut BitReader, codesize: usize, depth: u32) -> Result<Vec<u16>> {
106    let zero_flag = br.next_bits(1);
107    let width = br.next_bits(2) + 2;
108    let offset = br.next_bits(3) + 1;
109    let size = 1usize << width;
110    let repeat = size - 1;
111    // With the flag set, the value one below `repeat` stands for an absent symbol.
112    let absent = if zero_flag != 0 {
113        Some(repeat - 1)
114    } else {
115        None
116    };
117    let mut code = vec![0u8; codesize];
118    let mut i = 0usize;
119    if br.next_bits(2) & 1 != 0 {
120        if depth >= MAX_TREE_DEPTH {
121            bail!("ShrinkWrap: code-length trees nested too deeply");
122        }
123        let sub = read_tree(br, size, depth + 1)?;
124        while i < codesize {
125            let l = walk(br, sub.as_slice(), size)?;
126            if Some(l) == absent {
127                code[i] = 0;
128                i += 1;
129            } else if l == repeat {
130                let n = walk(br, sub.as_slice(), size)? + 3;
131                i = repeat_previous(code.as_mut_slice(), i, n)?;
132            } else {
133                code[i] = (l as u32 + offset) as u8;
134                i += 1;
135            }
136        }
137    } else {
138        while i < codesize {
139            let l = br.next_bits(width) as usize;
140            if Some(l) == absent {
141                code[i] = 0;
142                i += 1;
143            } else if l == repeat {
144                let n = br.next_bits(width) as usize + 3;
145                i = repeat_previous(code.as_mut_slice(), i, n)?;
146            } else {
147                code[i] = (l as u32 + offset) as u8;
148                i += 1;
149            }
150        }
151    }
152    let tree = build_tree(code.as_slice())?;
153    br.align_to_byte();
154    Ok(tree)
155}
156
157/// Copy the previous code length `n` times from `i`; returns the new position.
158fn repeat_previous(code: &mut [u8], i: usize, n: usize) -> Result<usize> {
159    if i == 0 {
160        bail!("ShrinkWrap: code-length repeat with no previous length");
161    }
162    let end = (i + n).min(code.len());
163    let prev = code[i - 1];
164    for c in &mut code[i..end] {
165        *c = prev;
166    }
167    Ok(end)
168}
169
170/// Canonical codes in method 14's order (its quicksort is unstable, so it is kept verbatim).
171fn build_tree(code: &[u8]) -> Result<Vec<u16>> {
172    let n = code.len();
173    if code.iter().any(|&c| c > 32) {
174        bail!("ShrinkWrap: code length over 32 bits");
175    }
176    let mut sorted = code.to_vec();
177    let mut order: Vec<u16> = (0..n as u16).collect();
178    sort_codes(0, n, sorted.as_mut_slice(), order.as_mut_slice());
179
180    let mut bits = vec![0u64; n];
181    let mut i = sorted.iter().position(|&c| c != 0).unwrap_or(n);
182    let mut j = 0u64;
183    while i < n {
184        if i > 0 {
185            j <<= sorted[i] - sorted[i - 1];
186        }
187        let mut m = 0u64;
188        let mut l = j;
189        for _ in 0..sorted[i] {
190            m = (m << 1) | (l & 1);
191            l >>= 1;
192        }
193        bits[order[i] as usize] = m;
194        i += 1;
195        j += 1;
196    }
197
198    let mut tree = vec![0u16; n * 2];
199    let mut next = 2usize;
200    for (sym, &len) in code.iter().enumerate() {
201        let mut node = 0usize;
202        let mut m = bits[sym];
203        for k in 0..len {
204            node += (m & 1) as usize;
205            let slot = match tree.get_mut(node) {
206                Some(s) => s,
207                None => bail!("ShrinkWrap: over-subscribed Huffman code"),
208            };
209            if k + 1 == len {
210                *slot = (n * 2 + sym) as u16;
211            } else {
212                if *slot == 0 {
213                    if next >= n * 2 {
214                        bail!("ShrinkWrap: Huffman code needs more nodes than it has symbols");
215                    }
216                    *slot = next as u16;
217                    next += 2;
218                }
219                node = *slot as usize;
220            }
221            m >>= 1;
222        }
223    }
224    Ok(tree)
225}
226
227/// `SIT14_Update`: the in-place quicksort over `[first, last)` that sets the code order.
228fn sort_codes(mut first: usize, mut last: usize, code: &mut [u8], order: &mut [u16]) {
229    while last - first > 1 {
230        let mut i = first;
231        let mut j = last;
232        loop {
233            i += 1;
234            while i < last && code[first] > code[i] {
235                i += 1;
236            }
237            j -= 1;
238            while j > first && code[first] < code[j] {
239                j -= 1;
240            }
241            if j > i {
242                code.swap(i, j);
243                order.swap(i, j);
244            } else {
245                break;
246            }
247        }
248        if first != j {
249            code.swap(first, j);
250            order.swap(first, j);
251            let i = j + 1;
252            if last - i <= j - first {
253                sort_codes(i, last, code, order);
254                last = j;
255            } else {
256                sort_codes(first, j, code, order);
257                first = i;
258            }
259        } else {
260            first += 1;
261        }
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    /// LSB-first bit writer, the mirror of the decoder's reader.
270    struct BitWriter {
271        out: Vec<u8>,
272        acc: u32,
273        n: u32,
274    }
275
276    impl BitWriter {
277        fn put(&mut self, v: u32, bits: u32) {
278            for b in 0..bits {
279                self.acc |= ((v >> b) & 1) << self.n;
280                self.n += 1;
281                if self.n == 8 {
282                    self.out.push(self.acc as u8);
283                    self.acc = 0;
284                    self.n = 0;
285                }
286            }
287        }
288        fn align(&mut self) {
289            if self.n > 0 {
290                self.out.push(self.acc as u8);
291                self.acc = 0;
292                self.n = 0;
293            }
294        }
295    }
296
297    /// Codeword bits per symbol for `lengths`, the way [`build_tree`] assigns them.
298    fn codes_for(lengths: &[u8]) -> Vec<u64> {
299        let n = lengths.len();
300        let mut sorted = lengths.to_vec();
301        let mut order: Vec<u16> = (0..n as u16).collect();
302        sort_codes(0, n, &mut sorted, &mut order);
303        let mut bits = vec![0u64; n];
304        let mut i = sorted.iter().position(|&c| c != 0).unwrap_or(n);
305        let mut j = 0u64;
306        while i < n {
307            if i > 0 {
308                j <<= sorted[i] - sorted[i - 1];
309            }
310            let (mut m, mut l) = (0u64, j);
311            for _ in 0..sorted[i] {
312                m = (m << 1) | (l & 1);
313                l >>= 1;
314            }
315            bits[order[i] as usize] = m;
316            i += 1;
317            j += 1;
318        }
319        bits
320    }
321
322    /// A flat (non-nested) tree header: every symbol one `width`-bit length (`len - 1`).
323    fn put_flat_tree(w: &mut BitWriter, lengths: &[u8], width: u32) {
324        w.put(0, 1); // no "absent" value
325        w.put(width - 2, 2);
326        w.put(0, 3); // offset 1
327        w.put(0, 2); // flat
328        for &l in lengths {
329            w.put(l as u32 - 1, width);
330        }
331        w.align();
332    }
333
334    /// Encode `data` as literals plus one back-reference, then decode it back.
335    #[test]
336    fn round_trips_literals_and_a_match() {
337        // A complete code, as real streams carry: 236 symbols of 8 bits and 40 of 9.
338        let litlen: Vec<u8> = (0..LITLEN_SYMBOLS)
339            .map(|s| if s < 236 { 8 } else { 9 })
340            .collect();
341        let dist: Vec<u8> = vec![4; DIST_SYMBOLS];
342        let (lc, dc) = (codes_for(&litlen), codes_for(&dist));
343        let mut w = BitWriter {
344            out: Vec::new(),
345            acc: 0,
346            n: 0,
347        };
348        put_flat_tree(&mut w, &litlen, 4);
349        put_flat_tree(&mut w, &dist, 3);
350        let text = b"NDIF chunk, ";
351        for &b in text.iter() {
352            w.put(lc[b as usize] as u32, 8);
353        }
354        // Length symbol 258 is 6 with no extra bits; distance symbol 3 is 8 plus 3 extra bits.
355        w.put(lc[258] as u32, 9);
356        w.put(dc[3] as u32, 4);
357        w.put(4, 3);
358        w.align();
359
360        let mut expect = text.to_vec();
361        for _ in 0..6 {
362            expect.push(expect[expect.len() - 12]);
363        }
364        let mut chunk = (w.out.len() as u32 + 16).to_le_bytes().to_vec();
365        chunk.extend_from_slice(&(expect.len() as u32).to_le_bytes());
366        chunk.extend_from_slice(&(expect.len() as u32).to_le_bytes());
367        chunk.extend_from_slice(&[0xff; 4]);
368        chunk.extend_from_slice(&w.out);
369        assert_eq!(decompress(&chunk, 1 << 20).unwrap(), expect);
370        assert!(decompress(&chunk, 4).is_err(), "over the caller's cap");
371        assert!(
372            decompress(&chunk[..10], 1 << 20).is_err(),
373            "truncated header"
374        );
375    }
376
377    /// Garbage must error out rather than panic or spin.
378    #[test]
379    fn garbage_errors_cleanly() {
380        let mut chunk = vec![
381            0u8, 0, 0, 0, 0, 0x10, 0, 0, 0, 0x10, 0, 0, 0xff, 0xff, 0xff, 0xff,
382        ];
383        chunk.extend((0..200u32).map(|i| (i.wrapping_mul(37) ^ 0x5a) as u8));
384        let _ = decompress(&chunk, 1 << 20);
385        let _ = decompress(&[0u8; 16], 1 << 20);
386    }
387}
arc_mac/shrinkwrap.rs

arc_mac/
shrinkwrap.rs
