Allow dynamically sized unions

Author: Schamper

dissect/cstruct/types/structure.py

@@ -373,25 +373,35 @@ def _calculate_size_and_offsets(cls, fields: list[Field], align: bool = False) -
                 # If a field already has an alignment, it's leading
                 field.alignment = field.type.alignment
 
-            size = max(len(field.type), size)
+            if size is not None:
+                try:
+                    size = max(len(field.type), size)
+                except TypeError:
+                    size = None
+
             alignment = max(field.alignment, alignment)
 
         return size, alignment
 
     def _read(cls, stream: BinaryIO, context: dict[str, Any] = None) -> Union:
-        buf = io.BytesIO(memoryview(stream.read(len(cls))))
         result = {}
         sizes = {}
 
+        if cls.size is None:
+            offset = stream.tell()
+            buf = stream
+        else:
+            offset = 0
+            buf = io.BytesIO(stream.read(cls.size))
+
         for field in cls.fields:
-            start = 0
-            buf.seek(0)
             field_type = cls.cs.resolve(field.type)
 
+            start = 0
             if field.offset is not None:
-                buf.seek(field.offset)
                 start = field.offset
 
+            buf.seek(offset + start)
             value = field_type._read(buf, result)
 
             if isinstance(field_type, StructureMetaType) and field_type.anonymous:

tests/test_types_union.py

@@ -187,7 +187,7 @@ def test_union_definition(cs: cstruct):
     assert cs.test().dumps() == b"\x00\x00\x00\x00\x00\x00\x00\x00"
 
 
-def test_union_definitiont_nested(cs: cstruct, compiled: bool):
+def test_union_definition_nested(cs: cstruct, compiled: bool):
     cdef = """
     struct test {
         char magic[4];
@@ -247,3 +247,25 @@ def test_union_definition_anonymous(cs: cstruct, compiled: bool):
     assert obj.c == b"\x02"
     assert obj.d == 0x04040303
     assert obj.dumps() == buf
+
+
+def test_union_definition_dynamic(cs: cstruct):
+    cdef = """
+    struct dynamic {
+        uint8   size;
+        char    data[size];
+    };
+
+    union test {
+        dynamic a;
+        uint64  b;
+    };
+    """
+    cs.load(cdef, compiled=False)
+
+    buf = b"\x09aaaaaaaaa"
+    obj = cs.test(buf)
+
+    assert obj.a.size == 9
+    assert obj.a.data == b"aaaaaaaaa"
+    assert obj.b == 0x6161616161616109