Add LEB128 variable-lengh integer support (#73)

Author: joost-j

dissect/cstruct/__init__.py

@@ -14,6 +14,7 @@
 from dissect.cstruct.types.enum import Enum, EnumInstance
 from dissect.cstruct.types.flag import Flag, FlagInstance
 from dissect.cstruct.types.instance import Instance
+from dissect.cstruct.types.leb128 import LEB128
 from dissect.cstruct.types.packedtype import PackedType
 from dissect.cstruct.types.pointer import Pointer, PointerInstance
 from dissect.cstruct.types.structure import Field, Structure, Union
@@ -44,6 +45,7 @@
     "Union",
     "Field",
     "Instance",
+    "LEB128",
     "Structure",
     "Expression",
     "PackedType",

dissect/cstruct/cstruct.py

@@ -7,6 +7,7 @@
 from dissect.cstruct.exceptions import ResolveError
 from dissect.cstruct.parser import CStyleParser, TokenParser
 from dissect.cstruct.types import (
+    LEB128,
     Array,
     BaseType,
     BytesInteger,
@@ -59,6 +60,9 @@ def __init__(self, endian: str = "<", pointer: Optional[str] = None):
             "int128": BytesInteger(self, "int128", 16, True, alignment=16),
             "uint128": BytesInteger(self, "uint128", 16, False, alignment=16),
 
+            "uleb128": LEB128(self, 'uleb128', None, False),
+            "ileb128": LEB128(self, 'ileb128', None, True),
+
             "void": VoidType(),
 
             # Common C types not covered by internal types

dissect/cstruct/types/__init__.py

@@ -4,6 +4,7 @@
 from dissect.cstruct.types.enum import Enum, EnumInstance
 from dissect.cstruct.types.flag import Flag, FlagInstance
 from dissect.cstruct.types.instance import Instance
+from dissect.cstruct.types.leb128 import LEB128
 from dissect.cstruct.types.packedtype import PackedType
 from dissect.cstruct.types.pointer import Pointer, PointerInstance
 from dissect.cstruct.types.structure import Field, Structure, Union
@@ -21,6 +22,7 @@
     "Flag",
     "FlagInstance",
     "Instance",
+    "LEB128",
     "PackedType",
     "Pointer",
     "PointerInstance",

dissect/cstruct/types/leb128.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, BinaryIO
+
+from dissect.cstruct.types.base import RawType
+
+if TYPE_CHECKING:
+    from dissect.cstruct import cstruct
+
+
+class LEB128(RawType):
+    """Variable-length code compression to store an arbitrarily large integer in a small number of bytes.
+
+    See https://en.wikipedia.org/wiki/LEB128 for more information and an explanation of the algorithm.
+    """
+
+    signed: bool
+
+    def __init__(self, cstruct: cstruct, name: str, size: int, signed: bool, alignment: int = 1):
+        self.signed = signed
+        super().__init__(cstruct, name, size, alignment)
+
+    def _read(self, stream: BinaryIO, context: dict[str, Any] = None) -> LEB128:
+        result = 0
+        shift = 0
+        while True:
+            b = stream.read(1)
+            if b == b"":
+                raise EOFError("EOF reached, while final LEB128 byte was not yet read")
+
+            b = ord(b)
+            result |= (b & 0x7F) << shift
+            shift += 7
+            if (b & 0x80) == 0:
+                break
+
+        if self.signed:
+            if b & 0x40 != 0:
+                result |= ~0 << shift
+
+        return result
+
+    def _read_0(self, stream: BinaryIO, context: dict[str, Any] = None) -> LEB128:
+        result = []
+
+        while True:
+            if (value := self._read(stream, context)) == 0:
+                break
+
+            result.append(value)
+
+        return result
+
+    def _write(self, stream: BinaryIO, data: int) -> int:
+        # only write negative numbers when in signed mode
+        if data < 0 and not self.signed:
+            raise ValueError("Attempt to encode a negative integer using unsigned LEB128 encoding")
+
+        result = bytearray()
+        while True:
+            # low-order 7 bits of value
+            byte = data & 0x7F
+            data = data >> 7
+
+            # function works similar for signed- and unsigned integers, except for the check when to stop
+            # the encoding process.
+            if (self.signed and (data == 0 and byte & 0x40 == 0) or (data == -1 and byte & 0x40 != 0)) or (
+                not self.signed and data == 0
+            ):
+                result.append(byte)
+                break
+
+            # Set high-order bit of byte
+            result.append(0x80 | byte)
+
+        stream.write(result)
+        return len(result)
+
+    def _write_0(self, stream: BinaryIO, data: list[int]) -> int:
+        return self._write_array(stream, data + [0])
+
+    def default(self) -> int:
+        return 0
+
+    def default_array(self, count: int) -> list[int]:
+        return [0] * count

tests/test_types_leb128.py

@@ -0,0 +1,210 @@
+import io
+
+import pytest
+from dissect import cstruct
+
+
+def test_leb128_unsigned_read_EOF():
+    cs = cstruct.cstruct()
+
+    with pytest.raises(EOFError, match="EOF reached, while final LEB128 byte was not yet read"):
+        cs.uleb128(b"\x8b")
+
+
+def test_leb128_unsigned_read():
+    cs = cstruct.cstruct()
+
+    assert cs.uleb128(b"\x02") == 2
+    assert cs.uleb128(b"\x8b\x25") == 4747
+    assert cs.uleb128(b"\xc9\x8f\xb0\x06") == 13371337
+    assert cs.uleb128(b"\x7e") == 126
+    assert cs.uleb128(b"\xf5\x5a") == 11637
+    assert cs.uleb128(b"\xde\xd6\xcf\x7c") == 261352286
+
+
+def test_leb128_signed_read():
+    cs = cstruct.cstruct()
+
+    assert cs.ileb128(b"\x02") == 2
+    assert cs.ileb128(b"\x8b\x25") == 4747
+    assert cs.ileb128(b"\xc9\x8f\xb0\x06") == 13371337
+    assert cs.ileb128(b"\x7e") == -2
+    assert cs.ileb128(b"\xf5\x5a") == -4747
+    assert cs.ileb128(b"\xde\xd6\xcf\x7c") == -7083170
+
+
+@pytest.mark.parametrize("compiled", [True, False])
+def test_leb128_struct_unsigned(compiled):
+    cdef = """
+    struct test {
+        uleb128 len;
+        char    data[len];
+    };
+    """
+    cs = cstruct.cstruct()
+    cs.load(cdef, compiled=compiled)
+
+    buf = b"\xaf\x18"
+    buf += b"\x41" * 3119
+    obj = cs.test(buf)
+
+    assert obj.len == 3119
+    assert obj.data == (b"\x41" * 3119)
+    assert len(obj.data) == 3119
+    assert len(buf) == 3119 + 2
+
+    assert obj.dumps() == buf
+
+
+@pytest.mark.parametrize("compiled", [True, False])
+def test_leb128_struct_unsigned_zero(compiled):
+    cdef = """
+    struct test {
+        uleb128 numbers[];
+    };
+    """
+    cs = cstruct.cstruct()
+    cs.load(cdef, compiled=compiled)
+
+    buf = b"\xaf\x18\x8b\x25\xc9\x8f\xb0\x06\x00"
+    obj = cs.test(buf)
+
+    assert len(obj.numbers) == 3
+    assert obj.numbers[0] == 3119
+    assert obj.numbers[1] == 4747
+    assert obj.numbers[2] == 13371337
+
+    assert obj.dumps() == buf
+
+
+@pytest.mark.parametrize("compiled", [True, False])
+def test_leb128_struct_signed_zero(compiled):
+    cdef = """
+    struct test {
+        ileb128 numbers[];
+    };
+    """
+    cs = cstruct.cstruct()
+    cs.load(cdef, compiled=compiled)
+
+    buf = b"\xaf\x18\xf5\x5a\xde\xd6\xcf\x7c\x00"
+    obj = cs.test(buf)
+
+    assert len(obj.numbers) == 3
+    assert obj.numbers[0] == 3119
+    assert obj.numbers[1] == -4747
+    assert obj.numbers[2] == -7083170
+
+    assert obj.dumps() == buf
+
+
+@pytest.mark.parametrize("compiled", [True, False])
+def test_leb128_nested_struct_unsigned(compiled):
+    cdef = """
+    struct entry {
+        uleb128 len;
+        char    data[len];
+        uint32  crc;
+    };
+    struct nested {
+        uleb128 name_len;
+        char    name[name_len];
+        uleb128 n_entries;
+        entry   entries[n_entries];
+    };
+    """
+    cs = cstruct.cstruct()
+    cs.load(cdef, compiled=compiled)
+
+    # Dummy file format specifying 300 entries
+    buf = b"\x08\x54\x65\x73\x74\x66\x69\x6c\x65\xac\x02"
+
+    # Each entry has 4 byte data + 4 byte CRC
+    buf += b"\x04\x41\x41\x41\x41\x42\x42\x42\x42" * 300
+
+    obj = cs.nested(buf)
+
+    assert obj.name_len == 8
+    assert obj.name == b"\x54\x65\x73\x74\x66\x69\x6c\x65"
+    assert obj.n_entries == 300
+
+    assert obj.dumps() == buf
+
+
+@pytest.mark.parametrize("compiled", [True, False])
+def test_leb128_nested_struct_signed(compiled):
+    cdef = """
+    struct entry {
+        ileb128 len;
+        char    data[len];
+        uint32  crc;
+    };
+    struct nested {
+        ileb128 name_len;
+        char    name[name_len];
+        ileb128 n_entries;
+        entry   entries[n_entries];
+    };
+    """
+    cs = cstruct.cstruct()
+    cs.load(cdef, compiled=compiled)
+
+    # Dummy file format specifying 300 entries
+    buf = b"\x08\x54\x65\x73\x74\x66\x69\x6c\x65\xac\x02"
+
+    # Each entry has 4 byte data + 4 byte CRC
+    buf += b"\x04\x41\x41\x41\x41\x42\x42\x42\x42" * 300
+
+    obj = cs.nested(buf)
+
+    assert obj.name_len == 8
+    assert obj.name == b"\x54\x65\x73\x74\x66\x69\x6c\x65"
+    assert obj.n_entries == 300
+
+    assert obj.dumps() == buf
+
+
+def test_leb128_unsigned_write():
+    cs = cstruct.cstruct()
+
+    assert cs.uleb128.dumps(2) == b"\x02"
+    assert cs.uleb128.dumps(4747) == b"\x8b\x25"
+    assert cs.uleb128.dumps(13371337) == b"\xc9\x8f\xb0\x06"
+    assert cs.uleb128.dumps(126) == b"\x7e"
+    assert cs.uleb128.dumps(11637) == b"\xf5\x5a"
+    assert cs.uleb128.dumps(261352286) == b"\xde\xd6\xcf\x7c"
+
+
+def test_leb128_signed_write():
+    cs = cstruct.cstruct()
+
+    assert cs.ileb128.dumps(2) == b"\x02"
+    assert cs.ileb128.dumps(4747) == b"\x8b\x25"
+    assert cs.ileb128.dumps(13371337) == b"\xc9\x8f\xb0\x06"
+    assert cs.ileb128.dumps(-2) == b"\x7e"
+    assert cs.ileb128.dumps(-4747) == b"\xf5\x5a"
+    assert cs.ileb128.dumps(-7083170) == b"\xde\xd6\xcf\x7c"
+
+
+def test_leb128_write_negatives():
+    cs = cstruct.cstruct()
+
+    with pytest.raises(ValueError, match="Attempt to encode a negative integer using unsigned LEB128 encoding"):
+        cs.uleb128.dumps(-2)
+    assert cs.ileb128.dumps(-2) == b"\x7e"
+
+
+def test_leb128_unsigned_write_amount_written():
+    cs = cstruct.cstruct()
+
+    out1 = io.BytesIO()
+    bytes_written1 = cs.uleb128.write(out1, 2)
+    assert bytes_written1 == out1.tell()
+
+    out2 = io.BytesIO()
+    bytes_written2 = cs.uleb128.write(out2, 4747)
+    assert bytes_written2 == out2.tell()
+
+    out3 = io.BytesIO()
+    bytes_written3 = cs.uleb128.write(out3, 13371337)
+    assert bytes_written3 == out3.tell()