Repository: d0c-s4vage/pfp
Pull request #94 (2019-11-11 15:07:24)

Fixes forward declared struct declarations


Patch:
diff --git a/pfp/fields.py b/pfp/fields.py
index 6ad101a..5a6add1 100644
--- a/pfp/fields.py
+++ b/pfp/fields.py
@@ -33,6 +33,8 @@ def false():
 
 def get_value(field):
     if isinstance(field, Field):
+        if isinstance(field, Array):
+            return field._array_to_str()
         return field._pfp__value
     else:
         return field
@@ -2028,7 +2030,7 @@ class Array(Field):
         return res
 
     def __eq__(self, other):
-        if self.is_stringable() and other.__class__ in [String, WString, str]:
+        if self.is_stringable() and other.__class__ in [String, WString, str, bytes]:
             res = self._array_to_str()
             return utils.binary(res) == utils.binary(PYSTR(other))
         else:
diff --git a/pfp/interp.py b/pfp/interp.py
index 07697b8..c483c1a 100644
--- a/pfp/interp.py
+++ b/pfp/interp.py
@@ -1047,18 +1047,20 @@ class PfpInterp(object):
 
         # one pass to define all functions. Functions may only live at the
         # top-level (functions may not be nested or contained within structs,
-        # if/else statements, or other code block types).
+        # if/else statements, or other code block types). aka hoisting
         for child in children:
             if type(child) is tuple:
                 child = child[1]
-            if not isinstance(child, (AST.FuncDef, AST.Typedef)):
+            if not isinstance(child, (AST.FuncDef, AST.Typedef)) \
+                    and not  is_forward_declared_struct(child):
                 continue
             self._handle_node(child, scope, ctxt, stream)
 
         for child in children:
             if type(child) is tuple:
                 child = child[1]
-            if isinstance(child, (AST.FuncDef, AST.Typedef)):
+            if isinstance(child, (AST.FuncDef, AST.Typedef)) or \
+                    is_forward_declared_struct(child):
                 continue
             self._handle_node(child, scope, ctxt, stream)
 
@@ -1151,6 +1153,7 @@ class PfpInterp(object):
         field = self._handle_node(node.type, scope, ctxt, stream)
         bitsize = None
         bitfield_rw = None
+
         if getattr(node, "bitsize", None) is not None:
             bitsize = self._handle_node(node.bitsize, scope, ctxt, stream)
             has_prev = len(ctxt._pfp__children) > 0
@@ -1181,7 +1184,10 @@ class PfpInterp(object):
                 bitfield_rw = fields.BitfieldRW(self, field)
                 bitfield_rw.reserve_bits(bitsize, stream)
 
-        if getattr(node, "is_func_param", False):
+        if is_forward_declared_struct(node):
+            scope.add_type_class(node.type.name, field)
+
+        elif getattr(node, "is_func_param", False):
             # we want to keep this as a class and not instantiate it
             # instantiation will be done in functions.ParamListDef.instantiate
             field = (field_name, field)
@@ -1551,19 +1557,19 @@ class PfpInterp(object):
             for param in node.args.params:
                 param.is_func_param = True
 
-        # it's actually being defined
         if node.decls is not None:
             struct_cls = StructUnionDef("struct", self, node)
-
             if node.name is not None:
                 scope.add_type_class(node.name, struct_cls)
-
             return struct_cls
 
         # it's declaring a struct field. E.g.
         #    struct IFD subDir;
         else:
-            return scope.get_type(node.name)
+            res = scope.get_type(node.name)
+            if res is None:
+                res = StructUnionDef(node.name, self, node)
+            return res
 
     def _handle_struct_decls(self, node, scope, ctxt, stream):
         self._dlog("handling struct decls")
@@ -2565,3 +2571,12 @@ class PfpInterp(object):
 
         cls = getattr(fields, res)
         return cls
+
+def is_forward_declared_struct(node):
+    return (
+        isinstance(node, AST.Decl)
+        and node.init is None
+        and isinstance(node.type, AST.Struct)
+        and node.type.decls is None
+    )
+


