Pad array writes to the correct number of entries

Author: Schamper

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)
 

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)
 
 

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
 

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]))
 
 

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 {

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"")

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 = ">"