From 50c3b904c8a558198300e1ca324b3ac398c6c7b2 Mon Sep 17 00:00:00 2001
From: Schamper <1254028+Schamper@users.noreply.github.com>
Date: Fri, 23 Oct 2026 15:34:14 +0200
Subject: [PATCH 1/2] Add separate huffman decoder

---
 dissect/util/compression/huffman.py          |  96 +++++++++++++
 dissect/util/compression/lzxpress_huffman.py | 143 +++----------------
 dissect/util/stream.py                       |  78 +++++++---
 tests/compression/test_huffman.py            |  73 ++++++++++
 tests/compression/test_lzxpress_huffman.py   |  23 +++
 tests/test_stream.py                         |  53 +++++++
 6 files changed, 329 insertions(+), 137 deletions(-)
 create mode 100644 dissect/util/compression/huffman.py
 create mode 100644 tests/compression/test_huffman.py

diff --git a/dissect/util/compression/huffman.py b/dissect/util/compression/huffman.py
new file mode 100644
index 0000000..9a5183b
--- /dev/null
+++ b/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)
diff --git a/dissect/util/compression/lzxpress_huffman.py b/dissect/util/compression/lzxpress_huffman.py
index a1e2569..5eb4a64 100644
--- a/dissect/util/compression/lzxpress_huffman.py
+++ b/dissect/util/compression/lzxpress_huffman.py
@@ -3,124 +3,24 @@
 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]
+    lengths = []
+    for c in buf:
+        lengths.append(c & 0x0F)
+        lengths.append((c >> 4) & 0x0F)
 
-        node = nodes[tree_index]
-        node.symbol = s.symbol
-        node.is_leaf = True
-
-        mask = (mask << s.length - bits) & 0xFFFFFFFF
-        bits = s.length
-
-        tree_index = _add_leaf(nodes, tree_index, mask, bits)
-        mask += 1
-
-    return root
-
-
-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
+    decoder = HuffmanDecoder(512, 15)
+    decoder.import_lengths(lengths)
+    return decoder
 
 
 def decompress(src: bytes | BinaryIO) -> bytes:
@@ -144,15 +44,18 @@ def decompress(src: bytes | BinaryIO) -> bytes:
     size = src.tell() - start_offset
     src.seek(start_offset, io.SEEK_SET)
 
-    bitstring = BitString()
-
     while src.tell() - start_offset < size:
-        root = _build_tree(src.read(256))
-        bitstring.init(src)
+        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 chunk are read directly after these buffered bits, so we do the same
+        stream = BitStream(src, word_size=2, byteorder="little")
 
         chunk_size = 0
         while chunk_size < 65536 and src.tell() - start_offset < size:
-            symbol = bitstring.decode(root)
+            symbol = decoder.decode(stream)
+            stream.ensure(16)
+
             if symbol < 256:
                 dst.append(symbol)
                 chunk_size += 1
@@ -161,15 +64,15 @@ def decompress(src: bytes | BinaryIO) -> bytes:
                 length = symbol & 0x0F
                 symbol >>= 4
 
-                offset = (1 << symbol) + bitstring.lookup(symbol)
+                offset = (1 << symbol) + stream.read(symbol)
 
                 if length == 15:
-                    length = ord(bitstring.read(1)) + 15
+                    length = ord(stream.read_bytes(1)) + 15
 
                     if length == 270:
-                        length = _read_16_bit(bitstring.source)
+                        length = int.from_bytes(stream.read_bytes(2), "little")
 
-                bitstring.skip(symbol)
+                stream.ensure(16)
 
                 length += 3
 
diff --git a/dissect/util/stream.py b/dissect/util/stream.py
index a0af22e..ee1be99 100644
--- a/dissect/util/stream.py
+++ b/dissect/util/stream.py
@@ -674,12 +674,18 @@ def readall(self) -> bytes:
 class BitStream:
     """Bit-level stream reader.
 
+    Bits are read MSB first from words of ``word_size`` bytes.
+
     Args:
         fh: File-like object to read bits from.
+        word_size: The size of a word in bytes.
+        byteorder: The byte order of a word, either ``"big"`` or ``"little"``.
     """
 
-    def __init__(self, fh: BinaryIO):
+    def __init__(self, fh: BinaryIO, word_size: int = 1, byteorder: str = "big"):
         self.fh = fh
