Add separate Huffman decoder and improve LZXpress Huffman

dissect/util/compression/huffman.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from dissect.util.stream import BitStream
+
+
+class HuffmanDecoder:
+    """Canonical Huffman decoder using a lookup table.
+
+    Formats that describe their Huffman tree in some other way are expected to convert it to a list of code lengths
+    and import those using :meth:`import_lengths`.
+
+    Decoding works on any bit stream that provides ``peek`` and ``remove`` methods, such as
+    :class:`~dissect.util.stream.BitStream`.
+
+    Args:
+        num_codes: The number of codes in the alphabet.
+        max_bits: The maximum length of a code in bits.
+    """
+
+    def __init__(self, num_codes: int, max_bits: int):
+        self.num_codes = num_codes
+        self.max_bits = max_bits
+
+        self.lengths = [0] * num_codes
+        self.codes = [0] * num_codes
+        self.table: list[tuple[int, int] | None] = [None] * (1 << max_bits)
+
+    def import_lengths(self, lengths: list[int], longest_first: bool = False) -> None:
+        """Import the code lengths of all symbols and build the lookup table.
+
+        Codes are assigned canonically, ordered by code length and symbol. By default the lowest codes are assigned
+        to the shortest code lengths, like DEFLATE. Some formats (e.g. MAME) assign the lowest codes to the longest
+        code lengths instead.
+
+        The code lengths must describe a complete tree, unless fewer than two symbols are used.
+
+        Args:
+            lengths: The code length of each symbol, where 0 means that the symbol is unused.
+            longest_first: Whether to assign the lowest codes to the longest code lengths.
+        """
+        if len(lengths) != self.num_codes:
+            raise ValueError("Invalid number of Huffman code lengths")
+
+        self.lengths = lengths
+        self._assign_canonical_codes(longest_first)
+        self._build_lookup_table()
+
+    def decode(self, stream: BitStream) -> int:
+        """Decode a single symbol.
+
+        Args:
+            stream: The bit stream to decode from.
+        """
+        if (entry := self.table[stream.peek(self.max_bits)]) is None:
+            raise ValueError("Invalid Huffman code")
+
+        symbol, length = entry
+        stream.remove(length)
+        return symbol
+
+    def _assign_canonical_codes(self, longest_first: bool) -> None:
+        histogram = [0] * (self.max_bits + 1)
+        for length in self.lengths:
+            if length > self.max_bits:
+                raise ValueError("Invalid Huffman code length")
+            histogram[length] += 1
+
+        # Make sure the codes exactly fill the code space, a tree with fewer than two codes can't be complete
+        used = sum(histogram[length] << (self.max_bits - length) for length in range(1, self.max_bits + 1))
+        if used > 1 << self.max_bits or (used < 1 << self.max_bits and sum(histogram[1:]) > 1):
+            raise ValueError("Invalid Huffman tree")
+
+        # For each code length, determine the starting code
+        start = 0
+        if longest_first:
+            for length in range(self.max_bits, 0, -1):
+                histogram[length], start = start, (start + histogram[length]) >> 1
+        else:
+            for length in range(1, self.max_bits + 1):
+                histogram[length], start = start, (start + histogram[length]) << 1
+
+        self.codes = [0] * self.num_codes
+        for symbol, length in enumerate(self.lengths):
+            if length > 0:
+                self.codes[symbol] = histogram[length]
+                histogram[length] += 1
+
+    def _build_lookup_table(self) -> None:
+        self.table = [None] * (1 << self.max_bits)
+        for symbol, (length, code) in enumerate(zip(self.lengths, self.codes, strict=True)):
+            if length > 0:
+                shift = self.max_bits - length
+                self.table[code << shift : (code + 1) << shift] = [(symbol, length)] * (1 << shift)

dissect/util/compression/lzxpress_huffman.py

@@ -1,135 +1,39 @@
 # https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-frs2/8cb5bae9-edf3-4833-9f0a-9d7e24218d3d
