util: ctypes: handle vla_counted_by

Handle automatically constructing VLA structs with the `vla_counted_by`
attribute.

Signed-off-by: Jordan Yates <jordan@embeint.com>

Author: JordanYates

src/infuse_iot/util/ctypes.py

@@ -43,23 +43,46 @@ class property is not `None`, it will consume the remainder of
         """
 
         base = cls.from_buffer_copy(source, offset)
+        vla_val: list | VLACompatLittleEndianStruct
         if cls.vla_field is None:
             return base
 
         remainder = source[ctypes.sizeof(cls) :]
         vla_field_name, vla_field_type = cls.vla_field  # type: ignore
 
         if issubclass(vla_field_type, ctypes.Array):
-            array_base: ctypes._CData = vla_field_type._type_  # type: ignore
+            array_base: ctypes._PyCSimpleType = vla_field_type._type_  # type: ignore
+            if hasattr(array_base, "vla_counted_by"):
+                # This is an array of VLA arrays where the sub-arrys define their own length
+                vla_val = []
+                # Consume all remaining buffer bytes
+                while len(remainder) > 0:
+                    sub_vla_field_name, sub_vla_field_type = array_base.vla_field  # type: ignore
+                    sub_array_base: ctypes._CData = sub_vla_field_type._type_  # type: ignore
+                    sub_base = array_base.from_buffer_copy(remainder)
+                    sub_base_size = ctypes.sizeof(sub_base)
+                    sub_count = getattr(sub_base, array_base.vla_counted_by)
+                    if sub_count < 0:
+                        # Assume that negative length is an error code and use 0
+                        vla_val.append(sub_base)
+                    else:
+                        sub_vla_type = sub_count * sub_array_base
+                        # Don't use ctypes.sizeof on constructed type, it returns the wrong value
+                        sub_vla_size = sub_count * ctypes.sizeof(sub_array_base)
+                        sub_vla_val = sub_vla_type.from_buffer_copy(remainder[sub_base_size:])
+                        setattr(sub_base, sub_vla_field_name, sub_vla_val)
+                        vla_val.append(sub_base)
+                    remainder = remainder[sub_base_size + sub_vla_size :]
+            else:
+                # Determine the number of VLA elements on "source"
+                vla_byte_len = (len(source) - offset) - ctypes.sizeof(cls)
+                vla_element_size = ctypes.sizeof(array_base)
+                if vla_byte_len % vla_element_size != 0:
+                    raise TypeError(f"Unaligned VLA buffer for {cls} (len {len(source)})")
+                vla_num = vla_byte_len // vla_element_size
+                vla_type = vla_num * array_base
+                vla_val = vla_type.from_buffer_copy(remainder)
 
-            # Determine the number of VLA elements on "source"
-            vla_byte_len = (len(source) - offset) - ctypes.sizeof(cls)
-            vla_element_size = ctypes.sizeof(array_base)
-            if vla_byte_len % vla_element_size != 0:
-                raise TypeError(f"Unaligned VLA buffer for {cls} (len {len(source)})")
-            vla_num = vla_byte_len // vla_element_size
-            vla_type = vla_num * array_base
-            vla_val = vla_type.from_buffer_copy(remainder)
         elif issubclass(vla_field_type, VLACompatLittleEndianStruct):
             vla_val = vla_field_type.vla_from_buffer_copy(remainder)
         else:

tests/util/test_ctypes.py

@@ -11,20 +11,40 @@ class VLABase(VLACompatLittleEndianStruct):
         ("first", ctypes.c_uint32),
     ]
     vla_field = ("vla", 0 * ctypes.c_uint32)
+    _pack_ = 1
+
+
+class VLACountedBy(VLACompatLittleEndianStruct):
+    _fields_ = [
+        ("count", ctypes.c_uint8),
+    ]
+    vla_field = ("data", 0 * ctypes.c_uint8)
+    vla_counted_by = "count"
+    _pack_ = 1
 
 
 class VLANested(VLACompatLittleEndianStruct):
     _fields_ = [
         ("first", ctypes.c_uint32),
     ]
     vla_field = ("vla", VLABase)
+    _pack_ = 1
+
+
+class VLADualNested(VLACompatLittleEndianStruct):
+    _fields_ = [
+        ("num", ctypes.c_uint8),
+    ]
+    vla_field = ("values", 0 * VLACountedBy)
+    _pack_ = 1
 
 
 class VLANone(VLACompatLittleEndianStruct):
     _fields_ = [
         ("first", ctypes.c_uint32),
         ("second", ctypes.c_uint32),
     ]
+    _pack_ = 1
 
 
 def test_vla_compat_struct():
@@ -47,6 +67,25 @@ def test_vla_compat_struct():
     assert none.first == 2
     assert none.second == 3
 
+    one_element = b"\x01\x01\xff"
+    two_element = b"\x02\x01\x33\x02\xaa\x55"
+    dual_nested_one = VLADualNested.vla_from_buffer_copy(one_element)
+    assert dual_nested_one.num == 1
+    assert len(dual_nested_one.values) == 1
+    assert dual_nested_one.values[0].count == 1
+    assert len(dual_nested_one.values[0].data) == 1
+    assert dual_nested_one.values[0].data[0] == 0xFF
+    dual_nested_two = VLADualNested.vla_from_buffer_copy(two_element)
+    assert dual_nested_two.num == 2
+    assert len(dual_nested_two.values) == 2
+    assert dual_nested_two.values[0].count == 1
+    assert len(dual_nested_one.values[0].data) == 1
+    assert dual_nested_two.values[0].data[0] == 0x33
+    assert dual_nested_two.values[1].count == 2
+    assert len(dual_nested_two.values[1].data) == 2
+    assert dual_nested_two.values[1].data[0] == 0xAA
+    assert dual_nested_two.values[1].data[1] == 0x55
+
     unaligned = b"\x00" * 31
     try:
         VLABase.vla_from_buffer_copy(unaligned)