Skip to content

Commit

Permalink
Add support for EOF length identifier
Browse files Browse the repository at this point in the history
  • Loading branch information
Schamper committed Aug 2, 2023
1 parent 5458791 commit bf7bd5c
Show file tree
Hide file tree
Showing 5 changed files with 111 additions and 13 deletions.
27 changes: 24 additions & 3 deletions dissect/cstruct/types/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Check warning on line 10 in dissect/cstruct/types/base.py

View check run for this annotation

Codecov / codecov/patch

dissect/cstruct/types/base.py#L9-L10

Added lines #L9 - L10 were not covered by tests


EOF = -0xE0F # Negative counts are illegal anyway, so abuse that for our EOF sentinel


class MetaType(type):
Expand Down Expand Up @@ -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]:
Expand Down Expand Up @@ -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

Check warning on line 188 in dissect/cstruct/types/base.py

View check run for this annotation

Codecov / codecov/patch

dissect/cstruct/types/base.py#L188

Added line #L188 was not covered by tests
num = EOF
else:
num = max(0, cls.num_entries)

return cls.type._read_array(stream, num, context)


class Array(list, BaseType, metaclass=ArrayMetaType):
Expand Down
6 changes: 3 additions & 3 deletions dissect/cstruct/types/char.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -17,8 +17,8 @@ def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] = Non
if count == 0:
return b""

Check warning on line 18 in dissect/cstruct/types/char.py

View check run for this annotation

Codecov / codecov/patch

dissect/cstruct/types/char.py#L18

Added line #L18 was not covered by tests

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)
Expand Down
12 changes: 9 additions & 3 deletions dissect/cstruct/types/packed.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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:
Expand Down
10 changes: 6 additions & 4 deletions dissect/cstruct/types/wchar.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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]))
Expand Down
69 changes: 69 additions & 0 deletions tests/test_types_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

0 comments on commit bf7bd5c

Please sign in to comment.