-# https://winprotocoldoc.blob.core.windows.net/productionwindowsarchives/MS-XCA/[MS-XCA].pdf
+# https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-xca/a8b7cb0a-92a6-4187-a23b-5e14273b96f8
 from __future__ import annotations
 
 import io
-import struct
-from typing import BinaryIO, NamedTuple
+from typing import BinaryIO
 
+from dissect.util.compression.huffman import HuffmanDecoder
+from dissect.util.stream import BitStream
 
-class Symbol(NamedTuple):
-    length: int
-    symbol: int
 
-
-def _read_16_bit(fh: BinaryIO) -> int:
-    return struct.unpack("<H", fh.read(2).rjust(2, b"\x00"))[0]
-
-
-class Node:
-    __slots__ = ("children", "is_leaf", "symbol")
-
-    def __init__(self, symbol: Symbol | None = None, is_leaf: bool = False):
-        self.symbol = symbol
-        self.is_leaf = is_leaf
-        self.children = [None, None]
-
-
-def _add_leaf(nodes: list[Node], idx: int, mask: int, bits: int) -> int:
-    node = nodes[0]
-    i = idx + 1
-
-    while bits > 1:
-        bits -= 1
-        childidx = (mask >> bits) & 1
-        if node.children[childidx] is None:
-            node.children[childidx] = nodes[i]
-            nodes[i].is_leaf = False
-            i += 1
-        node = node.children[childidx]
-
-    node.children[mask & 1] = nodes[idx]
-    return i
-
-
-def _build_tree(buf: bytes) -> Node:
+def _build_decoder(buf: bytes) -> HuffmanDecoder:
     if len(buf) != 256:
         raise ValueError("Not enough data for Huffman code tree")
 
-    nodes = [Node() for _ in range(1024)]
-    symbols: list[Symbol] = []
-
-    for i, c in enumerate(buf):
-        symbols.append(Symbol(c & 0x0F, i * 2))
-        symbols.append(Symbol((c >> 4) & 0x0F, i * 2 + 1))
-
-    symbols = sorted(symbols)
-
-    symbol_index_start = 0
-    for s in symbols:
-        if s.length > 0:
-            break
-        symbol_index_start += 1
-
-    mask = 0
-    bits = 1
-
-    root = nodes[0]
-
-    tree_index = 1
-    for symbol_index in range(symbol_index_start, 512):
-        s = symbols[symbol_index]
-
-        node = nodes[tree_index]
-        node.symbol = s.symbol
-        node.is_leaf = True
+    lengths = []
+    for c in buf:
+        lengths.append(c & 0x0F)
+        lengths.append((c >> 4) & 0x0F)
 
-        mask = (mask << s.length - bits) & 0xFFFFFFFF
-        bits = s.length
+    decoder = HuffmanDecoder(512, 15)
+    decoder.import_lengths(lengths)
+    return decoder
 
-        tree_index = _add_leaf(nodes, tree_index, mask, bits)
-        mask += 1
 
-    return root
+def decompress(src: bytes | BinaryIO, size: int | None = None) -> bytes:
+    """LZXPRESS Huffman decompress from a file-like object or bytes.
 
+    Decompresses until the decompressed size is reached or, if it's not given, until the end-of-stream symbol.
 
-class BitString:
-    def __init__(self):
-        self.source = None
-        self.mask = 0
-        self.bits = 0
-
-    @property
-    def index(self) -> int:
-        return self.source.tell()
-
-    def init(self, fh: BinaryIO) -> None:
-        self.mask = (_read_16_bit(fh) << 16) + _read_16_bit(fh)
-        self.bits = 32
-        self.source = fh
-
-    def read(self, n: int) -> bytes:
-        return self.source.read(n)
-
-    def lookup(self, n: int) -> int:
-        if n == 0:
-            return 0
-
-        return self.mask >> (32 - n)
-
-    def skip(self, n: int) -> None:
-        self.mask = (self.mask << n) & 0xFFFFFFFF
-        self.bits -= n
-        if self.bits < 16:
-            self.mask += _read_16_bit(self.source) << (16 - self.bits)
-            self.bits += 16
-
-    def decode(self, root: Node) -> Symbol:
-        node = root
-        while not node.is_leaf:
-            bit = self.lookup(1)
-            self.skip(1)
-            node = node.children[bit]
-        return node.symbol
-
-
-def decompress(src: bytes | BinaryIO) -> bytes:
-    """LZXPRESS decompress from a file-like object or bytes.
-
-    Decompresses until EOF of the input data.
+    The end-of-stream symbol is also a valid match symbol, so without a decompressed size it's only recognized as such
+    when nothing but zero bits follow it. Pass the decompressed size whenever it's known.
 
     Args:
         src: File-like object or bytes to decompress.
