Support struct setting and getting in angr (callbacks wip) (#173)

* Support struct setting and getting in angr (callbacks wip)

* working refresh

Author: mahaloz

libbs/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.2.0"
+__version__ = "3.3.0"
 
 
 import logging

libbs/decompilers/angr/interface.py

@@ -12,7 +12,8 @@
     DecompilerInterface,
 )
 from libbs.artifacts import (
-    Function, FunctionHeader, Comment, StackVariable, FunctionArgument, Artifact, Decompilation, Context
+    Function, FunctionHeader, Comment, StackVariable, FunctionArgument, Artifact, Decompilation, Context,
+    Struct, StructMember
 )
 from .artifact_lifter import AngrArtifactLifter
 
@@ -256,8 +257,9 @@ def _set_function(self, func: Function, **kwargs) -> bool:
         # re-decompile a function if needed
         decompilation = self.decompile_function(angr_func).codegen
         changes = super()._set_function(func, decompilation=decompilation, **kwargs)
-        if not self.headless:
-            self.refresh_decompilation(func.addr)
+        if not self.headless and changes:
+            # Use "retype_variable" event to trigger proper UI refresh including type reflow
+            self.refresh_decompilation(func.addr, event="retype_variable")
 
         return changes
 
@@ -348,12 +350,34 @@ def _set_stack_variable(self, svar: StackVariable, decompilation=None, **kwargs)
             return changed
 
         dec_svar = AngrInterface.find_stack_var_in_codegen(decompilation, svar.offset)
-        if dec_svar and svar.name and svar.name != dec_svar.name:
-            # TODO: set the types of the stack vars
+        if not dec_svar:
+            return changed
+
+        # Set the name if provided and different
+        if svar.name and svar.name != dec_svar.name:
             dec_svar.name = svar.name
             dec_svar.renamed = True
             changed = True
 
+        # Set the type if provided
+        if svar.type:
+            try:
+                from angr.sim_type import parse_type
+                types_store = self.main_instance.project.kb.types
+                arch = self.main_instance.project.arch
+
+                # Parse the type string into a SimType
+                sim_type = parse_type(svar.type, predefined_types=types_store, arch=arch)
+                sim_type = sim_type.with_arch(arch)
+
+                # Get the variable manager and set the type
+                variable_kb = decompilation._variable_kb if hasattr(decompilation, '_variable_kb') else self.main_instance.project.kb
+                variable_manager = variable_kb.variables[svar.addr]
+                variable_manager.set_variable_type(dec_svar, sim_type, all_unified=True, mark_manual=True)
+                changed = True
+            except Exception as e:
+                l.warning(f"Failed to set stack variable type for {svar.name}: {e}")
+
         return changed
 
     def _set_comment(self, comment: Comment, decompilation=None, **kwargs) -> bool:
@@ -382,6 +406,109 @@ def _set_comment(self, comment: Comment, decompilation=None, **kwargs) -> bool:
         func_addr = comment.func_addr or self.get_closest_function(comment.addr)
         return changed & self.refresh_decompilation(func_addr)
 
+    # structs
+    def _structs(self) -> Dict[str, Struct]:
+        """
+        Returns a dict of libbs.Struct that contain the name and size of each struct in the decompiler.
+        """
+        from angr.sim_type import SimStruct, TypeRef
+        structs = {}
+        types_store = self.main_instance.project.kb.types
+
+        for type_ref in types_store.iter_own():
+            if not isinstance(type_ref, TypeRef):
+                continue
+            sim_type = type_ref.type
+            if isinstance(sim_type, SimStruct):
+                structs[type_ref.name] = Struct(type_ref.name, sim_type.size // 8 if sim_type.size else 0, {})
+
+        return structs
+
+    def _get_struct(self, name) -> Optional[Struct]:
+        """
+        Get a struct by name from the TypesStore.
+        """
+        from angr.sim_type import SimStruct, TypeRef
+        types_store = self.main_instance.project.kb.types
+
+        try:
+            type_ref = types_store[name]
+        except KeyError:
+            return None
+
+        if not isinstance(type_ref, TypeRef):
+            return None
+
+        sim_struct = type_ref.type
+        if not isinstance(sim_struct, SimStruct):
+            return None
+
+        return self._angr_struct_to_libbs(name, sim_struct)
+
+    def _set_struct(self, struct: Struct, header=True, members=True, **kwargs) -> bool:
+        """
+        Create or update a struct in the TypesStore.
+        """
+        from angr.sim_type import SimStruct, TypeRef, parse_type
+        from collections import OrderedDict
+
+        types_store = self.main_instance.project.kb.types
+        arch = self.main_instance.project.arch
+
+        # Build the fields OrderedDict from LibBS struct members
+        fields = OrderedDict()
+        if members and struct.members:
+            sorted_members = sorted(struct.members.items(), key=lambda x: x[0])
+            for offset, member in sorted_members:
+                # Parse the member type string into a SimType
+                try:
+                    sim_type = parse_type(member.type, predefined_types=types_store, arch=arch)
+                except Exception:
+                    # Fallback to a simple int type with the right size if parsing fails
+                    from angr.sim_type import SimTypeInt
+                    sim_type = SimTypeInt(signed=False).with_arch(arch)
+
+                fields[member.name] = sim_type.with_arch(arch)
+
+        # Create the SimStruct
+        sim_struct = SimStruct(fields, name=struct.name, pack=True)
+        sim_struct = sim_struct.with_arch(arch)
+
+        # Wrap it in a TypeRef and store it
+        type_ref = TypeRef(struct.name, sim_struct)
+        types_store[struct.name] = type_ref
+
+        return True
+
+    def _del_struct(self, name) -> bool:
+        """
+        Delete a struct from the TypesStore.
+        """
+        types_store = self.main_instance.project.kb.types
+
+        if name in types_store.data:
+            del types_store.data[name]
+            return True
+
+        return False
+
+    @staticmethod
+    def _angr_struct_to_libbs(name: str, sim_struct: "angr.sim_type.SimStruct") -> Struct:
+        """
+        Convert an angr SimStruct to a LibBS Struct.
+        """
+        members = {}
+        if sim_struct._arch is not None:
+            offsets = sim_struct.offsets
+            for field_name, sim_type in sim_struct.fields.items():
+                offset = offsets.get(field_name, 0)
+                type_str = sim_type.c_repr() if sim_type else None
+                size = sim_type.size // 8 if sim_type and sim_type.size else 0
+                members[offset] = StructMember(field_name, offset, type_str, size)
+
+        size = sim_struct.size // 8 if sim_struct.size else 0
+        return Struct(name, size, members)
+
     #
     #   Utils
     #
@@ -431,13 +558,16 @@ def addr_starts_instruction(self, addr) -> bool:
 
         return addr in node.instruction_addrs
 
-    def refresh_decompilation(self, func_addr):
+    def refresh_decompilation(self, func_addr, event=None):
         if self.headless:
             return False
 
         self.workspace.jump_to(func_addr)
         view = self.workspace._get_or_create_view("pseudocode", CodeView)
-        view.codegen.am_event()
+        if event:
+            view.codegen.am_event(event=event)
+        else:
+            view.codegen.am_event()
         view.focus()
         return True
 

tests/test_decompilers.py

@@ -436,13 +436,100 @@ def test_angr(self):
             headless=True,
             binary_path=self.FAUXWARE_PATH
         )
