angr 9.2.192__cp311-cp311-macosx_10_12_x86_64.whl

angr/analyses/decompiler/optimization_passes/peephole_simplifier.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from angr import ailment
+from angr.analyses.decompiler.utils import (
+    peephole_optimize_expr,
+)
+from angr.analyses.decompiler.sequence_walker import SequenceWalker
+from angr.analyses.decompiler.peephole_optimizations import (
+    PeepholeOptimizationExprBase,
+    EXPR_OPTS,
+)
+from .optimization_pass import OptimizationPassStage, SequenceOptimizationPass
+
+
+class ExpressionSequenceWalker(SequenceWalker):
+    """
+    Walks sequences with generic expression handling.
+    """
+
+    def _handle(self, node, **kwargs):
+        if isinstance(node, ailment.Expr.Expression):
+            handler = self._handlers.get(ailment.Expr.Expression, None)
+            if handler:
+                return handler(node, **kwargs)
+        return super()._handle(node, **kwargs)
+
+
+class PostStructuringPeepholeOptimizationPass(SequenceOptimizationPass):
+    """
+    Perform a post-structuring peephole optimization pass to simplify node statements and expressions.
+    """
+
+    ARCHES = None
+    PLATFORMS = None
+    STAGE = OptimizationPassStage.AFTER_STRUCTURING
+    NAME = "Post-Structuring Peephole Optimization"
+    DESCRIPTION = (__doc__ or "").strip()
+
+    def __init__(self, func, peephole_optimizations=None, **kwargs):
+        super().__init__(func, **kwargs)
+        self._peephole_optimizations = peephole_optimizations
+        self._expr_peephole_opts = [
+            cls(self.project, self.kb, self._func.addr)
+            for cls in (self._peephole_optimizations or EXPR_OPTS)
+            if issubclass(cls, PeepholeOptimizationExprBase)
+        ]
+        self.analyze()
+
+    def _check(self):
+        return True, None
+
+    def _analyze(self, cache=None):
+        walker = ExpressionSequenceWalker(
+            handlers={ailment.Expr.Expression: self._optimize_expr, ailment.Block: self._optimize_block}
+        )
+        walker.walk(self.seq)
+        self.out_seq = self.seq
+
+    def _optimize_expr(self, expr, **_):
+        new_expr = peephole_optimize_expr(expr, self._expr_peephole_opts)
+        return new_expr if expr is not new_expr else None
+
+    def _optimize_block(self, block, **_):
+        old_block, new_block = None, block
+        while old_block != new_block:
+            old_block = new_block
+            # Note: AILBlockSimplifier updates expressions in place
+            simp = self.project.analyses.AILBlockSimplifier(
+                new_block,
+                func_addr=self._func.addr,
+                peephole_optimizations=self._peephole_optimizations,
+            )
+            assert simp.result_block is not None
+            new_block = simp.result_block
+        return new_block if block != new_block else None

angr/analyses/decompiler/optimization_passes/register_save_area_simplifier.py