+        self.word_size = word_size
+        self.byteorder = byteorder
         self._byte_offset = fh.tell()
 
         self.buffer = 0
@@ -703,19 +709,21 @@ def seek(self, pos: int, whence: int = io.SEEK_SET) -> int:
         Returns:
             The new stream position after seeking.
         """
+        word_bits = self.word_size * 8
+
         if whence == io.SEEK_SET:
-            byte_pos, bit_pos = divmod(pos, 8)
+            word_pos, bit_pos = divmod(pos, word_bits)
         elif whence == io.SEEK_CUR:
             current_pos = self.tell()
-            byte_pos, bit_pos = divmod(current_pos + pos, 8)
+            word_pos, bit_pos = divmod(current_pos + pos, word_bits)
         elif whence == io.SEEK_END:
             self.fh.seek(0, io.SEEK_END)
             end_pos = self.fh.tell() * 8
-            byte_pos, bit_pos = divmod(end_pos + pos, 8)
+            word_pos, bit_pos = divmod(end_pos + pos, word_bits)
         else:
             raise IOError("invalid whence value")
 
-        self._byte_offset = byte_pos
+        self._byte_offset = word_pos * self.word_size
         self.bits = 0
         self.buffer = 0
 
@@ -738,6 +746,49 @@ def read(self, n: int) -> int:
         self.remove(n)
         return value
 
+    def read_bytes(self, n: int) -> bytes:
+        """Read n bytes from the underlying stream, directly following the buffered bits.
+
+        The buffered bits are not affected, which allows reading formats that interleave bytes with the bit stream.
+
+        Args:
+            n: Number of bytes to read.
+        """
+        self.fh.seek(self._byte_offset)
+        buf = self.fh.read(n)
+        self._byte_offset += len(buf)
+        return buf
+
+    def ensure(self, n: int) -> None:
+        """Ensure at least n bits are buffered, reading whole words from the underlying stream.
+
+        Args:
+            n: Number of bits to buffer.
+        """
+        word_bits = self.word_size * 8
+
+        while self.bits < n:
+            num_words = (n - self.bits + word_bits - 1) // word_bits
+            self.fh.seek(self._byte_offset)
+            if not (buf := self.fh.read(num_words * self.word_size)):
+                break
+
+            if len(buf) % self.word_size:
+                # Pad a partial word at the end of the stream with zero bytes
+                buf = buf.ljust(len(buf) + self.word_size - len(buf) % self.word_size, b"\x00")
+
+            if self.byteorder == "big" or self.word_size == 1:
+                new_bits = int.from_bytes(buf, "big")
+            else:
+                new_bits = 0
+                for i in range(0, len(buf), self.word_size):
+                    new_bits = (new_bits << word_bits) | int.from_bytes(buf[i : i + self.word_size], self.byteorder)
+
+            num_new_bits = len(buf) * 8
+            self.buffer = (self.buffer << num_new_bits) | new_bits
+            self.bits += num_new_bits
+            self._byte_offset += len(buf)
+
     def peek(self, n: int) -> int:
         """Peek n bits from the stream without advancing.
 
@@ -747,20 +798,13 @@ def peek(self, n: int) -> int:
         if n == 0:
             return 0
 
-        if n > self.bits:
-            while self.bits < n:
-                num_bytes = (n - self.bits + 7) // 8
-                self.fh.seek(self._byte_offset)
-                if not (buf := self.fh.read(num_bytes)):
-                    break
+        self.ensure(n)
 
-                new_bits = int.from_bytes(buf, "big")
-                num_new_bits = len(buf) * 8
-                self.buffer = (self.buffer << num_new_bits) | new_bits
-                self.bits += num_new_bits
-                self._byte_offset += len(buf)
+        if n > self.bits:
+            # Pad with zero bits when peeking past the end of the stream
+            return self.buffer << (n - self.bits)
 
