Add support for EOF length identifier

Author: Schamper

dissect/cstruct/types/base.py

@@ -4,10 +4,13 @@
 from typing import TYPE_CHECKING, Any, BinaryIO, Optional, Union
 
 from dissect.cstruct.exceptions import ArraySizeError
-from dissect.cstruct.expression import Expression
 
 if TYPE_CHECKING:
     from dissect.cstruct.cstruct import cstruct
+    from dissect.cstruct.expression import Expression
+
+
+EOF = -0xE0F  # Negative counts are illegal anyway, so abuse that for our EOF sentinel
 
 
 class MetaType(type):
@@ -93,6 +96,15 @@ def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] = Non
             count: The amount of values to read.
             context: Optional reading context.
         """
+        if count == EOF:
+            result = []
+            while True:
+                try:
+                    result.append(cls._read(stream, context))
+                except EOFError:
+                    break
+            return result
+
         return [cls._read(stream, context) for _ in range(count)]
 
     def _read_0(cls, stream: BinaryIO, context: dict[str, Any] = None) -> list[BaseType]:
@@ -168,8 +180,17 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] = None) -> Array:
         if cls.null_terminated:
             return cls.type._read_0(stream, context)
 
-        num = cls.num_entries.evaluate(context) if cls.dynamic else cls.num_entries
-        return cls.type._read_array(stream, max(0, num), context)
+        if cls.dynamic:
+            try:
+                num = max(0, cls.num_entries.evaluate(context))
+            except Exception:
+                if cls.num_entries.expression != "EOF":
+                    raise
+                num = EOF
+        else:
+            num = max(0, cls.num_entries)
+
+        return cls.type._read_array(stream, num, context)
 
 
 class Array(list, BaseType, metaclass=ArrayMetaType):

dissect/cstruct/types/char.py

@@ -2,7 +2,7 @@
 
 from typing import Any, BinaryIO, Union
 
-from dissect.cstruct.types.base import ArrayMetaType, BaseType
+from dissect.cstruct.types.base import EOF, ArrayMetaType, BaseType
 
 
 class Char(bytes, BaseType):
@@ -17,8 +17,8 @@ def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] = Non
         if count == 0:
             return b""
 
-        data = stream.read(count)
-        if len(data) != count:
+        data = stream.read(-1 if count == EOF else count)
+        if count != EOF and len(data) != count:
             raise EOFError(f"Read {len(data)} bytes, but expected {count}")
 
         return type.__call__(cls, data)

dissect/cstruct/types/packed.py

@@ -4,7 +4,7 @@
 from struct import Struct
 from typing import Any, BinaryIO
 
-from dissect.cstruct.types.base import BaseType
+from dissect.cstruct.types.base import EOF, BaseType
 
 
 @lru_cache(1024)
@@ -23,8 +23,14 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] = None) -> Packed:
 
     @classmethod
     def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] = None) -> list[Packed]:
-        length = cls.size * count
-        data = stream.read(length)
+        if count == EOF:
+            data = stream.read()
+            length = len(data)
+            count = length // cls.size
+        else:
+            length = cls.size * count
+            data = stream.read(length)
+
         fmt = _struct(cls.cs.endian, f"{count}{cls.packchar}")
 
         if len(data) != length:

dissect/cstruct/types/wchar.py

@@ -3,7 +3,7 @@
 import sys
 from typing import Any, BinaryIO
 
-from dissect.cstruct.types.base import ArrayMetaType, BaseType
+from dissect.cstruct.types.base import EOF, ArrayMetaType, BaseType
 
 
 class Wchar(str, BaseType):
@@ -23,12 +23,14 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] = None) -> Wchar:
 
     @classmethod
     def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] = None) -> Wchar:
-        count *= 2
         if count == 0:
             return ""
 
-        data = stream.read(count)
-        if len(data) != count:
+        if count != EOF:
+            count *= 2
+
+        data = stream.read(-1 if count == EOF else count)
+        if count != EOF and len(data) != count:
             raise EOFError(f"Read {len(data)} bytes, but expected {count}")
 
         return type.__call__(cls, data.decode(cls.__encoding_map__[cls.cs.endian]))

tests/test_types_base.py

@@ -3,9 +3,78 @@
 from dissect.cstruct.cstruct import cstruct
 from dissect.cstruct.exceptions import ArraySizeError
 
+from .utils import verify_compiled
+
 
 def test_array_size_mismatch(cs: cstruct):
     with pytest.raises(ArraySizeError):
         cs.uint8[2]([1, 2, 3]).dumps()
 
     assert cs.uint8[2]([1, 2]).dumps()
+
+
+def test_eof(cs: cstruct, compiled: bool):
+    cdef = """
+    struct test_char {
+        char data[EOF];
+    };
+
+    struct test_wchar {
+        wchar data[EOF];
+    };
+
+    struct test_packed {
+        uint16 data[EOF];
+    };
+
+    struct test_int {
+        uint24 data[EOF];
+    };
+
+    enum Test : uint16 {
+        A = 1
+    };
+
+    struct test_enum {
+        Test data[EOF];
+    };
+
+    struct test_eof_field {
+        uint8 EOF;
+        char data[EOF];
+        uint8 remainder;
+    };
+    """
+    cs.load(cdef, compiled=compiled)
+
+    assert verify_compiled(cs.test_char, compiled)
+    assert verify_compiled(cs.test_wchar, compiled)
+    assert verify_compiled(cs.test_packed, compiled)
+    assert verify_compiled(cs.test_int, compiled)
+    assert verify_compiled(cs.test_enum, compiled)
+    assert verify_compiled(cs.test_eof_field, compiled)
+
+    test_char = cs.test_char(b"abc")
+    assert test_char.data == b"abc"
+    assert test_char.dumps() == b"abc"
+
+    test_wchar = cs.test_wchar("abc".encode("utf-16-le"))
+    assert test_wchar.data == "abc"
+    assert test_wchar.dumps() == "abc".encode("utf-16-le")
+
+    test_packed = cs.test_packed(b"\x01\x00\x02\x00")
+    assert test_packed.data == [1, 2]
+    assert test_packed.dumps() == b"\x01\x00\x02\x00"
+
+    test_int = cs.test_int(b"\x01\x00\x00\x02\x00\x00")
+    assert test_int.data == [1, 2]
+    assert test_int.dumps() == b"\x01\x00\x00\x02\x00\x00"
+
+    test_enum = cs.test_enum(b"\x01\x00")
+    assert test_enum.data == [cs.Test.A]
+    assert test_enum.dumps() == b"\x01\x00"
+
+    test_eof_field = cs.test_eof_field(b"\x01a\x02")
+    assert test_eof_field.data == b"a"
+    assert test_eof_field.remainder == 2
+    assert test_eof_field.dumps() == b"\x01a\x02"