Pad array writes to the correct number of entries

Author: Schamper

dissect/cstruct/types/base.py

@@ -273,8 +273,10 @@ def _write(cls, stream: BinaryIO, data: list[Any]) -> int:
         if cls.null_terminated:
             return cls.type._write_0(stream, data)
 
-        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 += [cls.type.__default__()] * remaining
 
         return cls.type._write_array(stream, data)
 

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:
@@ -20,15 +21,21 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:
         return type.__call__(cls, super()._read(stream, context))
 
     @classmethod
-    def _write(cls, stream: BinaryIO, data: bytes) -> int:
+    def _write(cls, stream: BinaryIO, data: bytes | str | list[int]) -> int:
         if isinstance(data, list) and data and isinstance(data[0], int):
             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)
 
 

dissect/cstruct/types/structure.py

@@ -259,7 +259,7 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:
             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)
 
@@ -325,7 +325,7 @@ def _write(cls, stream: BinaryIO, data: Structure) -> 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
 

dissect/cstruct/types/wchar.py

@@ -3,6 +3,7 @@
 import sys
 from typing import Any, BinaryIO, ClassVar
 
+from dissect.cstruct.exception import ArraySizeError
 from dissect.cstruct.types.base import EOF, BaseArray, BaseType
 
 
@@ -23,6 +24,12 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Wchar
     def _write(cls, stream: BinaryIO, data: str) -> 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__[cls.cs.endian]))
 
 

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])
+    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)])
+    assert buf.getvalue() == b"\x01\x02" + b"\x00\x00" * 3
+
+
 def test_eof(cs: cstruct, compiled: bool) -> None:
     cdef = """
     struct test_char {

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")
+    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")
+
+
 def test_char_eof(cs: cstruct) -> None:
     with pytest.raises(EOFError):
         cs.char(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")
+    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")
+
+
 def test_wchar_be_read(cs: cstruct) -> None:
     cs.endian = ">"