angr 9.2.165__cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

angr/analyses/decompiler/optimization_passes/return_duplicator_low.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+import logging
+import inspect
+from typing import Any
+
+import networkx
+
+from angr.ailment import Block
+from angr.ailment.statement import ConditionalJump, Label
+
+from .return_duplicator_base import ReturnDuplicatorBase
+from .optimization_pass import StructuringOptimizationPass
+
+_l = logging.getLogger(name=__name__)
+
+
+class ReturnDuplicatorLow(StructuringOptimizationPass, ReturnDuplicatorBase):
+    """
+    An optimization pass that reverts a subset of Irreducible Statement Condensing (ISC) optimizations, as described
+    in the USENIX 2024 paper SAILR. This is the heavy/goto version of the ReturnDuplicator optimization pass.
+
+    Some compilers, including GCC, Clang, and MSVC, apply various optimizations to reduce the number of statements in
+    code. These optimizations will take equivalent statements, or a subset of them, and replace them with a single
+    copy that is jumped to by gotos -- optimizing for space and sometimes speed.
+
+    This optimization pass will revert those gotos by re-duplicating the condensed blocks. Since Return statements
+    are the most common, we use this optimization pass to revert only gotos to return statements. Additionally, we
+    perform some additional readability fixups, like not re-duplicating returns to shared components.
+
+    Args:
+        func: The function to optimize.
+        node_idx_start: The index to start at when creating new nodes. This is used by Clinic to ensure that
+            node indices are unique across multiple passes.
+        max_opt_iters: The maximum number of optimization iterations to perform.
+        max_calls_in_regions: The maximum number of calls that can be in a region. This is used to prevent
+            duplicating too much code.
+        prevent_new_gotos: If True, this optimization pass will prevent new gotos from being created.
+        minimize_copies_for_regions: If True, this optimization pass will minimize the number of copies by doing only
+            a single copy for connected in_edges that form a region.
+    """
+
+    ARCHES = None
+    PLATFORMS = None
+    NAME = "Duplicate returns connect with gotos (low)"
+    DESCRIPTION = inspect.cleandoc(__doc__[: __doc__.index("Args:")])  # pylint:disable=unsubscriptable-object
+
+    def __init__(
+        self,
+        func,
+        # settings
+        max_opt_iters: int = 4,
+        max_calls_in_regions: int = 2,
+        prevent_new_gotos: bool = True,
+        minimize_copies_for_regions: bool = True,
+        region_identifier=None,
+        vvar_id_start: int | None = None,
+        scratch: dict[str, Any] | None = None,
+        max_func_blocks: int = 500,
+        **kwargs,
+    ):
+        StructuringOptimizationPass.__init__(
+            self,
+            func,
+            max_opt_iters=max_opt_iters,
+            prevent_new_gotos=prevent_new_gotos,
+            require_gotos=True,
+            vvar_id_start=vvar_id_start,
+            scratch=scratch,
+            region_identifier=region_identifier,
+            **kwargs,
+        )
+        ReturnDuplicatorBase.__init__(
+            self,
+            func,
+            max_calls_in_regions=max_calls_in_regions,
+            minimize_copies_for_regions=minimize_copies_for_regions,
+            ri=region_identifier,
+            vvar_id_start=vvar_id_start,
+            scratch=scratch,
+            max_func_blocks=max_func_blocks,
+        )
+        self.analyze()
+
+    def _check(self):
+        return ReturnDuplicatorBase._check(self)
+
+    def _should_duplicate_dst(self, src, dst, graph, dst_is_const_ret=False):
+        return self._is_goto_edge(src, dst, graph=graph)
+
+    def _is_goto_edge(
+        self,
+        src: Block,
+        dst: Block,
+        graph: networkx.DiGraph = None,
+        max_level_check=1,
+    ):
+        """
+        TODO: Implement a more principled way of checking if an edge is a goto edge with Phoenix's structuring info
+        This function only exists because a long-standing bug that sometimes reports the if-stmt addr
+        above a goto edge as the goto src.
+        """
+        # Do a simple and fast check first
+        is_simple_goto = self._goto_manager.is_goto_edge(src, dst)
+        if is_simple_goto:
+            return True
+
+        if graph is not None:
+            # Special case 1:
+            # We need to check for predecessors above the goto and see if they are a goto.
+            # This needs to include Jump to deal with loops.
+            blocks = [src]
+            level_blocks = [src]
+            for _ in range(max_level_check):
+                new_level_blocks = []
+                for lblock in level_blocks:
+                    new_level_blocks += list(graph.predecessors(lblock))
+
+                blocks += new_level_blocks
+                level_blocks = new_level_blocks
+
+            src_direct_parents = list(graph.predecessors(src))
+            for block in blocks:
+                if not block or not block.statements:
+                    continue
+
+                # special case if-stmts that are next to each other
+                if block in src_direct_parents and isinstance(block.statements[-1], ConditionalJump):
+                    continue
+
+                if self._goto_manager.is_goto_edge(block, dst):
+                    return True
+
+            # Special case 2: A "goto edge" that ReturnDuplicator wants to test might be an edge that Phoenix
+            # includes in its loop region (during the cyclic refinement). In fact, Phoenix tends to include as many
+            # nodes as possible into the loop region, and generate a goto edge (which ends up in the structured code)
+            # from `dst` to the loop successor.
+            # an example of this is captured by the test case `TestDecompiler.test_stty_recover_mode_ret_dup_region`.
+            # until someone (ideally @mahaloz) implements a more principled way of translating "goto statements" that
+            # Phoenix generates and "goto edges" that ReturnDuplicator tests, we rely on the following stopgap to
+            # handle this case.
+            node = dst
+            while True:
+                succs = list(graph.successors(node))
+                if len(succs) != 1:
+                    break
+                succ = succs[0]
+                if succ is node:
+                    # loop!
+                    break
+                succ_preds = list(graph.predecessors(succ))
+                if len(succ_preds) != 1:
+                    break
+                if self._goto_manager.is_goto_edge(node, succ):
+                    return True
+                # keep testing the next edge
+                node = succ
+
+            # Special case 3: In Phoenix, regions full of only if-stmts can be collapsed and moved. This causes
+            # the goto manager to report gotos that are at the top of the region instead of ones in the middle of it.
+            # Because of this, we need to gather all the nodes above the original src and check if any of them
+            # go to the destination. Additionally, we need to do this on the supergraph to get rid of
+            # goto edges that are removed by Phoenix.
+            # This case is observed in the test case `TestDecompiler.test_tail_tail_bytes_ret_dup`.
+            if self._supergraph is None:
+                return False
+
+            super_to_og_nodes = {n: self._supergraph.nodes[n]["original_nodes"] for n in self._supergraph.nodes}
+            og_to_super_nodes = {og: super_n for super_n, ogs in super_to_og_nodes.items() for og in ogs}
+            super_src = og_to_super_nodes.get(src)
+            super_dst = og_to_super_nodes.get(dst)
+            if super_src is None or super_dst is None:
+                return False
+
+            # collect all nodes which have only an if-stmt in them that are ancestors of super_src
+            check_blks = {super_src}
+            level_blocks = {super_src}
+            for _ in range(10):
+                done = False
+                if_blks = set()
+                for lblock in level_blocks:
+                    preds = list(self._supergraph.predecessors(lblock))
+                    for pred in preds:
+                        only_cond_jump = all(isinstance(s, (ConditionalJump, Label)) for s in pred.statements)
+                        if only_cond_jump:
+                            if_blks.add(pred)
+
+                    done = len(if_blks) == 0
+
+                if done:
+                    break
+
+                check_blks |= if_blks
+                level_blocks = if_blks
+
+            # convert all the found if-only super-blocks back into their original blocks
+            og_check_blocks = set()
+            for blk in check_blks:
+                og_check_blocks |= set(super_to_og_nodes[blk])
+
+            # check if any of the original blocks are gotos to the destination
+            goto_hits = 0
+            for block in og_check_blocks:
+                if self._goto_manager.is_goto_edge(block, dst):
+                    goto_hits += 1
+
+            # Although it is good to find a goto in the if-only block region, having more than a single goto
+            # existing that goes to the same dst is a bad sign. This can be seen in the the following test:
+            # TestDecompiler.test_dd_iread_ret_dup_region
+            #
+            # It occurs when you have something like:
+            # ```
+            # if (a || c)
+            #     goto target;
+            # target:
+            # return 0;
+            # ```
+            #
+            #
+            # This looks like an edge from (a, target) and (c, target) but it is actually a single edge.
+            # If you allow both to duplicate you get the following:
+            # ```
+            # if (a):
+            #    return
+            # if (c):
+            #    return
+            # ```
+            # This is not the desired behavior.
+            # So we need to check if there is only a single goto that goes to the destination.
+            return goto_hits == 1
+
+        return False
+
+    def _analyze(self, cache=None):
+        """
+        This analysis is run in a loop in analyze() for a maximum of max_opt_iters times.
+        """
+        return self._analyze_core(self.out_graph)