@@ -0,0 +1,263 @@
+# pylint:disable=too-many-boolean-expressions
+from __future__ import annotations
+from collections.abc import Iterable
+import logging
+
+import archinfo
+import angr.ailment as ailment
+
+from angr.calling_conventions import SimRegArg
+from angr.code_location import CodeLocation
+from angr.analyses.decompiler.stack_item import StackItem, StackItemType
+from .optimization_pass import OptimizationPass, OptimizationPassStage
+
+
+_l = logging.getLogger(name=__name__)
+
+
+class RegisterSaveAreaSimplifier(OptimizationPass):
+    """
+    Optimizes away register spilling effects, including callee-saved registers.
+
+    This optimization runs between SSA-level0 and SSA-level1, which means registers are converted to vvars but stack
+    accesses stay unchanged.
+    """
+
+    ARCHES = None
+    PLATFORMS = None
+    STAGE = OptimizationPassStage.AFTER_SINGLE_BLOCK_SIMPLIFICATION
+    NAME = "Simplify register save areas"
+    DESCRIPTION = __doc__.strip()
+
+    def __init__(self, func, **kwargs):
+        super().__init__(func, **kwargs)
+
+        self.analyze()
+
+    def _check(self):
+        # Check the first block to see what external registers are stored on the stack
+        stored_info = self._find_registers_stored_on_stack()
+        if not stored_info:
+            return False, None
+
+        # Check all return sites to see what external registers are restored to registers from the stack
+        restored_info = self._find_registers_restored_from_stack()
+        if not restored_info:
+            return False, None
+
+        # Find common registers and stack offsets
+        info = self._intersect_register_info(stored_info, restored_info)
+
+        return bool(info), {"info": info}
+
+    @staticmethod
+    def _modify_statement(
+        old_block, stmt_idx_: int, updated_blocks_, stack_offset: int | None = None
+    ):  # pylint:disable=unused-argument
+        if old_block not in updated_blocks_:
+            block = old_block.copy()
+            updated_blocks_[old_block] = block
+        else:
+            block = updated_blocks_[old_block]
+        block.statements[stmt_idx_] = None
+
+    def _analyze(self, cache=None):
+
+        if cache is None:
+            return
+
+        info: dict[int, dict[str, list[tuple[int, CodeLocation]]]] = cache["info"]
+        updated_blocks = {}
+
+        for data in info.values():
+            # remove storing statements
+            for stack_offset, codeloc in data["stored"]:
+                old_block = self._get_block(codeloc.block_addr, idx=codeloc.block_idx)
+                self._modify_statement(old_block, codeloc.stmt_idx, updated_blocks, stack_offset=stack_offset)
+            for stack_offset, codeloc in data["restored"]:
+                old_block = self._get_block(codeloc.block_addr, idx=codeloc.block_idx)
+                self._modify_statement(old_block, codeloc.stmt_idx, updated_blocks, stack_offset=stack_offset)
+
+        for old_block, new_block in updated_blocks.items():
+            # remove all statements that are None
+            new_block.statements = [stmt for stmt in new_block.statements if stmt is not None]
+            # update it
+            self._update_block(old_block, new_block)
+
+        if updated_blocks:
+            # update stack_items
+            for data in info.values():
+                for stack_offset, _ in data["stored"]:
+                    self.stack_items[stack_offset] = StackItem(
+                        stack_offset, self.project.arch.bytes, "regs", StackItemType.SAVED_REGS
+                    )
+
+    def _find_registers_stored_on_stack(self) -> list[tuple[int, int, CodeLocation]]:
+        first_block = self._get_block(self._func.addr)
+        if first_block is None:
+            return []
+
+        results = []
+
+        # there are cases where sp is moved to another register at the beginning of the function before it is
+        # subtracted, then it is used for stack accesses.
+        # for example:
+        # 132B0 mov     r11, rsp
+        # 132B3 sub     rsp, 88h
+        # 132BA ...
+        # 132C1 ...
+        # 132C6 mov     [r11-18h], r13
+        # 132CA mov     [r11-20h], r14
+        # we support such cases because we will have rewritten r11-18h to SpOffset(-N) when this optimization runs.
+
+        # identify which registers have been updated in this function; they are no longer saved
+        ignored_regs: set[int] = set()
+
+        for idx, stmt in enumerate(first_block.statements):
+            if (
+                isinstance(stmt, ailment.Stmt.Assignment)
+                and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                and stmt.dst.was_reg
+            ):
+                ignored_regs.add(stmt.dst.reg_offset)
+            if (
+                isinstance(stmt, ailment.Stmt.Store)
+                and isinstance(stmt.addr, ailment.Expr.StackBaseOffset)
+                and isinstance(stmt.addr.offset, int)
+            ):
+                if isinstance(stmt.data, ailment.Expr.VirtualVariable) and stmt.data.was_reg:
+                    if stmt.data.reg_offset not in ignored_regs:
+                        # it's storing registers to the stack!
+                        stack_offset = stmt.addr.offset
+                        reg_offset = stmt.data.reg_offset
+                        codeloc = CodeLocation(
+                            first_block.addr, idx, block_idx=first_block.idx, ins_addr=stmt.tags["ins_addr"]
+                        )
+                        results.append((reg_offset, stack_offset, codeloc))
+                elif (
+                    self.project.arch.name == "AMD64"
+                    and isinstance(stmt.data, ailment.Expr.Convert)
+                    and isinstance(stmt.data.operand, ailment.Expr.VirtualVariable)
+                    and stmt.data.operand.was_reg
+                    and stmt.data.from_bits == 256
+                    and stmt.data.to_bits == 128
+                ):
+                    # storing xmm registers to the stack
+                    stack_offset = stmt.addr.offset
+                    reg_offset = stmt.data.operand.reg_offset
+                    codeloc = CodeLocation(
+                        first_block.addr, idx, block_idx=first_block.idx, ins_addr=stmt.tags["ins_addr"]
+                    )
+                    results.append((reg_offset, stack_offset, codeloc))
+
+        return results
+
+    def _find_registers_restored_from_stack(self) -> list[list[tuple[int, int, CodeLocation]]]:
+        all_results = []
+        for ret_site in self._func.ret_sites + self._func.jumpout_sites:
+
+            ret_blocks = list(self._get_blocks(ret_site.addr))
+            if len(ret_blocks) == 1 and self.project.simos is not None and self.project.simos.name == "Win32":
+                # PE files may call __security_check_cookie (which terminates the program if the stack canary is
+                # corrupted) before returning.
+                preds = list(self._graph.predecessors(ret_blocks[0]))
+                if len(preds) == 1:
+                    pred = preds[0]
+                    if pred.statements and isinstance(pred.statements[-1], ailment.Stmt.Call):
+                        last_stmt = pred.statements[-1]
+                        if isinstance(last_stmt.target, ailment.Expr.Const):
+                            callee_addr = last_stmt.target.value
+                            if self.project.kb.functions.contains_addr(callee_addr):
+                                callee_func = self.project.kb.functions.get_by_addr(callee_addr)
+                                if callee_func.name == "_security_check_cookie":
+                                    ret_blocks.append(pred)
+
+            for block in ret_blocks:
+                results = []
+                for idx, stmt in enumerate(block.statements):
+                    if (
+                        isinstance(stmt, ailment.Stmt.Assignment)
+                        and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                        and stmt.dst.was_reg
+                        and isinstance(stmt.src, ailment.Expr.Load)
+                        and isinstance(stmt.src.addr, ailment.Expr.StackBaseOffset)
+                    ):
+                        stack_offset = stmt.src.addr.offset
+                        reg_offset = stmt.dst.reg_offset
+                        codeloc = CodeLocation(block.addr, idx, block_idx=block.idx, ins_addr=stmt.tags["ins_addr"])
+                        results.append((reg_offset, stack_offset, codeloc))
+
+                if results:
+                    all_results.append(results)
+
+        return all_results
+
+    def _intersect_register_info(
+        self,
+        stored: list[tuple[int, int, CodeLocation]],
+        restored: Iterable[list[tuple[int, int, CodeLocation]]],
+    ) -> dict[int, dict[str, list[tuple[int, CodeLocation]]]]:
+        def _collect(info: list[tuple[int, int, CodeLocation]], output, keystr: str):
+            for reg_offset, stack_offset, codeloc in info:
+                if reg_offset not in output:
+                    output[reg_offset] = {}
+                if keystr not in output[reg_offset]:
+                    output[reg_offset][keystr] = []
+                output[reg_offset][keystr].append((stack_offset, codeloc))
+
+        result: dict[int, dict[str, list[tuple[int, CodeLocation]]]] = {}
+        _collect(stored, result, "stored")
+        for item in restored:
+            _collect(item, result, "restored")
+
+        # remove registers that are
+        # (a) stored but not restored
+        # (b) restored but not stored
+        # (c) from different offsets
+        # (d) the same as the return value register
+
+        cc = self._func.calling_convention
+        if cc is not None and isinstance(cc.RETURN_VAL, SimRegArg):
+            ret_val_reg_offset = self.project.arch.registers[cc.RETURN_VAL.reg_name][0]
+        else:
+            ret_val_reg_offset = None
+
+        # link register
+        if archinfo.arch_arm.is_arm_arch(self.project.arch):
+            lr_reg_offset = self.project.arch.registers["lr"][0]
+        elif self.project.arch.name in {"MIPS32", "MIPS64"}:
+            lr_reg_offset = self.project.arch.registers["ra"][0]
+        elif self.project.arch.name in {"PPC32", "PPC64"} or self.project.arch.name.startswith("PowerPC:"):
+            lr_reg_offset = self.project.arch.registers["lr"][0]
+        else:
+            lr_reg_offset = None
+
+        for reg in list(result.keys()):
+            # stored link register should always be removed
+            if lr_reg_offset is not None and reg == lr_reg_offset:
+                if "restored" not in result[reg]:
+                    # add a dummy one
+                    result[reg]["restored"] = []
+                continue
+
+            if ret_val_reg_offset is not None and reg == ret_val_reg_offset:
+                # (d)
+                del result[reg]
+                continue
+
+            info = result[reg]
+            if len(info.keys()) != 2:
+                # (a) or (b)
+                del result[reg]
+                continue
+
+            stack_offsets = {stack_offset for stack_offset, _ in info["stored"]} | {
+                stack_offset for stack_offset, _ in info["restored"]
+            }
+
+            if len(stack_offsets) != 1:
+                # (c)
+                del result[reg]
+                continue
+
+        return result

