diff --git a/angr/ailment/tagged_object.py b/angr/ailment/tagged_object.py index d6bbb1b9d..384cb2e56 100644 --- a/angr/ailment/tagged_object.py +++ b/angr/ailment/tagged_object.py @@ -31,6 +31,7 @@ class TagDict(TypedDict, total=False): keep_in_slice: bool orig_ins_addr: int reg_name: str + return_addr_assignment_removed: bool uninitialized: bool vex_block_addr: int vex_stmt_idx: int diff --git a/angr/analyses/decompiler/callsite_maker.py b/angr/analyses/decompiler/callsite_maker.py index 53c4e7f20..cc1d19f78 100644 --- a/angr/analyses/decompiler/callsite_maker.py +++ b/angr/analyses/decompiler/callsite_maker.py @@ -270,7 +270,9 @@ class CallSiteMaker: new_stmts = self.block.statements[:-1] # remove the statement that stores the return address - if self.project.arch.call_pushes_ret: + # The marker keeps a later CallSiteMaker pass from mistaking an earlier live write for the removed assignment. + return_addr_assignment_removed = call_expr.tags.get("return_addr_assignment_removed", False) + if not return_addr_assignment_removed and self.project.arch.call_pushes_ret: # check if the last statement is storing the return address onto the top of the stack for stmt_idx_r, the_stmt in enumerate(reversed(new_stmts)): stmt_idx = len(new_stmts) - 1 - stmt_idx_r @@ -291,16 +293,16 @@ class CallSiteMaker: if varid is not None: self.removed_vvar_ids.add(varid) new_stmts = new_stmts[:stmt_idx] + new_stmts[stmt_idx + 1 :] + return_addr_assignment_removed = True break - else: + elif not return_addr_assignment_removed: # if there is an lr register... - lr_offset = None - if archinfo.arch_arm.is_arm_arch(self.project.arch) or self.project.arch.name in {"PPC32", "PPC64"}: - lr_offset = self.project.arch.registers["lr"][0] - elif self.project.arch.name in {"MIPS32", "MIPS64"}: + lr_offset = self.project.arch.lr_offset + if lr_offset is None and self.project.arch.name in {"MIPS32", "MIPS64"}: lr_offset = self.project.arch.registers["ra"][0] # remove the assignment to the lr register - if lr_offset is not None: + if lr_offset is not None and self.block.original_size is not None: + expected_return_addr = self.block.addr + self.block.original_size for stmt_idx_r, the_stmt in enumerate(reversed(new_stmts)): stmt_idx = len(new_stmts) - 1 - stmt_idx_r if isinstance(the_stmt, Stmt.SideEffectStatement): @@ -317,10 +319,13 @@ class CallSiteMaker: varid = the_stmt.dst.varid else: continue + if not isinstance(the_stmt.src, Expr.Const) or the_stmt.src.value != expected_return_addr: + continue # found it new_stmts = new_stmts[:stmt_idx] + new_stmts[stmt_idx + 1 :] if varid is not None: self.removed_vvar_ids.add(varid) + return_addr_assignment_removed = True break # calculate stack offsets for arguments that are put on the stack. these offsets will be consumed by @@ -386,6 +391,8 @@ class CallSiteMaker: tags = call_expr.tags.copy() tags.pop("arg_vvars", None) + if return_addr_assignment_removed: + tags["return_addr_assignment_removed"] = True if func is not None: tags["is_prototype_guessed"] = func.is_prototype_guessed new_call = Expr.Call( diff --git a/angr/analyses/variable_recovery/engine_base.py b/angr/analyses/variable_recovery/engine_base.py index 8619340f8..82e798e74 100644 --- a/angr/analyses/variable_recovery/engine_base.py +++ b/angr/analyses/variable_recovery/engine_base.py @@ -311,10 +311,11 @@ class SimEngineVRBase[VRStateType: VariableRecoveryStateBase, BlockType: BlockPr :return: """ - if ( - offset in (self.project.arch.ip_offset, self.project.arch.sp_offset, self.project.arch.lr_offset) - or not create_variable - ): + excluded_register_offsets = ( + self.project.arch.ip_offset, + self.project.arch.sp_offset, + ) + if offset in excluded_register_offsets or not create_variable: # only store the value. don't worry about variables. v = MultiValues(richr.data) self.state.register_region.store(offset, v) @@ -398,7 +399,7 @@ class SimEngineVRBase[VRStateType: VariableRecoveryStateBase, BlockType: BlockPr if ( vvar.category == ailment.expression.VirtualVariableCategory.REGISTER - and vvar.oident in (self.project.arch.ip_offset, self.project.arch.sp_offset, self.project.arch.lr_offset) + and vvar.oident in (self.project.arch.ip_offset, self.project.arch.sp_offset) ) or not create_variable: # only store the value. don't worry about variables. self.vvar_region[vvar_id] = richr.data diff --git a/tests/analyses/test_callsite_maker.py b/tests/analyses/test_callsite_maker.py index 7a891089a..4e6426503 100755 --- a/tests/analyses/test_callsite_maker.py +++ b/tests/analyses/test_callsite_maker.py @@ -5,6 +5,7 @@ __package__ = __package__ or "tests.analyses" # pylint:disable=redefined-builti import os import unittest +from typing import cast import angr from angr import ailment @@ -18,6 +19,286 @@ test_location = os.path.join(bin_location, "tests") # pylint: disable=missing-class-docstring # pylint: disable=no-self-use class TestCallsiteMaker(unittest.TestCase): + @staticmethod + def _register_destination(manager, project, offset, varid, use_vvar): + if use_vvar: + return ailment.Expr.VirtualVariable( + manager.next_atom(), + varid, + project.arch.bits, + ailment.Expr.VirtualVariableCategory.REGISTER, + oident=offset, + ) + return ailment.Expr.Register(manager.next_atom(), offset, project.arch.bits) + + @staticmethod + def _assignment(manager, dst, src, ins_addr): + return ailment.Stmt.Assignment(manager.next_atom(), dst, src, ins_addr=ins_addr) + + def _make_callsite(self, project, manager, block_addr, block_size, assignments): + call = ailment.Expr.Call( + manager.next_atom(), + ailment.Expr.Const(manager.next_atom(), 0x2000, project.arch.bits), + args=[], + bits=project.arch.bits, + ins_addr=block_addr + 2, + ) + call_stmt = ailment.Stmt.SideEffectStatement(manager.next_atom(), call, ins_addr=block_addr + 2) + block = ailment.Block( + block_addr, + original_size=block_size, + statements=[*assignments, call_stmt], + ) + + callsite_maker = CallSiteMaker(project, block, ail_manager=manager) + self.assertIsNotNone(callsite_maker.result_block) + return callsite_maker, cast(ailment.Block, callsite_maker.result_block) + + def _assert_return_addr_assignment_removed(self, block): + call_stmt = block.statements[-1] + self.assertIsInstance(call_stmt, ailment.Stmt.SideEffectStatement) + self.assertTrue(call_stmt.expr.tags.get("return_addr_assignment_removed", False)) + + def test_synthetic_link_register_assignment(self): + for arch_dir, binary, block_addr, block_size in ( + ("armel", "checkbyte", 0x1001, 4), + ("aarch64", "test_loops", 0x1000, 4), + ("ppc", "checkbyte", 0x1000, 4), + ("mips", "checkbyte", 0x1000, 8), + ("riscv64", "test_return_type_riscv64.elf", 0x1000, 4), + ("s390x", "checkbyte", 0x1000, 6), + ): + project = angr.Project( + os.path.join(test_location, arch_dir, binary), + auto_load_libs=False, + ) + self.assertFalse(project.arch.call_pushes_ret) + lr_offset = project.arch.lr_offset + if lr_offset is None and project.arch.name in {"MIPS32", "MIPS64"}: + lr_offset = project.arch.registers["ra"][0] + self.assertIsNotNone(lr_offset) + + for use_vvar in (False, True): + with self.subTest(arch=project.arch.name, use_vvar=use_vvar): + self._check_synthetic_link_register_assignment(project, block_addr, block_size, lr_offset, use_vvar) + + def _check_synthetic_link_register_assignment(self, project, block_addr, block_size, lr_offset, use_vvar): + manager = ailment.Manager(arch=project.arch) + scratch_dst = self._register_destination(manager, project, lr_offset, 1, use_vvar) + link_dst = self._register_destination(manager, project, lr_offset, 2, use_vvar) + + scratch_assignment = self._assignment( + manager, + scratch_dst, + ailment.Expr.Const(manager.next_atom(), 1, project.arch.bits), + block_addr, + ) + link_assignment = self._assignment( + manager, + link_dst, + ailment.Expr.Const(manager.next_atom(), block_addr + block_size, project.arch.bits), + block_addr + 2, + ) + callsite_maker, result_block = self._make_callsite( + project, + manager, + block_addr, + block_size, + [scratch_assignment, link_assignment], + ) + self.assertEqual(len(result_block.statements), 2) + self.assertIs(result_block.statements[0], scratch_assignment) + self.assertIsInstance(result_block.statements[1], ailment.Stmt.SideEffectStatement) + self.assertEqual(callsite_maker.removed_vvar_ids, {2} if use_vvar else set()) + self._assert_return_addr_assignment_removed(result_block) + + def test_real_lifted_arm_link_register_assignment(self): + project = angr.Project( + os.path.join(test_location, "armel", "checkbyte"), + auto_load_libs=False, + ) + manager = ailment.Manager(arch=project.arch) + block = project.factory.block(0x8370, size=8) + self.assertEqual(block.vex.jumpkind, "Ijk_Call") + ail_block = ailment.IRSBConverter.convert(block.vex, manager) + + link_assignment = ail_block.statements[-2] + self.assertIsInstance(link_assignment, ailment.Stmt.Assignment) + link_assignment = cast(ailment.Stmt.Assignment, link_assignment) + self.assertIsInstance(link_assignment.dst, ailment.Expr.Register) + link_dst = cast(ailment.Expr.Register, link_assignment.dst) + self.assertEqual(link_dst.reg_offset, project.arch.lr_offset) + self.assertIsInstance(link_assignment.src, ailment.Expr.Const) + link_src = cast(ailment.Expr.Const, link_assignment.src) + original_size = cast(int, ail_block.original_size) + self.assertEqual(link_src.value, ail_block.addr + original_size) + + callsite_maker = CallSiteMaker(project, ail_block, ail_manager=manager) + + self.assertIsNotNone(callsite_maker.result_block) + result_block = cast(ailment.Block, callsite_maker.result_block) + self.assertEqual(len(result_block.statements), len(ail_block.statements) - 1) + self.assertNotIn(link_assignment, result_block.statements) + self._assert_return_addr_assignment_removed(result_block) + + def test_non_synthetic_link_register_assignment_survives(self): + project = angr.Project( + os.path.join(test_location, "armel", "checkbyte"), + auto_load_libs=False, + ) + block_addr = 0x1001 + block_size = 4 + lr_offset = project.arch.lr_offset + self.assertIsNotNone(lr_offset) + + for use_vvar in (False, True): + with self.subTest(use_vvar=use_vvar): + manager = ailment.Manager(arch=project.arch) + scratch_dst = self._register_destination(manager, project, lr_offset, 1, use_vvar) + scratch_assignment = self._assignment( + manager, + scratch_dst, + ailment.Expr.Const(manager.next_atom(), 0x4242, project.arch.bits), + block_addr, + ) + callsite_maker, result_block = self._make_callsite( + project, + manager, + block_addr, + block_size, + [scratch_assignment], + ) + self.assertEqual(len(result_block.statements), 2) + self.assertIs(result_block.statements[0], scratch_assignment) + self.assertEqual(callsite_maker.removed_vvar_ids, set()) + + def test_synthetic_link_register_assignment_removal_is_idempotent(self): + project = angr.Project( + os.path.join(test_location, "armel", "checkbyte"), + auto_load_libs=False, + ) + block_addr = 0x1001 + block_size = 4 + lr_offset = project.arch.lr_offset + r0_offset = project.arch.registers["r0"][0] + self.assertIsNotNone(lr_offset) + + for use_vvar in (False, True): + with self.subTest(use_vvar=use_vvar): + manager = ailment.Manager(arch=project.arch) + scratch_dst = self._register_destination(manager, project, lr_offset, 1, use_vvar) + scratch_use = self._register_destination(manager, project, lr_offset, 1, use_vvar) + use_dst = self._register_destination(manager, project, r0_offset, 3, use_vvar) + link_dst = self._register_destination(manager, project, lr_offset, 2, use_vvar) + return_addr = ailment.Expr.Const(manager.next_atom(), block_addr + block_size, project.arch.bits) + scratch_assignment = self._assignment(manager, scratch_dst, return_addr, block_addr) + use_assignment = self._assignment(manager, use_dst, scratch_use, block_addr + 2) + link_assignment = self._assignment(manager, link_dst, return_addr, block_addr + 2) + first_callsite_maker, first_result = self._make_callsite( + project, + manager, + block_addr, + block_size, + [scratch_assignment, use_assignment, link_assignment], + ) + self.assertEqual(len(first_result.statements), 3) + self.assertEqual(first_callsite_maker.removed_vvar_ids, {2} if use_vvar else set()) + self._assert_return_addr_assignment_removed(first_result) + + second_callsite_maker = CallSiteMaker(project, first_result, ail_manager=manager) + + self.assertIsNotNone(second_callsite_maker.result_block) + second_result = cast(ailment.Block, second_callsite_maker.result_block) + self.assertEqual(len(second_result.statements), 3) + self.assertIs(second_result.statements[0], scratch_assignment) + self.assertIs(second_result.statements[1], use_assignment) + self.assertEqual(second_callsite_maker.removed_vvar_ids, set()) + self._assert_return_addr_assignment_removed(second_result) + + def test_stack_return_address_assignment_removal_is_idempotent(self): + project = angr.Project( + os.path.join(test_location, "x86_64", "all"), + auto_load_libs=False, + ) + manager = ailment.Manager(arch=project.arch) + block_addr = 0x1000 + block_size = 5 + return_addr = ailment.Expr.Const(manager.next_atom(), block_addr + block_size, project.arch.bits) + scratch_store = ailment.Stmt.Store( + manager.next_atom(), + ailment.Expr.Const(manager.next_atom(), 0x3000, project.arch.bits), + return_addr, + project.arch.bytes, + project.arch.memory_endness, + ins_addr=block_addr, + ) + link_store = ailment.Stmt.Store( + manager.next_atom(), + ailment.Expr.Const(manager.next_atom(), 0x4000, project.arch.bits), + return_addr, + project.arch.bytes, + project.arch.memory_endness, + ins_addr=block_addr, + ) + + _, first_result = self._make_callsite( + project, + manager, + block_addr, + block_size, + [scratch_store, link_store], + ) + self.assertEqual(len(first_result.statements), 2) + self.assertIs(first_result.statements[0], scratch_store) + self._assert_return_addr_assignment_removed(first_result) + + second_callsite_maker = CallSiteMaker(project, first_result, ail_manager=manager) + + self.assertIsNotNone(second_callsite_maker.result_block) + second_result = cast(ailment.Block, second_callsite_maker.result_block) + self.assertEqual(len(second_result.statements), 2) + self.assertIs(second_result.statements[0], scratch_store) + self.assertEqual(second_callsite_maker.removed_vvar_ids, set()) + self._assert_return_addr_assignment_removed(second_result) + + def test_s390x_noncanonical_link_preserves_lr_scratch(self): + project = angr.Project( + os.path.join(test_location, "s390x", "checkbyte"), + auto_load_libs=False, + ) + manager = ailment.Manager(arch=project.arch) + block_addr = 0x1000 + block_size = 6 + lr_offset = project.arch.lr_offset + other_link_offset = project.arch.registers["r0"][0] + self.assertIsNotNone(lr_offset) + + scratch_dst = self._register_destination(manager, project, lr_offset, 1, False) + link_dst = self._register_destination(manager, project, other_link_offset, 2, False) + scratch_assignment = self._assignment( + manager, + scratch_dst, + ailment.Expr.Const(manager.next_atom(), 0x4242, project.arch.bits), + block_addr, + ) + link_assignment = self._assignment( + manager, + link_dst, + ailment.Expr.Const(manager.next_atom(), block_addr + block_size, project.arch.bits), + block_addr + 2, + ) + callsite_maker, result_block = self._make_callsite( + project, + manager, + block_addr, + block_size, + [scratch_assignment, link_assignment], + ) + self.assertEqual(len(result_block.statements), 3) + self.assertIs(result_block.statements[0], scratch_assignment) + self.assertIs(result_block.statements[1], link_assignment) + self.assertEqual(callsite_maker.removed_vvar_ids, set()) + def test_callsite_maker(self): project = angr.Project( os.path.join(test_location, "x86_64", "all"), diff --git a/tests/analyses/test_variablerecovery.py b/tests/analyses/test_variablerecovery.py index d3e921ebc..0c26d2a2b 100755 --- a/tests/analyses/test_variablerecovery.py +++ b/tests/analyses/test_variablerecovery.py @@ -2,12 +2,23 @@ from __future__ import annotations __package__ = __package__ or "tests.analyses" # pylint:disable=redefined-builtin +# pylint:disable=missing-class-docstring,protected-access import logging import os import unittest +from collections import defaultdict +from types import SimpleNamespace +from typing import Any, cast + +import claripy import angr +from angr import ailment +from angr.analyses.typehoon.typevars import TypeVariableManager +from angr.analyses.variable_recovery.engine_base import RichR, SimEngineVRBase +from angr.analyses.variable_recovery.variable_recovery_fast import VariableRecoveryFastState +from angr.knowledge_plugins.functions import Function from angr.knowledge_plugins.variables import VariableType from angr.sim_variable import SimRegisterVariable, SimStackVariable from tests.common import bin_location, print_decompilation_result @@ -17,12 +28,118 @@ test_location = os.path.join(bin_location, "tests") l = logging.getLogger("test_variablerecovery") +class _VariableRecoveryEngine(SimEngineVRBase): + def process(self, state, *, block=None, **kwargs): # pylint:disable=unused-argument + return None + + # # Utility methods # class TestVariableRecovery(unittest.TestCase): + def _arm_variable_recovery_engine(self): + project = angr.Project(os.path.join(test_location, "armel", "checkbyte"), auto_load_libs=False) + function = cast( + Function, + project.kb.functions.function(addr=project.entry, name="test_lr_assignment", create=True), + ) + self.assertIsNotNone(function) + + analysis: Any = SimpleNamespace( + _dominance_frontiers=defaultdict(set), + variable_manager=project.kb.variables, + get_variable_definitions=lambda _: set(), + ) + tv_manager = TypeVariableManager(function.addr) + state = VariableRecoveryFastState( + function.addr, + analysis, + project.arch, + function, + project, + tv_manager=tv_manager, + ) + engine = _VariableRecoveryEngine(project, project.kb, tv_manager=tv_manager) + engine.state = state + engine.block = ailment.Block(function.addr, statements=[]) + engine.stmt_idx = 0 + engine.ins_addr = function.addr + return project, function, engine + + def test_lr_assignment_creates_register_variable(self): + project, function, engine = self._arm_variable_recovery_engine() + lr = cast(int, project.arch.lr_offset) + lr_dst = ailment.Expr.Register(0, lr, project.arch.bits) + + engine._assign_to_register( + lr, + RichR(claripy.BVV(1, project.arch.bits)), + project.arch.bytes, + dst=lr_dst, + ) + + lr_variables = project.kb.variables[function.addr].find_variables_by_atom( + function.addr, engine.stmt_idx, lr_dst + ) + self.assertEqual(len(lr_variables), 1) + lr_variable, _ = next(iter(lr_variables)) + self.assertIsInstance(lr_variable, SimRegisterVariable) + lr_variable = cast(SimRegisterVariable, lr_variable) + self.assertEqual(lr_variable.reg, lr) + + for stmt_idx, excluded_offset in enumerate((project.arch.ip_offset, project.arch.sp_offset), start=1): + with self.subTest(offset=excluded_offset): + excluded_offset = cast(int, excluded_offset) + engine.stmt_idx = stmt_idx + dst = ailment.Expr.Register(stmt_idx, excluded_offset, project.arch.bits) + engine._assign_to_register( + excluded_offset, + RichR(claripy.BVV(stmt_idx, project.arch.bits)), + project.arch.bytes, + dst=dst, + ) + variables = project.kb.variables[function.addr].find_variables_by_atom(function.addr, stmt_idx, dst) + self.assertFalse(variables) + + def test_lr_vvar_assignment_creates_register_variable(self): + project, function, engine = self._arm_variable_recovery_engine() + category = ailment.Expr.VirtualVariableCategory.REGISTER + lr = cast(int, project.arch.lr_offset) + lr_vvar = ailment.Expr.VirtualVariable(0, 1, project.arch.bits, category, oident=lr) + + lr_variable = engine._assign_to_vvar( + lr_vvar, + RichR(claripy.BVV(1, project.arch.bits)), + dst=lr_vvar, + ) + + self.assertIsInstance(lr_variable, SimRegisterVariable) + lr_variable = cast(SimRegisterVariable, lr_variable) + self.assertEqual(lr_variable.reg, lr) + self.assertEqual( + project.kb.variables[function.addr].find_variables_by_atom(function.addr, engine.stmt_idx, lr_vvar), + {(lr_variable, None)}, + ) + + for stmt_idx, excluded_offset in enumerate((project.arch.ip_offset, project.arch.sp_offset), start=1): + with self.subTest(offset=excluded_offset): + excluded_offset = cast(int, excluded_offset) + engine.stmt_idx = stmt_idx + vvar = ailment.Expr.VirtualVariable( + stmt_idx, stmt_idx + 1, project.arch.bits, category, excluded_offset + ) + variable = engine._assign_to_vvar( + vvar, + RichR(claripy.BVV(stmt_idx, project.arch.bits)), + dst=vvar, + ) + self.assertIsNone(variable) + self.assertFalse( + project.kb.variables[function.addr].find_variables_by_atom(function.addr, stmt_idx, vvar) + ) + def _compare_memory_variable(self, variable, variable_info): if variable_info["location"] == "stack": if not isinstance(variable, SimStackVariable):