angr/analyses/decompiler/optimization_passes/stack_canary_simplifier.py

@@ -0,0 +1,333 @@
+# pylint:disable=too-many-boolean-expressions
+from __future__ import annotations
+from collections import defaultdict
+import logging
+
+import angr.ailment as ailment
+
+from angr.utils.bits import s2u
+from angr.analyses.decompiler.stack_item import StackItem, StackItemType
+from .optimization_pass import OptimizationPass, OptimizationPassStage
+
+
+_l = logging.getLogger(name=__name__)
+
+
+class StackCanarySimplifier(OptimizationPass):
+    """
+    Removes stack canary checks from decompilation results.
+    """
+
+    ARCHES = [
+        "X86",
+        "AMD64",
+    ]  # TODO: fs is x86 only. Figure out how stack canary is loaded in other architectures
+    PLATFORMS = ["cgc", "linux"]
+    STAGE = OptimizationPassStage.AFTER_GLOBAL_SIMPLIFICATION
+    NAME = "Simplify stack canaries"
+    DESCRIPTION = __doc__.strip()
+
+    def __init__(self, func, **kwargs):
+        super().__init__(func, **kwargs)
+
+        self.analyze()
+
+    def _check(self):
+        # Check the first block and see if there is any statement reading data from fs:0x28h
+        init_stmt = self._find_canary_init_stmt()
+
+        return init_stmt is not None, {"init_stmt": init_stmt}
+
+    def _analyze(self, cache=None):
+        init_stmt = None
+        if cache is not None:
+            init_stmt = cache.get("init_stmt", None)
+
+        if init_stmt is None:
+            init_stmt = self._find_canary_init_stmt()
+
+        if init_stmt is None:
+            return
+
+        # Look for the statement that loads back canary value from the stack
+        first_block, stmt_idx = init_stmt
+        canary_init_stmt = first_block.statements[stmt_idx]
+        # where is the stack canary stored?
+        if not (isinstance(canary_init_stmt.dst, ailment.Expr.VirtualVariable) and canary_init_stmt.dst.was_stack):
+            _l.debug(
+                "Unsupported canary storing location %s. Expects a stack VirtualVariable.",
+                canary_init_stmt.addr,
+            )
+            return
+
+        store_offset = canary_init_stmt.dst.stack_offset
+        if not isinstance(store_offset, int):
+            _l.debug("Unsupported canary storing offset %s. Expects an int.", store_offset)
+
+        # The function should have at least one end point with an if-else statement
+        # Find all nodes with 0 out-degrees
+        all_endpoint_addrs = [node.addr for node in self._func.graph.nodes() if self._func.graph.out_degree(node) == 0]
+
+        # Before node duplication, each pair of canary-check-success and canary-check-failure nodes have a common
+        # predecessor.
+        # map endpoint addrs to their common predecessors
+        pred_addr_to_endpoint_addrs: dict[int, set[int]] = defaultdict(set)
+        for node_addr in all_endpoint_addrs:
+            preds = self._func.graph.predecessors(self._func.get_node(node_addr))
+            for pred in preds:
+                pred_addr_to_endpoint_addrs[pred.addr].add(node_addr)
+
+        found_endpoints = False
+        for pred_addr in pred_addr_to_endpoint_addrs:
+            endpoint_addrs = pred_addr_to_endpoint_addrs[pred_addr]
+
+            if len(endpoint_addrs) != 2:
+                # we expect there to be only two nodes: one for canary-check-success, and the other for
+                # canary-check-failure. if not, we check the next predecessor
+                continue
+
+            # because other optimization passes may duplicate nodes, we may have more than one node for each function
+            # endpoint.
+            endpoint_addrs_list = list(endpoint_addrs)
+            endnodes_0 = list(self._get_blocks(endpoint_addrs_list[0]))
+            endnodes_1 = list(self._get_blocks(endpoint_addrs_list[1]))
+
+            if not endnodes_0 or not endnodes_1:
+                _l.warning("Unexpected situation: endnodes_0 or endnodes_1 is empty.")
+                continue
+
+            # One of the end nodes calls __stack_chk_fail
+            stack_chk_fail_callers = None
+            other_nodes = None
+            for endnodes, o in [(endnodes_0, endnodes_1), (endnodes_1, endnodes_0)]:
+                if self._calls_stack_chk_fail(endnodes[0]):
+                    stack_chk_fail_callers = endnodes
+                    other_nodes = o
+                    break
+            else:
+                _l.debug("Cannot find the node that calls __stack_chk_fail().")
+                continue
+
+            # Match stack_chk_fail_caller, ret_node, and predecessor
+            nodes_to_process = []
+            for stack_chk_fail_caller in stack_chk_fail_callers:
+                all_preds = set(self._graph.predecessors(stack_chk_fail_caller))
+                preds_for_other_nodes = set()
+                for o in other_nodes:
+                    preds_for_other_nodes |= set(self._graph.predecessors(o))
+                preds = list(all_preds.intersection(preds_for_other_nodes))
+                if len(preds) != 1:
+                    _l.debug("Expect 1 predecessor. Found %d.", len(preds))
+                    continue
+                pred = preds[0]
+
+                # More sanity checks
+                if len(pred.statements) < 1:
+                    _l.debug("The predecessor node is empty.")
+                    continue
+
+                # Check the last statement
+                if not isinstance(pred.statements[-1], ailment.Stmt.ConditionalJump):
+                    _l.debug("The predecessor does not end with a conditional jump.")
+                    continue
+
+                # Find the statement that computes real canary value xor stored canary value
+                canary_check_stmt_idx = self._find_canary_comparison_stmt(pred, store_offset)
+                if canary_check_stmt_idx is None:
+                    _l.debug("Cannot find the canary check statement in the predecessor.")
+                    continue
+
+                succs = list(self._graph.successors(pred))
+                if len(succs) != 2:
+                    _l.debug("Expect 2 successors. Found %d.", len(succs))
+                    continue
+                ret_node = succs[1] if stack_chk_fail_caller is succs[0] else succs[0]
+                nodes_to_process.append((pred, canary_check_stmt_idx, stack_chk_fail_caller, ret_node))
+
+            # Awesome. Now patch this function.
+            for pred, _canary_check_stmt_idx, stack_chk_fail_caller, ret_node in nodes_to_process:
+                # Patch the pred so that it jumps to the one that is not stack_chk_fail_caller
+                pred_copy = pred.copy()
+                pred_copy.statements[-1] = ailment.Stmt.Jump(
+                    len(pred_copy.statements) - 1,
+                    ailment.Expr.Const(None, None, ret_node.addr, self.project.arch.bits),
+                    ins_addr=pred_copy.statements[-1].ins_addr,
+                )
+
+                self._graph.remove_edge(pred, stack_chk_fail_caller)
+                self._update_block(pred, pred_copy)
+
+                if self._graph.in_degree[stack_chk_fail_caller] == 0:
+                    # Remove the block that calls stack_chk_fail_caller
+                    self._remove_block(stack_chk_fail_caller)
+
+                found_endpoints = True
+
+        if found_endpoints:
+            # Remove the statement that loads the stack canary from fs
+            first_block_copy = first_block.copy()
+            first_block_copy.statements.pop(stmt_idx)
+            self._update_block(first_block, first_block_copy)
+
+            # update stack_items
+            self.stack_items[store_offset] = StackItem(
+                store_offset, canary_init_stmt.dst.size, "canary", StackItemType.STACK_CANARY
+            )
+
+        # Done!
+
+    def _find_canary_init_stmt(self):
+        block_addr = self._func.addr
+        traversed = set()
+
+        while True:
+            traversed.add(block_addr)
+            try:
+                first_block = next(self._get_blocks(block_addr))
+            except StopIteration:
+                break
+
+            if first_block is None:
+                break
+
+            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_stack
+                    and isinstance(stmt.dst.stack_offset, int)
+                    and isinstance(stmt.src, ailment.Expr.Load)
+                    and self._is_add(stmt.src.addr)
+                ):
+                    # Check addr: must be fs+0x28
+                    op0, op1 = stmt.src.addr.operands
+                    if isinstance(op1, ailment.Expr.VirtualVariable) and op1.was_reg:
+                        op0, op1 = op1, op0
+                    if (
+                        isinstance(op0, ailment.Expr.VirtualVariable)
+                        and op0.was_reg
+                        and isinstance(op1, ailment.Expr.Const)
+                        and op0.reg_offset == self.project.arch.get_register_offset("fs")
+                        and op1.value == 0x28
+                    ):
+                        return first_block, idx
+
+            succs = list(self._graph.successors(first_block))
+            if len(succs) == 1:
+                block_addr = succs[0].addr
+                if block_addr not in traversed:
+                    continue
+            break
+
+        return None
+
+    def _find_canary_comparison_stmt(self, block, canary_value_stack_offset):
+        for idx, stmt in enumerate(block.statements):
+            if isinstance(stmt, ailment.Stmt.ConditionalJump):
+                if isinstance(stmt.condition, ailment.Expr.UnaryOp) and stmt.condition.op == "Not":
+                    # !(s_10 ^ fs:0x28 == 0)
+                    negated = True
+                    condition = stmt.condition.operand
+                else:
+                    negated = False
+                    condition = stmt.condition
+                if isinstance(condition, ailment.Expr.BinaryOp) and (
+                    (not negated and condition.op == "CmpEQ") or (negated and condition.op == "CmpNE")
+                ):
+                    pass
+                else:
+                    continue
+
+                expr0, expr1 = condition.operands
+                if isinstance(expr0, ailment.Expr.BinaryOp) and expr0.op == "Xor":
+                    # a ^ b
+                    op0, op1 = expr0.operands
+                    if isinstance(op0, ailment.Expr.VirtualVariable) and op0.was_reg:
+                        # maybe op0 holds the value of another stack variable, like the following:
+                        #
+                        # 00 | 0x404e75 | LABEL_404e75:
+                        # 01 | 0x404e75 | vvar_62{reg 16} = vvar_79{stack -64}
+                        # 02 | 0x404e7a | vvar_66{reg 16} = (vvar_62{reg 16} ^ Load(addr=(0x28<64> + vvar_31{reg 208}),
+                        #                 size=8, endness=Iend_LE))
+                        # 03 | 0x404e83 | if (((vvar_62{reg 16} ^ Load(addr=(0x28<64> + vvar_31{reg 208}), size=8,
+                        #                 endness=Iend_LE)) == 0x0<64>)) { Goto ... } else { Goto ... }
+                        op0_v = self._get_vvar_value(block, op0.varid)
+                        if isinstance(op0_v, ailment.Expr.VirtualVariable) and op0_v.was_stack:
+                            op0 = op0_v
+
+                    if not (
+                        (
+                            self._is_stack_canary_load_expr(op0, self.project.arch.bits, canary_value_stack_offset)
+                            and self._is_random_number_load_expr(op1, self.project.arch.get_register_offset("fs"))
+                        )
+                        or (
+                            self._is_stack_canary_load_expr(op1, self.project.arch.bits, canary_value_stack_offset)
+                            and self._is_random_number_load_expr(op0, self.project.arch.get_register_offset("fs"))
+                        )
+                    ):
+                        continue
+                elif (
+                    isinstance(expr0, ailment.Expr.Load)
+                    and isinstance(expr1, ailment.Expr.Load)
+                    and condition.op == "CmpEQ"
+                ):
+                    # a == b
+                    if not (
+                        (
+                            self._is_stack_canary_load_expr(expr0, self.project.arch.bits, canary_value_stack_offset)
+                            and self._is_random_number_load_expr(expr1, self.project.arch.get_register_offset("fs"))
+                        )
+                        or (
+                            self._is_stack_canary_load_expr(expr1, self.project.arch.bits, canary_value_stack_offset)
+                            and self._is_random_number_load_expr(expr0, self.project.arch.get_register_offset("fs"))
+                        )
+                    ):
+                        continue
+
+                # Found it
+                return idx
+
+        return None
+
+    def _calls_stack_chk_fail(self, node):
+        for stmt in node.statements:
+            if isinstance(stmt, ailment.Stmt.Call) and isinstance(stmt.target, ailment.Expr.Const):
+                const_target = stmt.target.value
+                if const_target in self.kb.functions:
+                    func = self.kb.functions.function(addr=const_target)
+                    if func.name == "__stack_chk_fail":
+                        return True
+
+        return False
+
+    @staticmethod
+    def _is_stack_canary_load_expr(expr, bits: int, canary_value_stack_offset: int) -> bool:
+        if not (
+            isinstance(expr, ailment.Expr.VirtualVariable) and expr.was_stack and isinstance(expr.stack_offset, int)
+        ):
+            return False
+        return s2u(expr.stack_offset, bits) == s2u(canary_value_stack_offset, bits)
+
+    @staticmethod
+    def _is_random_number_load_expr(expr, fs_reg_offset: int) -> bool:
+        return (
+            isinstance(expr, ailment.Expr.Load)
+            and isinstance(expr.addr, ailment.Expr.BinaryOp)
+            and expr.addr.op == "Add"
+            and isinstance(expr.addr.operands[0], ailment.Expr.Const)
+            and expr.addr.operands[0].value == 0x28
+            and isinstance(expr.addr.operands[1], ailment.Expr.VirtualVariable)
+            and expr.addr.operands[1].was_reg
+            and expr.addr.operands[1].reg_offset == fs_reg_offset
+        )
+
+    @staticmethod
+    def _get_vvar_value(block: ailment.Block, vvar_id: int) -> ailment.Expression | None:
+        for stmt in block.statements:
+            if (
+                isinstance(stmt, ailment.Stmt.Assignment)
+                and isinstance(stmt.dst, ailment.Expr.VirtualVariable)
+                and stmt.dst.varid == vvar_id
+            ):
+                return stmt.src
+        return None