angr/analyses/decompiler/optimization_passes/register_save_area_simplifier_adv.py

@@ -0,0 +1,198 @@
+# pylint:disable=too-many-boolean-expressions
+from __future__ import annotations
+import logging
+
+
+from angr.ailment.statement import Assignment
+from angr.ailment.expression import VirtualVariable
+from angr.code_location import CodeLocation, ExternalCodeLocation
+from angr.analyses.decompiler.stack_item import StackItem, StackItemType
+from angr.utils.ail import is_phi_assignment
+from .optimization_pass import OptimizationPass, OptimizationPassStage
+
+_l = logging.getLogger(name=__name__)
+
+
+class RegisterSaveAreaSimplifierAdvanced(OptimizationPass):
+    """
+    Optimizes away registers that are stored to or restored on the stack space.
+
+    This analysis is more complex than RegisterSaveAreaSimplifier because it handles:
+    (1) Registers that are stored in the stack shadow space (sp+N) according to the Windows x64 calling convention.
+    (2) Registers that are aliases of sp.
+    """
+
+    ARCHES = None
+    PLATFORMS = None
+    STAGE = OptimizationPassStage.AFTER_MAKING_CALLSITES
+    NAME = "Simplify register save areas (advanced)"
+    DESCRIPTION = __doc__.strip()  # type:ignore
+
+    def __init__(self, func, **kwargs):
+        super().__init__(func, **kwargs)
+        self._srda = None
+
+        self.analyze()
+
+    def _check(self):
+
+        self._srda = self.project.analyses.SReachingDefinitions(
+            subject=self._func, func_graph=self._graph, func_args=self._arg_vvars
+        )
+        info = self._find_reg_store_and_restore_locations()
+        if not info:
+            return False, None
+
+        return True, {"info": info}
+
+    @staticmethod
+    def _modify_statement(
+        old_block, stmt_idx_: int, updated_blocks_, stack_offset: int | None = None
+    ):  # pylint:disable=unused-argument
+        if old_block not in updated_blocks_:
+            block = old_block.copy()
+            updated_blocks_[old_block] = block
+        else:
+            block = updated_blocks_[old_block]
+        block.statements[stmt_idx_] = None
+
+    def _analyze(self, cache=None):
+
+        if cache is None:
+            return
+
+        info: list[tuple[int, CodeLocation, int, CodeLocation, int]] = cache["info"]
+        updated_blocks = {}
+
+        for _regvar, regvar_loc, _stackvar, stackvar_loc, _ in info:
+            # remove storing statements
+            old_block = self._get_block(regvar_loc.block_addr, idx=regvar_loc.block_idx)
+            assert regvar_loc.stmt_idx is not None
+            self._modify_statement(old_block, regvar_loc.stmt_idx, updated_blocks)
+            old_block = self._get_block(stackvar_loc.block_addr, idx=stackvar_loc.block_idx)
+            assert stackvar_loc.stmt_idx is not None
+            self._modify_statement(old_block, stackvar_loc.stmt_idx, updated_blocks)
+
+        for old_block, new_block in updated_blocks.items():
+            # remove all statements that are None
+            new_block.statements = [stmt for stmt in new_block.statements if stmt is not None]
+            # update it
+            self._update_block(old_block, new_block)
+
+        if updated_blocks:
+            # update stack_items
+            for _, _, _, _, stack_offset in info:
+                self.stack_items[stack_offset] = StackItem(
+                    stack_offset, self.project.arch.bytes, "regs", StackItemType.SAVED_REGS
+                )
+
+    def _find_reg_store_and_restore_locations(self) -> list[tuple[int, CodeLocation, int, CodeLocation, int]]:
+        results = []
+
+        assert self._srda is not None
+        srda_model = self._srda.model
+        # find all registers that are defined externally and used exactly once
+        saved_vvars: set[tuple[int, CodeLocation]] = set()
+        for vvar_id, loc in srda_model.all_vvar_definitions.items():
+            if isinstance(loc, ExternalCodeLocation):
+                uses = srda_model.all_vvar_uses.get(vvar_id, [])
+                if len(uses) == 1:
+                    vvar, used_loc = next(iter(uses))
+                    if vvar is not None and vvar.was_reg:
+                        saved_vvars.add((vvar_id, used_loc))
+
+        if not saved_vvars:
+            return results
+
+        # for each candidate, we check to ensure:
+        # - it is stored onto the stack (into a stack virtual variable)
+        # - the stack virtual variable is only used once and restores the value to the same register
+        # - the restore location is in the dominance frontier of the store location
+        for vvar_id, used_loc in saved_vvars:
+            def_block = self._get_block(used_loc.block_addr, idx=used_loc.block_idx)
+            assert def_block is not None and used_loc.stmt_idx is not None
+            stmt = def_block.statements[used_loc.stmt_idx]
+            if not (
+                isinstance(stmt, Assignment)
+                and isinstance(stmt.dst, VirtualVariable)
+                and stmt.dst.was_stack
+                and isinstance(stmt.src, VirtualVariable)
+                and stmt.src.was_reg
+                and stmt.src.varid == vvar_id
+            ):
+                continue
+            stack_vvar = stmt.dst
+            all_stack_vvar_uses = srda_model.all_vvar_uses.get(stack_vvar.varid, [])
+            # eliminate the use location if it's a phi statement
+            stack_vvar_uses = set()
+            for vvar_, loc_ in all_stack_vvar_uses:
+                use_block = self._get_block(loc_.block_addr, idx=loc_.block_idx)
+                if use_block is None or loc_.stmt_idx is None:
+                    continue
+                use_stmt = use_block.statements[loc_.stmt_idx]
+                if is_phi_assignment(use_stmt):
+                    continue
+                stack_vvar_uses.add((vvar_, loc_))
+            if len(stack_vvar_uses) != 1:
+                continue
+            _, stack_vvar_use_loc = next(iter(stack_vvar_uses))
+            restore_block = self._get_block(stack_vvar_use_loc.block_addr, idx=stack_vvar_use_loc.block_idx)
+            assert restore_block is not None
+            restore_stmt = restore_block.statements[stack_vvar_use_loc.stmt_idx]
+
+            if not (
+                isinstance(restore_stmt, Assignment)
+                and isinstance(restore_stmt.src, VirtualVariable)
+                and restore_stmt.src.varid == stack_vvar.varid
+                and isinstance(restore_stmt.dst, VirtualVariable)
+                and restore_stmt.dst.was_reg
+                and restore_stmt.dst.reg_offset == stmt.src.reg_offset
+            ):
+                continue
+            # this is the dumb version of the dominance frontier check
+            if self._within_dominance_frontier(def_block, restore_block, True, True):
+                results.append(
+                    (stmt.src.varid, used_loc, stack_vvar.varid, stack_vvar_use_loc, stack_vvar.stack_offset)
+                )
+
+        return results
+
+    def _within_dominance_frontier(self, dom_node, node, use_preds: bool, use_succs: bool) -> bool:
+        if use_succs:
+            # scan forward
+            succs = [succ for succ in self._graph.successors(dom_node) if succ is not dom_node]
+            if len(succs) == 1:
+                succ = succs[0]
+                succ_preds = [pred for pred in self._graph.predecessors(succ) if pred is not succ]
+                if len(succ_preds) == 0:
+                    # the successor has no other predecessors
+                    r = self._within_dominance_frontier(succ, node, False, True)
+                    if r:
+                        return True
+
+                else:
+                    # the successor has other predecessors; gotta step back
+                    preds = [pred for pred in self._graph.predecessors(node) if pred is not node]
+                    if len(preds) == 1 and preds[0] is node:
+                        return True
+            elif len(succs) == 2:
+                return any(succ is node for succ in succs)
+
+        if use_preds:
+            # scan backward
+            preds = [pred for pred in self._graph.predecessors(dom_node) if pred is not dom_node]
+            if len(preds) == 1:
+                pred = preds[0]
+                pred_succs = [succ for succ in self._graph.successors(pred) if succ is not pred]
+                if len(pred_succs) == 0:
+                    # the predecessor has no other successors
+                    return self._within_dominance_frontier(pred, node, True, False)
+
+                # the predecessor has other successors; gotta step forward
+                succs = [succ for succ in self._graph.successors(node) if succ is not node]
+                if len(succs) == 1:
+                    return self._graph.has_edge(node, succs[0])
+            elif len(preds) == 2:
+                return False
+
+        return False