From 29ceaaafd040b023778c308749c23e5f658f7729 Mon Sep 17 00:00:00 2001 From: Schamper <1254028+Schamper@users.noreply.github.com> Date: Wed, 15 Jul 2026 19:08:20 +0200 Subject: [PATCH] Pad array writes to the correct number of entries --- dissect/cstruct/types/base.py | 6 ++++-- dissect/cstruct/types/char.py | 14 ++++++++---- dissect/cstruct/types/structure.py | 4 ++-- dissect/cstruct/types/wchar.py | 7 ++++++ tests/test_types_base.py | 19 +++++++++++++++++ tests/test_types_char.py | 34 ++++++++++++++++++++++++++++++ tests/test_types_wchar.py | 29 +++++++++++++++++++++++++ 7 files changed, 105 insertions(+), 8 deletions(-) diff --git a/dissect/cstruct/types/base.py b/dissect/cstruct/types/base.py index 7bc2aded..1c306f4a 100644 --- a/dissect/cstruct/types/base.py +++ b/dissect/cstruct/types/base.py @@ -294,8 +294,10 @@ def _write(cls, stream: BinaryIO, data: list[Any], *, endian: Endianness) -> int if cls.null_terminated: return cls.type._write_0(stream, data, endian=endian) - if not cls.dynamic and cls.num_entries != (actual_size := len(data)): - raise ArraySizeError(f"Expected static array size {cls.num_entries}, got {actual_size} instead.") + if not cls.dynamic and (remaining := cls.num_entries - (actual_size := len(data))): + if remaining < 0: + raise ArraySizeError(f"Expected static array size {cls.num_entries}, got {actual_size} instead.") + data = data + [cls.type.__default__()] * remaining return cls.type._write_array(stream, data, endian=endian) diff --git a/dissect/cstruct/types/char.py b/dissect/cstruct/types/char.py index 3f723978..064c993d 100644 --- a/dissect/cstruct/types/char.py +++ b/dissect/cstruct/types/char.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, BinaryIO +from dissect.cstruct.exception import ArraySizeError from dissect.cstruct.types.base import EOF, BaseArray, BaseType if TYPE_CHECKING: @@ -22,15 +23,20 @@ def _read(cls, stream: BinaryIO, *, context: dict[str, Any] | None = None, endia return type.__call__(cls, super()._read(stream, context=context, endian=endian)) @classmethod - def _write(cls, stream: BinaryIO, data: bytes, *, endian: Endianness) -> int: - if isinstance(data, list) and data and isinstance(data[0], int): + def _write(cls, stream: BinaryIO, data: bytes | str | list[int], *, endian: Endianness) -> int: + if isinstance(data, list): data = bytes(data) - elif isinstance(data, str): data = data.encode("latin-1") if cls.null_terminated: - return stream.write(data + b"\x00") + data += b"\x00" + + if not cls.dynamic and (remaining := cls.num_entries - (actual_size := len(data))): + if remaining < 0: + raise ArraySizeError(f"Expected static array size {cls.num_entries}, got {actual_size} instead.") + data += b"\x00" * remaining + return stream.write(data) diff --git a/dissect/cstruct/types/structure.py b/dissect/cstruct/types/structure.py index e8b13425..6974a233 100644 --- a/dissect/cstruct/types/structure.py +++ b/dissect/cstruct/types/structure.py @@ -252,7 +252,7 @@ def _read(cls, stream: BinaryIO, *, context: dict[str, Any] | None = None, endia offset = stream.tell() if field.offset is not None and offset != struct_start + field.offset: - # Field is at a specific offset, either alligned or added that way + # Field is at a specific offset, either aligned or added that way offset = struct_start + field.offset stream.seek(offset) @@ -318,7 +318,7 @@ def _write(cls, stream: BinaryIO, data: Structure, *, endian: Endianness) -> int offset = stream.tell() if field.offset is not None and offset < struct_start + field.offset: - # Field is at a specific offset, either alligned or added that way + # Field is at a specific offset, either aligned or added that way stream.write(b"\x00" * (struct_start + field.offset - offset)) offset = struct_start + field.offset diff --git a/dissect/cstruct/types/wchar.py b/dissect/cstruct/types/wchar.py index 3c911202..fce2d2d2 100644 --- a/dissect/cstruct/types/wchar.py +++ b/dissect/cstruct/types/wchar.py @@ -3,6 +3,7 @@ import sys from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar +from dissect.cstruct.exception import ArraySizeError from dissect.cstruct.types.base import EOF, BaseArray, BaseType if TYPE_CHECKING: @@ -26,6 +27,12 @@ def _read(cls, stream: BinaryIO, *, context: dict[str, Any] | None = None, endia def _write(cls, stream: BinaryIO, data: str, *, endian: Endianness) -> int: if cls.null_terminated: data += "\x00" + + if not cls.dynamic and (remaining := cls.num_entries - (actual_size := len(data))): + if remaining < 0: + raise ArraySizeError(f"Expected static array size {cls.num_entries}, got {actual_size} instead.") + data += "\x00" * remaining + return stream.write(data.encode(Wchar.__encoding_map__[endian])) diff --git a/tests/test_types_base.py b/tests/test_types_base.py index ee124e7d..29c7af21 100644 --- a/tests/test_types_base.py +++ b/tests/test_types_base.py @@ -1,5 +1,6 @@ from __future__ import annotations +import io from typing import TYPE_CHECKING, BinaryIO import pytest @@ -20,6 +21,24 @@ def test_array_size_mismatch(cs: cstruct) -> None: assert cs.uint8[2]([1, 2]).dumps() +def test_array_write_padding(cs: cstruct) -> None: + buf = io.BytesIO() + cs.uint8[4]._write(buf, [1, 2], endian=cs.endian) + assert buf.getvalue() == b"\x01\x02\x00\x00" + + cdef = """ + struct point { + uint8 x; + uint8 y; + }; + """ + cs.load(cdef) + + buf = io.BytesIO() + cs.point[4]._write(buf, [cs.point(x=1, y=2)], endian=cs.endian) + assert buf.getvalue() == b"\x01\x02" + b"\x00\x00" * 3 + + def test_eof(cs: cstruct, compiled: bool) -> None: cdef = """ struct test_char { diff --git a/tests/test_types_char.py b/tests/test_types_char.py index 5eff05ad..378f8313 100644 --- a/tests/test_types_char.py +++ b/tests/test_types_char.py @@ -5,6 +5,8 @@ import pytest +from dissect.cstruct.exception import ArraySizeError + if TYPE_CHECKING: from dissect.cstruct.cstruct import cstruct @@ -36,6 +38,38 @@ def test_char_array_write(cs: cstruct) -> None: assert cs.char[None](buf).dumps() == b"AAAA\x00" +def test_char_array_write_padding(cs: cstruct) -> None: + buf = io.BytesIO() + cs.char[8]._write(buf, "hi", endian=cs.endian) + assert buf.getvalue() == b"hi\x00\x00\x00\x00\x00\x00" + + cdef = """ + struct test_struct { + int x; + char y[8]; + char z[16]; + }; + """ + cs.load(cdef) + + obj = cs.test_struct() + obj.x = 4 + obj.y = "hi" + obj.z = "bye" + + assert len(obj.dumps()) == len(cs.test_struct) + assert ( + obj.dumps() + == b"\x04\x00\x00\x00hi\x00\x00\x00\x00\x00\x00bye\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + ) + + +def test_char_array_write_size_error(cs: cstruct) -> None: + buf = io.BytesIO() + with pytest.raises(ArraySizeError): + cs.char[4]._write(buf, b"toolong", endian=cs.endian) + + def test_char_eof(cs: cstruct) -> None: with pytest.raises(EOFError): cs.char(b"") diff --git a/tests/test_types_wchar.py b/tests/test_types_wchar.py index 69fb861d..6b52b7ab 100644 --- a/tests/test_types_wchar.py +++ b/tests/test_types_wchar.py @@ -5,6 +5,8 @@ import pytest +from dissect.cstruct.exception import ArraySizeError + if TYPE_CHECKING: from dissect.cstruct.cstruct import cstruct @@ -38,6 +40,33 @@ def test_wchar_array_write(cs: cstruct) -> None: assert cs.wchar[None](buf).dumps() == b"A\x00A\x00A\x00A\x00\x00\x00" +def test_wchar_array_write_padding(cs: cstruct) -> None: + buf = io.BytesIO() + cs.wchar[8]._write(buf, "hi", endian=cs.endian) + assert buf.getvalue() == b"h\x00i\x00" + b"\x00" * 12 + + cdef = """ + struct test_struct { + int x; + wchar y[8]; + wchar z[16]; + }; + """ + cs.load(cdef) + obj = cs.test_struct() + obj.x = 4 + obj.y = "hi" + obj.z = "bye" + assert len(obj.dumps()) == len(cs.test_struct) + assert obj.dumps() == b"\x04\x00\x00\x00h\x00i\x00" + (b"\x00" * 12) + b"b\x00y\x00e\x00" + (b"\x00" * 26) + + +def test_wchar_array_write_size_error(cs: cstruct) -> None: + buf = io.BytesIO() + with pytest.raises(ArraySizeError): + cs.wchar[4]._write(buf, "toolong", endian=cs.endian) + + def test_wchar_be_read(cs: cstruct) -> None: cs.endian = ">"