+        self.deci = deci
         func_addr = deci.art_lifter.lift_addr(0x400664)
         main = deci.functions[func_addr]
         main.name = self.RENAMED_NAME
         deci.functions[func_addr] = main
         assert deci.functions[func_addr].name == self.RENAMED_NAME
         assert self.RENAMED_NAME in deci.main_instance.project.kb.functions
 
+        #
+        # Struct support
+        #
+
+        # test struct creation
+        new_struct = Struct()
+        new_struct.name = "my_angr_struct"
+        new_struct.add_struct_member('char_member', 0, 'char', 1)
+        new_struct.add_struct_member('int_member', 1, 'int', 4)
+        deci.structs[new_struct.name] = new_struct
+
+        # verify struct was created
+        updated = deci.structs[new_struct.name]
+        assert updated is not None, "Struct was not created"
+        assert updated.name == new_struct.name
+
+        # verify members are present
+        assert 0 in updated.members, "First member not found"
+        assert 1 in updated.members, "Second member not found"
+
+        # test struct listing
+        struct_items = list(deci.structs.items())
+        struct_names = [k for k, v in struct_items]
+        assert new_struct.name in struct_names, "Struct not found in listing"
+
+        #
+        # Stack variable type setting
+        #
+
+        # Get the main function which has stack variables
+        main_func_addr = deci.art_lifter.lift_addr(0x40071d)
+        main_func = deci.functions[main_func_addr]
+
+        # Check that we have stack variables
+        assert len(main_func.stack_vars) > 0, "No stack variables found in main function"
+
+        # Get the first stack variable and change its type to a primitive
+        first_offset = list(main_func.stack_vars.keys())[0]
+        original_svar = main_func.stack_vars[first_offset]
+
+        # Set a new type (change to int)
+        original_svar.type = "int"
+        deci.functions[main_func_addr] = main_func
+
+        # Verify the type was set by re-fetching the function
+        updated_func = deci.functions[main_func_addr]
+        updated_svar = updated_func.stack_vars.get(first_offset)
+        assert updated_svar is not None, "Stack variable not found after update"
+        # The type should contain "int" (may be formatted differently by angr)
+        assert "int" in updated_svar.type.lower() if updated_svar.type else False, \
+            f"Stack variable type was not updated to int, got: {updated_svar.type}"
+
+        #
+        # Stack variable type setting with struct type
+        #
+
+        # Re-fetch the function to get fresh stack variables
+        main_func = deci.functions[main_func_addr]
+
+        # Get a stack variable (use the same one or another if available)
+        svar_offsets = list(main_func.stack_vars.keys())
+        struct_test_offset = svar_offsets[0] if len(svar_offsets) == 1 else svar_offsets[1]
+        struct_test_svar = main_func.stack_vars[struct_test_offset]
+
+        # Set the type to a pointer to our struct
+        struct_ptr_type = f"struct {new_struct.name} *"
+        struct_test_svar.type = struct_ptr_type
+        deci.functions[main_func_addr] = main_func
+
+        # Verify the struct type was set
+        updated_func = deci.functions[main_func_addr]
+        updated_svar = updated_func.stack_vars.get(struct_test_offset)
+        assert updated_svar is not None, "Stack variable not found after struct type update"
+        assert updated_svar.type is not None, "Stack variable type is None after struct type update"
+        # The type should contain the struct name
+        assert new_struct.name in updated_svar.type, \
+            f"Stack variable type was not updated to struct pointer, got: {updated_svar.type}"
+
+        # Now test struct deletion (after we're done using it for stack var types)
+        del deci.structs[new_struct.name]
+        struct_items = list(deci.structs.items())
+        struct_keys = [k for k, v in struct_items]
+        assert new_struct.name not in struct_keys, "Struct was not deleted"
+
+        deci.shutdown()
+
     def test_binja(self):
         deci = DecompilerInterface.discover(
             force_decompiler=BINJA_DECOMPILER,