-        return self.buffer >> (self.bits - min(n, self.bits))
+        return self.buffer >> (self.bits - n)
 
     def remove(self, n: int) -> None:
         """Remove n bits from the stream.
diff --git a/tests/compression/test_huffman.py b/tests/compression/test_huffman.py
new file mode 100644
index 0000000..0446d4c
--- /dev/null
+++ b/tests/compression/test_huffman.py
@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import io
+
+import pytest
+
+from dissect.util.compression.huffman import HuffmanDecoder
+from dissect.util.stream import BitStream
+
+
+def test_huffman_import_lengths() -> None:
+    # Example from RFC 1951, section 3.2.2
+    decoder = HuffmanDecoder(8, 4)
+    decoder.import_lengths([3, 3, 3, 3, 3, 2, 4, 4])
+
+    assert decoder.codes == [0b010, 0b011, 0b100, 0b101, 0b110, 0b00, 0b1110, 0b1111]
+
+    stream = BitStream(io.BytesIO(bytes([0b00010011, 0b11101111])))
+    assert [decoder.decode(stream) for _ in range(5)] == [5, 0, 1, 6, 7]
+
+
+def test_huffman_import_lengths_longest_first() -> None:
+    decoder = HuffmanDecoder(8, 4)
+    decoder.import_lengths([3, 3, 3, 3, 3, 2, 4, 4], longest_first=True)
+
+    assert decoder.codes == [0b001, 0b010, 0b011, 0b100, 0b101, 0b11, 0b0000, 0b0001]
+
+    stream = BitStream(io.BytesIO(bytes([0b11001010, 0b00000001])))
+    assert [decoder.decode(stream) for _ in range(5)] == [5, 0, 1, 6, 7]
+
+
+def test_huffman_import_lengths_invalid() -> None:
+    decoder = HuffmanDecoder(4, 2)
+
+    with pytest.raises(ValueError, match="Invalid number of Huffman code lengths"):
+        decoder.import_lengths([1, 1])
+
+    with pytest.raises(ValueError, match="Invalid Huffman code length"):
+        decoder.import_lengths([1, 3, 0, 0])
+
+    with pytest.raises(ValueError, match="Invalid Huffman tree"):
+        decoder.import_lengths([1, 1, 1, 0])
+
+    with pytest.raises(ValueError, match="Invalid Huffman tree"):
+        decoder.import_lengths([1, 1, 1, 0], longest_first=True)
+
+    # Incomplete trees
+    with pytest.raises(ValueError, match="Invalid Huffman tree"):
+        decoder.import_lengths([1, 2, 0, 0])
+
+    with pytest.raises(ValueError, match="Invalid Huffman tree"):
+        decoder.import_lengths([1, 2, 0, 0], longest_first=True)
+
+
+def test_huffman_import_lengths_single_code() -> None:
+    decoder = HuffmanDecoder(4, 2)
+
+    # Trees with fewer than two codes can't be complete, but are allowed
+    decoder.import_lengths([0, 0, 0, 0])
+    decoder.import_lengths([0, 0, 2, 0])
+    assert decoder.codes == [0, 0, 0, 0]
+
+    decoder.import_lengths([0, 0, 2, 0], longest_first=True)
+    assert decoder.codes == [0, 0, 0, 0]
+
+
+def test_huffman_decode_invalid() -> None:
+    decoder = HuffmanDecoder(4, 2)
+    decoder.import_lengths([1, 0, 0, 0])
+
+    stream = BitStream(io.BytesIO(b"\xff"))
+    with pytest.raises(ValueError, match="Invalid Huffman code"):
+        decoder.decode(stream)
diff --git a/tests/compression/test_lzxpress_huffman.py b/tests/compression/test_lzxpress_huffman.py
index b4ef0ec..8cf7c0f 100644
--- a/tests/compression/test_lzxpress_huffman.py
+++ b/tests/compression/test_lzxpress_huffman.py
@@ -47,6 +47,29 @@
             "73d3dd96ca2e2f0144a117019256d770ee7c6febeaee09b24956c723ae22b529",
             id="large",
         ),
+        # Hand crafted with a code length of 9 for every symbol, spanning two chunks and both match length extensions
+        pytest.param(
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "0000404020305f160000fff7ff00009999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999999999999999999999999999999999999"
+            "9999999999999999999999999999999830718c7cfc346402ff2901ce9c5366d2"
+            "1900de0084200000",
+            "739c6b859e1e3ac7429d54ac622f6e0eceac09c9a5ae1180bbb70620b5d442b4",
+            id="length-extension",
+        ),
     ],
 )
 
diff --git a/tests/test_stream.py b/tests/test_stream.py
index 898ae4d..444c1a5 100644
--- a/tests/test_stream.py
+++ b/tests/test_stream.py
@@ -292,6 +292,10 @@ def test_bitstream() -> None:
     assert bitstream.seek(-8, io.SEEK_END) == 56
     assert bitstream.peek(8) == 0b01010101
 
+    fh = io.BytesIO(b"\x00\xe0")
+    bitstream = stream.BitStream(fh)
+    assert bitstream.read(16) == 0x00E0
+
 
 def test_bitstream_empty() -> None:
     """Test reading from an empty bit stream."""
@@ -302,6 +306,55 @@ def test_bitstream_empty() -> None:
     assert bitstream.read(8) == 0
 
 
+def test_bitstream_eof() -> None:
+    """Test reading past the end of a bit stream."""
+    fh = io.BytesIO(b"\xab")
+    bitstream = stream.BitStream(fh)
+
+    assert bitstream.read(4) == 0xA
+    assert bitstream.peek(8) == 0xB0
+    assert bitstream.read(8) == 0xB0
+    assert bitstream.tell() == 8
+
+
+@pytest.mark.parametrize(
+    ("byteorder", "expected", "expected_partial"),
+    [
+        ("big", 0b10101010_11001100_11110000_00001111, 0b11111111_00000000),
+        ("little", 0b11001100_10101010_00001111_11110000, 0b00000000_11111111),
+    ],
+)
+def test_bitstream_word_size(byteorder: str, expected: int, expected_partial: int) -> None:
+    """Test reading bits from 16-bit words."""
+    fh = io.BytesIO(bytes([0b10101010, 0b11001100, 0b11110000, 0b00001111, 0b11111111]))
+    bitstream = stream.BitStream(fh, word_size=2, byteorder=byteorder)
+
+    assert bitstream.peek(32) == expected
+    assert bitstream.read(4) == expected >> 28
+    assert bitstream.tell() == 4
+
+    # Words are always read whole
+    assert bitstream.seek(20) == 20
+    assert bitstream.bits == 12
+    assert bitstream.read(12) == expected & 0xFFF
+
+    # A partial word at the end of the stream is padded with zero bytes
+    assert bitstream.read(16) == expected_partial
+    assert bitstream.tell() == 48
+
+
+def test_bitstream_read_bytes() -> None:
+    """Test reading bytes interleaved with the bit stream."""
+    fh = io.BytesIO(b"\x12\x34\x56\x78\x9a\xbc")
+    bitstream = stream.BitStream(fh, word_size=2, byteorder="little")
+
+    assert bitstream.read(4) == 0x3
+    assert bitstream.read_bytes(1) == b"\x56"
+    assert bitstream.read(12) == 0x412
+    assert bitstream.read(16) == 0x9A78
+    assert bitstream.read_bytes(2) == b"\xbc"
+
+
 class NullStream(stream.AlignedStream):
     def __init__(self, size: int | None, align: int = stream.STREAM_BUFFER_SIZE):
         super().__init__(size)

From 5c77594013bc39c9ef5364d0be435ebfc7db0305 Mon Sep 17 00:00:00 2001
From: Schamper <1254028+Schamper@users.noreply.github.com>
Date: Tue, 27 Oct 2026 14:21:26 +0100
Subject: [PATCH 2/2] Several LZXpress Huffman fixes

---
 dissect/util/compression/lzxpress_huffman.py |  78 +++++++-----
 tests/compression/test_lzxpress_huffman.py   | 119 +++++++++++++++----
 2 files changed, 144 insertions(+), 53 deletions(-)

diff --git a/dissect/util/compression/lzxpress_huffman.py b/dissect/util/compression/lzxpress_huffman.py
index 5eb4a64..c239b24 100644
--- a/dissect/util/compression/lzxpress_huffman.py
+++ b/dissect/util/compression/lzxpress_huffman.py
@@ -1,5 +1,5 @@
 # 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
@@ -23,13 +23,17 @@ def _build_decoder(buf: bytes) -> HuffmanDecoder:
     return decoder
 
 
-def decompress(src: bytes | BinaryIO) -> bytes:
-    """LZXPRESS decompress from a file-like object or bytes.
+def decompress(src: bytes | BinaryIO, size: int | None = None) -> bytes:
+    """LZXPRESS Huffman decompress from a file-like object or bytes.
 
-    Decompresses until EOF of the input data.
+    Decompresses until the decompressed size is reached or, if it's not given, until the end-of-stream symbol.
+
+    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.
@@ -41,52 +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)
 
-    while src.tell() - start_offset < size:
+    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 chunk are read directly after these buffered bits, so we do the same
+        # 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)
 
-        chunk_size = 0
-        while chunk_size < 65536 and src.tell() - start_offset < size:
+        while len(dst) < block_end and stream.bits:
             symbol = decoder.decode(stream)
             stream.ensure(16)
 
             if symbol < 256:
                 dst.append(symbol)
-                chunk_size += 1
-            else:
-                symbol -= 256
-                length = symbol & 0x0F
-                symbol >>= 4
+                continue
 
-                offset = (1 << symbol) + stream.read(symbol)
+            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)
 
-                if length == 15:
-                    length = ord(stream.read_bytes(1)) + 15
+            symbol -= 256
+            length = symbol & 0x0F
+            symbol >>= 4
 
-                    if length == 270:
-                        length = int.from_bytes(stream.read_bytes(2), "little")
+            offset = (1 << symbol) + stream.read(symbol)
 
-                stream.ensure(16)
+            if length == 15:
+                length = ord(stream.read_bytes(1)) + 15
 
-                length += 3
+                if length == 270:
+                    length = int.from_bytes(stream.read_bytes(2), "little")
 
-                if offset > len(dst):
-                    raise ValueError("Invalid match distance")
+                    if length == 0:
+                        length = int.from_bytes(stream.read_bytes(4), "little")
 
-                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]
+            stream.ensure(16)
 
-                chunk_size += length
+            length += 3
+
+            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])
diff --git a/tests/compression/test_lzxpress_huffman.py b/tests/compression/test_lzxpress_huffman.py
index 8cf7c0f..8686c1a 100644
--- a/tests/compression/test_lzxpress_huffman.py
+++ b/tests/compression/test_lzxpress_huffman.py
@@ -47,28 +47,96 @@
             "73d3dd96ca2e2f0144a117019256d770ee7c6febeaee09b24956c723ae22b529",
             id="large",
         ),
-        # Hand crafted with a code length of 9 for every symbol, spanning two chunks and both match length extensions
+        # Compressed on Windows (level 1), exactly one full block followed by the end-of-stream symbol
         pytest.param(
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "0000404020305f160000fff7ff00009999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999999999999999999999999999999999999"
-            "9999999999999999999999999999999830718c7cfc346402ff2901ce9c5366d2"
-            "1900de0084200000",
-            "739c6b859e1e3ac7429d54ac622f6e0eceac09c9a5ae1180bbb70620b5d442b4",
-            id="length-extension",
+            "0000000000000000000000000000000003000000000000060000000000000000"
+            "0000000000000000000005000000000060665666565555455555555555050000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0500000000000000000000000000000000000000000000000000000000000000"
+            "0500000000000000000000000000005000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "e7379a0f503b05ea8f56c14a63330aa28087bd96e27cc2e380fbb76d0020ffcf"
+            "ff0000",
+            "311971d921f30a78ac1e087296cf20a9f8a2ac682de20efe6c99c89ba82c6b35",
+            id="exact-64k",
+        ),
+        # Compressed on Windows, multiple blocks, matches crossing blocks, both length extensions and symbol 256 matches
+        pytest.param(
+            "0000000000000000000000008000800084000000000008070000000000880080"
+            "0000000000000000000000000000000060666756566655665655555555050000"
+            "0000000000000000000000000000000000000000000000080000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0600000000000000000000000000000000000000000000000887800000000000"
+            "8007000000000000778687800000000077770780000000008766760000000000"
+            "7776070000000000088000000000008000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000700000000000000070"
+            "c271198958527a0d80279c241d220546d41b933c82158276d7f76bef89805305"
+            "0e8c25393a60e5723aa4a81847568af39d8ba120268ee1470a58a6905e509ca0"
+            "91a609a7648597e9bdba862ba4e242c7118bb1808008c6a2e6912a3e34c16fa4"
+            "ee11763e0811686eb7d9bcce99ff9dc7d4c0035355cdd4948d37bc61c4bdae2c"
+            "acdae140743cc369a2754fa484b5812b1642304796ac2ec805da138ea120083d"
+            "a77ca90560230ebe664901a566bb7ccbe6abcc8fd23e3758f960c7ad504e3b26"
+            "6cbe7ba3f53b7dfb72b5bc49162fb2df9bef36a2bde0103759fe097318323ddc"
+            "fbb5cff5d9cf0a4f1c895bb87e503f85be84ef86d4d6103e979c40f885acdb2b"
+            "58fe88f3ffaa58e84487c8ff2234516718dfff313dac838ebeffd72b7ad818f6"
+            "0000ff6f0c000000000000000080000008000000008400000000000707008000"
+            "0000880080000000000000000000000000000000086066675756565656665555"
+            "5656050000000000008000000000000000000000000008800000000000000000"
+            "0000000000800000000000000000000080008000000000000000000000000000"
+            "0000000000060000000000000000000000000000000000000000000000008780"
+            "0000000000000008000000000088860800000000008878078000000000000786"
+            "0800000000800000000000000000000000000000700000000000000070080000"
+            "0000000000700707000000007000770000000000607067880000000060800000"
+            "000000006083e1808a8151088746b0bda2152844e12b0a18e92b4202f4914c47"
+            "b48707a024c4754f841ae3d276e54d58f8307241276c9c53b2614f68b9c50656"
+            "0902263786fd96c9352855a71b3f7e0d5c0a8037918092677621147062998a81"
+            "63750cd4e87d0550d924506c68090791945eac8a674c0834802d19a0e6c22312"
+            "b902bf4bca792c82daee3908e3af5fe89867f9be5692ff9ba0d7760575a99082"
+            "0f9bea1ce92ab8f64d721f8eb1fff6136b848e51ff1c0b6190a70bd1c7f8a73d"
+            "f18c5e3d1a468df8d5003b232afe683f4d6b0f1d22ff5f04af57a5023fda1b55"
+            "1c006078725ef0c3ff9804679e382c8951eb7f85ff2115ec48fffd032b9122b0"
+            "7cddc1f1ff0102a2b0c47bb0cc1b495705c1244534672bff61e48424de614882"
+            "25aad4e05ea1d82cde820de6846d2cff07025e3bc762fb3dff4924fbe11e58ca"
+            "e1afadd042ffd6143f68fffd0328c2f4d7b78eff6b13302e5bb1ff11152e22ff"
+            "6f17915fff1c0bb172dc86ff370562acff21195eda6e3bff7209e1a5ff3a05f6"
+            "590000ff9c170000000000000000000000000000000000006400000000000006"
+            "0000000000000000000000000000000000000000000000000066665566565505"
+            "4656556566060000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000060000000000000000000000"
+            "0000000000000000050000000000000000000000000000000000000000000000"
+            "0600000000000000000600000000000000060660000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000005000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "604566000000005088de6c02955bd358a80959c0ff43057254952afde11748b7"
+            "b973e654975bb5b150f3522abbf1967e627b5c62dc21c8da6c81ec9d55552403"
+            "cfc2b464202f75ea39eee65e4e0d86eaab10d381701555cee08497ff3d030000"
+            "0000",
+            "de021bba2f87735aa82632e9ef264915c78eb784eae71291c0b2d41e1049de42",
+            id="multi-block",
+        ),
+        # Compressed on Windows, a match spanning multiple blocks using the 32-bit match length extension
+        pytest.param(
+            "0300000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000303300300000000030030000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000030000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "d652e0700000ff00003c0d030000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000003033"
+            "0030000000002003000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000003000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "0000000000000000000000000000000000000000000000000000000000000000"
+            "00000000000000000000000000227400dc0000",
+            "225eb609eb33b5263a2e655958a453be5b5e18b2c4a6b3a1862e21890faddd34",
+            id="long-match",
         ),
     ],
 )
@@ -90,6 +158,15 @@ def test_lzxpress_huffman_decompress_invalid_distance() -> None:
         lzxpress_huffman.decompress(bytes(table) + bytes.fromhex("0080000000000000"))
 
 
+@pytest.mark.parametrize(*PARAMS)
+def test_lzxpress_huffman_decompress_size(data: str, digest: str) -> None:
+    buf = bytes.fromhex(data)
+    result = lzxpress_huffman.decompress(buf)
+
+    assert hashlib.sha256(lzxpress_huffman.decompress(buf, len(result))).hexdigest() == digest
+    assert lzxpress_huffman.decompress(buf, 100) == result[:100]
+
+
 @pytest.mark.benchmark
 @pytest.mark.parametrize(*PARAMS)
 def test_benchmark_lzxpress_huffman_decompress(data: str, digest: str, benchmark: BenchmarkFixture) -> None: