Pass context when reading a compiled Structure or Array (#27)

fox-it · Feb 27, 2023 · 4d69954 · 4d69954
1 parent f57fe5c
commit 4d69954
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 3 deletions.
diff --git a/dissect/cstruct/compiler.py b/dissect/cstruct/compiler.py
@@ -135,22 +135,22 @@ def gen_struct_class(self, name: str, structure: Structure) -> str:
                         f"""
                         r['{field.name}'] = []
                         for _ in range({num}):
-                            r['{field.name}'].append(self.lookup['{field.name}'].type.type._read(stream))
+                            r['{field.name}'].append(self.lookup['{field.name}'].type.type._read(stream, context=r))
                         sizes['{field.name}'] = stream.tell() - s
                         """
                     )
                 elif isinstance(field_type, Structure) and field_type.anonymous:
                     struct_read += dedent(
                         f"""
-                        v = self.lookup["{field.name}"].type._read(stream)
+                        v = self.lookup["{field.name}"].type._read(stream, context=r)
                         r.update(v._values)
                         sizes.update(v._sizes)
                         """
                     )
                 else:
                     struct_read += dedent(
                         f"""
-                        r['{field.name}'] = self.lookup['{field.name}'].type._read(stream)
+                        r['{field.name}'] = self.lookup['{field.name}'].type._read(stream, context=r)
                         sizes['{field.name}'] = stream.tell() - s
                         """
                     )

diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -502,6 +502,37 @@ def test_array_three_dimensional(compiled):
     assert obj.dumps() == buf
 
 
+@pytest.mark.parametrize("compiled", [True, False])
+def test_nested_array_of_variable_size(compiled: bool):
+    cdef = """
+    struct test {
+        uint8   outer;
+        uint8   medior;
+        uint8   inner;
+        uint8   a[outer][medior][inner];
+    }
+    """
+    cs = cstruct.cstruct(endian="<")
+    cs.load(cdef, compiled=compiled)
+
+    assert verify_compiled(cs.test, compiled)
+
+    buf = b"\x02\x01\x03\x01\x02\x03\x04\x05\x06"
+    obj = cs.test(buf)
+
+    assert obj.outer == 2
+    assert obj.medior == 1
+    assert obj.inner == 3
+    assert obj.a[0][0][0] == 1
+    assert obj.a[0][0][1] == 2
+    assert obj.a[0][0][2] == 3
+    assert obj.a[1][0][0] == 4
+    assert obj.a[1][0][1] == 5
+    assert obj.a[1][0][2] == 6
+
+    assert obj.dumps() == buf
+
+
 def test_report_array_size_mismatch():
     cdef = """
     struct test {