Test patch:
diff --git a/tests/test_compat.py b/tests/test_compat.py
index e4260f4..b02221f 100644
--- a/tests/test_compat.py
+++ b/tests/test_compat.py
@@ -132,6 +132,7 @@ class TestCompatIO(utils.PfpTestCase):
         dom = self._test_parse_build(
             "\x80\x01",
             """
+                BigEndian();
                 local ushort blah = ReadUShort();
                 Printf("%d|", blah);
                 Printf("%d", FTell());
diff --git a/tests/test_control_flow.py b/tests/test_control_flow.py
index 89b5dab..99249eb 100644
--- a/tests/test_control_flow.py
+++ b/tests/test_control_flow.py
@@ -223,6 +223,7 @@ class TestControlFlow(utils.PfpTestCase):
         dom = self._test_parse_build(
             "\x00\x00\x00\x01\x00\x00\x00\x02",
             """
+                BigEndian();
                 local int a = 5;
                 switch(a) {
                     case 0:
diff --git a/tests/test_enums.py b/tests/test_enums.py
index 61961de..4b13ba5 100644
--- a/tests/test_enums.py
+++ b/tests/test_enums.py
@@ -27,6 +27,7 @@ class TestEnums(utils.PfpTestCase):
         dom = self._test_parse_build(
             "\x00\x00\x00\x01",
             """
+                BigEndian();
                 enum TEST_ENUM {
                     BLAH1,
                     BLAH2,
@@ -158,6 +159,7 @@ class TestEnums(utils.PfpTestCase):
         dom = self._test_parse_build(
             "\x31",
             """
+                BigEndian();
                 enum <uchar> BLAHS {
                     BLAH1,
                     BLAH2,
@@ -198,6 +200,7 @@ class TestEnums(utils.PfpTestCase):
         dom = self._test_parse_build(
             "\x31",
             """
+                BigEndian();
                 enum <uchar> {
                     BLAH1,
                     BLAH2,
diff --git a/tests/test_metadata.py b/tests/test_metadata.py
index 45f8fa6..dd2bd78 100644
--- a/tests/test_metadata.py
+++ b/tests/test_metadata.py
@@ -145,6 +145,7 @@ class TestMetadata(utils.PfpTestCase):
         dom = self._test_parse_build(
             "TYPA\x41\x410000TYPB\x42\x420000",
             """
+                BigEndian();
                 typedef struct {
                     uchar a;
                     uchar b;
@@ -187,6 +188,7 @@ class TestMetadata(utils.PfpTestCase):
         dom = self._test_parse_build(
             "yoyoyo\x00\x10x\x9cc```d```\x02\x00\x00\x0f\x00\x04",
             """
+                BigEndian();
                 typedef struct {
                     int a;
                     int b;
diff --git a/tests/test_struct_union.py b/tests/test_struct_union.py
index fa6839d..7452e64 100644
--- a/tests/test_struct_union.py
+++ b/tests/test_struct_union.py
@@ -21,6 +21,24 @@ class TestStructUnion(utils.PfpTestCase):
 
     def tearDown(self):
         pass
+    
+    def test_basic_struct(self):
+        dom = self._test_parse_build(
+            "\x00\x01\x02\x03",
+            """
+                struct SomeStruct {
+                    byte a;
+                    byte b;
+                    byte c;
+                    byte d;
+                };
+                struct SomeStruct test;
+            """
+        )
+        assert dom.test.a == 0
+        assert dom.test.b == 1
+        assert dom.test.c == 2
+        assert dom.test.d == 3
 
     def test_field_path(self):
         dom = self._test_parse_build(
@@ -62,6 +80,25 @@ class TestStructUnion(utils.PfpTestCase):
                 blah some_struct2;
             """,
         )
+    
+    def test_forward_declared_struct(self):
+        dom = self._test_parse_build(
+            "\x00\x01",
+            """
+                struct fp16;
+                void ReadFP16(fp16& f)
+                {
+                    Printf("%d", f.value);
+                }
+                typedef struct
+                {
+                    int16 value;
+                } fp16;
+                fp16 preferred_volume;
+                ReadFP16(preferred_volume);
+            """,
+            stdout="256",
+        )
 
     def test_struct_with_parameters(self):
         dom = self._test_parse_build(
diff --git a/tests/utils.py b/tests/utils.py
index 673313b..4e8255e 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -69,6 +69,9 @@ class PfpTestCase(unittest.TestCase):
         if _stream:
             data = six.StringIO(data)
 
+        # defaults to LittleEndian
+        template = "LittleEndian();" + template
+
         dom = pfp.parse(
             data, template, debug=debug, predefines=predefines, printf=printf
         )