From ef710a2a596e488f018797deb5be7a4749275bb3 Mon Sep 17 00:00:00 2001
From: MrYoranimo <1333535+MrYoranimo@users.noreply.github.com>
Date: Sun, 26 Feb 2023 00:52:34 +0100
Subject: [PATCH] Pass context when reading a compiled Structure or Array

After allowing nested arrays in previous commits, only the parent
type would have access to the context while resolving expressions.

This changeset updates the compiled so that other fields' values
are passed as context so that sublevel arrays can also resolve
their size from an expression correct. A test has been written
to verify this change.
---
 dissect/cstruct/compiler.py |  6 +++---
 tests/test_basic.py         | 31 +++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/dissect/cstruct/compiler.py b/dissect/cstruct/compiler.py
index c972fae..296d944 100644
--- a/dissect/cstruct/compiler.py
+++ b/dissect/cstruct/compiler.py
@@ -135,14 +135,14 @@ 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)
                         """
@@ -150,7 +150,7 @@ def gen_struct_class(self, name: str, structure: Structure) -> str:
                 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
index 7a50425..3b8b9ea 100644
--- 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 {