+        size: The decompressed size, if known.
 
     Returns:
         The decompressed data.
@@ -141,49 +45,62 @@ def decompress(src: bytes | BinaryIO) -> bytes:
 
     start_offset = src.tell()
     src.seek(0, io.SEEK_END)
-    size = src.tell() - start_offset
+    end_offset = src.tell()
     src.seek(start_offset, io.SEEK_SET)
 
-    bitstring = BitString()
+    while end_offset - src.tell() >= 256 and (size is None or len(dst) < size):
+        decoder = _build_decoder(src.read(256))
+
+        # The reference decoder keeps at least 16 bits buffered after every removal
+        # The extra length bytes and the next block are read directly after these buffered bits, so we do the same
+        stream = BitStream(src, word_size=2, byteorder="little")
+        stream.ensure(16)
+
+        # Matches can extend past the end of a block, which extends the block
+        block_end = len(dst) + 65536
+        if size is not None:
+            block_end = min(block_end, size)
 
-    while src.tell() - start_offset < size:
-        root = _build_tree(src.read(256))
-        bitstring.init(src)
+        while len(dst) < block_end and stream.bits:
+            symbol = decoder.decode(stream)
+            stream.ensure(16)
 
-        chunk_size = 0
-        while chunk_size < 65536 and src.tell() - start_offset < size:
-            symbol = bitstring.decode(root)
             if symbol < 256:
                 dst.append(symbol)
-                chunk_size += 1
-            else:
-                symbol -= 256
-                length = symbol & 0x0F
-                symbol >>= 4
+                continue
+
+            if symbol == 256 and size is None and src.tell() >= end_offset and not stream.buffer:
+                # End-of-stream symbol, the input is fully read and only zero bits remain
+                return bytes(dst)
 
-                offset = (1 << symbol) + bitstring.lookup(symbol)
+            symbol -= 256
+            length = symbol & 0x0F
+            symbol >>= 4
 
-                if length == 15:
-                    length = ord(bitstring.read(1)) + 15
+            offset = (1 << symbol) + stream.read(symbol)
 
-                    if length == 270:
-                        length = _read_16_bit(bitstring.source)
+            if length == 15:
+                length = ord(stream.read_bytes(1)) + 15
 
-                bitstring.skip(symbol)
+                if length == 270:
+                    length = int.from_bytes(stream.read_bytes(2), "little")
 
-                length += 3
+                    if length == 0:
+                        length = int.from_bytes(stream.read_bytes(4), "little")
 
-                if offset > len(dst):
-                    raise ValueError("Invalid match distance")
+            stream.ensure(16)
 
-                if length <= offset:
-                    dst += dst[-offset : (-offset + length) or None]
-                else:
-                    # The match overlaps the output it produces, so it repeats the last offset bytes
-                    chunk = dst[-offset:]
-                    dst += chunk * (length // offset)
-                    dst += chunk[: length % offset]
+            length += 3
 
-                chunk_size += length
+            if offset > len(dst):
+                raise ValueError("Invalid match distance")
+
+            if length <= offset:
+                dst += dst[-offset : (-offset + length) or None]
+            else:
+                # The match overlaps the output it produces, so it repeats the last offset bytes
+                chunk = dst[-offset:]
+                dst += chunk * (length // offset)
+                dst += chunk[: length % offset]
 
-    return bytes(dst)
+    return bytes(dst[:size])