angr 9.2.117__py3-none-manylinux2014_x86_64.whl → 9.2.118__py3-none-manylinux2014_x86_64.whl

angr/analyses/decompiler/optimization_passes/inlined_string_transformation_simplifier.py

@@ -1,11 +1,21 @@
 # pylint:disable=arguments-renamed,too-many-boolean-expressions,no-self-use
 from __future__ import annotations
-from typing import Any, DefaultDict
+from typing import Any
 from collections import defaultdict
 
 from archinfo import Endness
-from ailment.expression import Const, Register, Load, StackBaseOffset, Convert, BinaryOp
-from ailment.statement import Store, ConditionalJump, Jump
+from ailment.expression import (
+    Const,
+    Register,
+    Load,
+    StackBaseOffset,
+    Convert,
+    BinaryOp,
+    VirtualVariable,
+    Phi,
+    VirtualVariableCategory,
+)
+from ailment.statement import ConditionalJump, Jump, Assignment
 import claripy
 
 from angr.engines.light import SimEngineLightAILMixin
@@ -42,6 +52,7 @@ class InlinedStringTransformationState:
 
         self.registers = FasterMemory(memory_id="reg")
         self.memory = FasterMemory(memory_id="mem")
+        self.virtual_variables = {}
 
         self.registers.set_state(self)
         self.memory.set_state(self)
@@ -71,6 +82,14 @@ class InlinedStringTransformationState:
         except SimMemoryMissingError:
             return None
 
+    def vvar_store(self, vvar: VirtualVariable, value: claripy.Bits | None) -> None:
+        self.virtual_variables[vvar.varid] = value
+
+    def vvar_load(self, vvar: VirtualVariable) -> claripy.Bits | None:
+        if vvar.varid in self.virtual_variables:
+            return self.virtual_variables[vvar.varid]
+        return None
+
 
 class InlinedStringTransformationAILEngine(SimEngineLightAILMixin):
     """
@@ -90,10 +109,11 @@ class InlinedStringTransformationAILEngine(SimEngineLightAILMixin):
         self.MASK = 0xFFFF_FFFF if self.arch.bits == 32 else 0xFFFF_FFFF_FFFF_FFFF
 
         state = InlinedStringTransformationState(project)
-        self.stack_accesses: DefaultDict[int, list[tuple[str, CodeLocation, claripy.Bits]]] = defaultdict(list)
+        self.stack_accesses: defaultdict[int, list[tuple[str, CodeLocation, claripy.Bits]]] = defaultdict(list)
         self.finished: bool = False
 
         i = 0
+        self.last_pc = None
         self.pc = self.start
         while i < self.step_limit:
             if self.pc not in self.nodes:
@@ -125,10 +145,22 @@ class InlinedStringTransformationAILEngine(SimEngineLightAILMixin):
         return None
 
     def _handle_Assignment(self, stmt):
-        if isinstance(stmt.dst, Register):
-            val = self._expr(stmt.src)
-            if isinstance(val, claripy.Bits):
-                self.state.reg_store(stmt.dst, val)
+        if isinstance(stmt.dst, VirtualVariable):
+            if stmt.dst.was_reg:
+                val = self._expr(stmt.src)
+                if isinstance(val, claripy.Bits):
+                    self.state.vvar_store(stmt.dst, val)
+            elif stmt.dst.was_stack:
+                addr = (stmt.dst.stack_offset + self.STACK_BASE) & self.MASK
+                val = self._expr(stmt.src)
+                if isinstance(val, claripy.ast.BV):
+                    self.state.mem_store(addr, val, self.arch.memory_endness)
+                    # log it
+                    for i in range(val.size() // self.arch.byte_width):
+                        byte_off = i
+                        if self.arch.memory_endness == Endness.LE:
+                            byte_off = val.size() // self.arch.byte_width - i - 1
+                        self.stack_accesses[addr + i].append(("store", self._codeloc(), val.get_byte(byte_off)))
 
     def _handle_Store(self, stmt):
         addr_and_type = self._process_address(stmt.addr)
@@ -139,19 +171,21 @@ class InlinedStringTransformationAILEngine(SimEngineLightAILMixin):
                 self.state.mem_store(addr, val, stmt.endness)
                 # log it
                 if addr_type == "stack":
-                    for i in range(0, val.size() // self.arch.byte_width):
+                    for i in range(val.size() // self.arch.byte_width):
                         byte_off = i
                         if self.arch.memory_endness == Endness.LE:
                             byte_off = val.size() // self.arch.byte_width - i - 1
                         self.stack_accesses[addr + i].append(("store", self._codeloc(), val.get_byte(byte_off)))
 
     def _handle_Jump(self, stmt):
+        self.last_pc = self.pc
         if isinstance(stmt.target, Const):
             self.pc = stmt.target.value
         else:
             self.pc = None
 
     def _handle_ConditionalJump(self, stmt):
+        self.last_pc = self.pc
         self.pc = None
         if isinstance(stmt.true_target, Const) and isinstance(stmt.false_target, Const):
             cond = self._expr(stmt.condition)
@@ -171,7 +205,7 @@ class InlinedStringTransformationAILEngine(SimEngineLightAILMixin):
             v = self.state.mem_load(addr, expr.size, expr.endness)
             # log it
             if addr_type == "stack" and isinstance(v, claripy.ast.BV):
-                for i in range(0, expr.size):
+                for i in range(expr.size):
                     byte_off = i
                     if self.arch.memory_endness == Endness.LE:
                         byte_off = expr.size - i - 1
@@ -182,6 +216,28 @@ class InlinedStringTransformationAILEngine(SimEngineLightAILMixin):
     def _handle_Register(self, expr: Register):
         return self.state.reg_load(expr)
 
+    def _handle_VirtualVariable(self, expr: VirtualVariable):
+        if expr.was_stack:
+            addr = (expr.stack_offset + self.STACK_BASE) & self.MASK
+            v = self.state.mem_load(addr, expr.size, self.arch.memory_endness)
+            if isinstance(v, claripy.Bits):
+                # log it
+                for i in range(expr.size):
+                    byte_off = i
+                    if self.arch.memory_endness == Endness.LE:
+                        byte_off = expr.size - i - 1
+                    self.stack_accesses[addr + i].append(("load", self._codeloc(), v.get_byte(byte_off)))
+            return v
+        if expr.was_reg:
+            return self.state.vvar_load(expr)
+        return None
+
+    def _handle_Phi(self, expr: Phi):
+        for src, vvar in expr.src_and_vvars:
+            if src[0] == self.last_pc and vvar is not None:
+                return self.state.vvar_load(vvar)
+        return None
+
     def _handle_Convert(self, expr: Convert):
         v = self._expr(expr.operand)
         if isinstance(v, claripy.Bits):
@@ -189,10 +245,9 @@ class InlinedStringTransformationAILEngine(SimEngineLightAILMixin):
                 if not expr.is_signed:
                     return claripy.ZeroExt(expr.to_bits - expr.from_bits, v)
                 return claripy.SignExt(expr.to_bits - expr.from_bits, v)
-            elif expr.to_bits < expr.from_bits:
+            if expr.to_bits < expr.from_bits:
                 return claripy.Extract(expr.to_bits - 1, 0, v)
-            else:
-                return v
+            return v
         return None
 
     def _handle_CmpEQ(self, expr):
@@ -292,17 +347,22 @@ class InlinedStringTransformationSimplifier(OptimizationPass):
             store_statements = []
             for off, stack_accesses in enumerate(desc.stack_accesses):
                 # the last element is the final storing statement
-                stack_addr = StackBaseOffset(None, self.project.arch.bits, desc.beginning_stack_offset + off)
                 new_value_ast = stack_accesses[-1][2]
                 new_value = Const(None, None, new_value_ast.concrete_value, self.project.arch.byte_width)
-                stmt = Store(
+                stmt = Assignment(
                     None,
-                    stack_addr,
+                    VirtualVariable(
+                        None,
+                        self.vvar_id_start,
+                        self.project.arch.bits,
+                        category=VirtualVariableCategory.STACK,
+                        oident=desc.beginning_stack_offset + off,
+                        ins_addr=desc.store_block.addr + desc.store_block.original_size - 1,
+                    ),
                     new_value,
-                    1,
-                    "Iend_LE",
                     ins_addr=desc.store_block.addr + desc.store_block.original_size - 1,
                 )
+                self.vvar_id_start += 1
                 store_statements.append(stmt)
             if new_statements and isinstance(new_statements[-1], (ConditionalJump, Jump)):
                 new_statements = new_statements[:-1] + store_statements + new_statements[-1:]
@@ -368,11 +428,16 @@ class InlinedStringTransformationSimplifier(OptimizationPass):
                     stack_accesses = engine.stack_accesses[stack_addr]
                     if len(stack_accesses) == 3:
                         item0, item1, item2 = stack_accesses
-                        if item0[0] == "store" and item1[0] == "load" and item2[0] == "store":
-                            if item0[1] != item1[1] and item1[1] == item2[1]:
-                                if item0[2] is item1[2]:
-                                    # found one!
-                                    candidate_stack_addrs.append(stack_addr)
+                        if (
+                            item0[0] == "store"
+                            and item1[0] == "load"
+                            and item2[0] == "store"
+                            and item0[1] != item1[1]
+                            and item1[1] == item2[1]
+                            and item0[2] is item1[2]
+                        ):
+                            # found one!
+                            candidate_stack_addrs.append(stack_addr)
 
                 if (
                     len(candidate_stack_addrs) >= 2

angr/analyses/decompiler/optimization_passes/ite_expr_converter.py

@@ -1,6 +1,7 @@
 # pylint:disable=unnecessary-pass
+from __future__ import annotations
 import logging
-from typing import Optional, Any, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 from ailment.statement import ConditionalJump, Assignment, Statement
 from ailment.expression import Const, ITE, Expression
@@ -25,8 +26,6 @@ class NodeFoundNotification(Exception):
     A notification that the target node has been found.
     """
 
-    pass
-
 
 class BlockLocator(RegionWalker):
     """
@@ -44,7 +43,7 @@ class BlockLocator(RegionWalker):
     def walk_node(self, region, node):
         if node == self._block:
             self.region = region
-            raise NodeFoundNotification()
+            raise NodeFoundNotification
 
 
 class ExpressionReplacer(AILBlockWalker):
@@ -59,11 +58,10 @@ class ExpressionReplacer(AILBlockWalker):
         self._callback = callback
 
     def _handle_expr(
-        self, expr_idx: int, expr: Expression, stmt_idx: int, stmt: Statement | None, block: Optional["AILBlock"]
+        self, expr_idx: int, expr: Expression, stmt_idx: int, stmt: Statement | None, block: AILBlock | None
     ) -> Any:
         if expr == self._target_expr:
-            new_expr = self._callback(self._block_addr, stmt_idx, stmt.ins_addr, expr)
-            return new_expr
+            return self._callback(self._block_addr, stmt_idx, stmt.ins_addr, expr)
         return super()._handle_expr(expr_idx, expr, stmt_idx, stmt, block)
 
 
@@ -129,10 +127,10 @@ class ITEExprConverter(OptimizationPass):
         #
 
         # find their regions
-        block_0 = [b for b in blocks if b.addr == defs[0].codeloc.block_addr][0]
+        block_0 = next(b for b in blocks if b.addr == defs[0].codeloc.block_addr)
         region_0 = self._locate_block(block_0)
 
-        block_1 = [b for b in blocks if b.addr == defs[1].codeloc.block_addr][0]
+        block_1 = next(b for b in blocks if b.addr == defs[1].codeloc.block_addr)
         region_1 = self._locate_block(block_1)
 
         if region_0 is None or region_1 is None or region_0 != region_1:
@@ -217,7 +215,7 @@ class ITEExprConverter(OptimizationPass):
 
         return new_expr
 
-    def _locate_block(self, block: "AILBlock"):
+    def _locate_block(self, block: AILBlock):
         locator = BlockLocator(block)
         try:
             locator.walk(self._ri.region)

angr/analyses/decompiler/optimization_passes/ite_region_converter.py

@@ -1,9 +1,12 @@
 # pylint:disable=unnecessary-pass
+from __future__ import annotations
 import logging
 
-from ailment.statement import ConditionalJump, Assignment, Jump
-from ailment.expression import ITE, Const
+from ailment.block import Block
+from ailment.statement import Statement, Call, ConditionalJump, Assignment, Jump
+from ailment.expression import ITE, Const, VirtualVariable, Phi
 
+from angr.utils.ail import is_phi_assignment
 from ....utils.graph import subgraph_between_nodes
 from ..utils import remove_labels, to_ail_supergraph
 from .optimization_pass import OptimizationPass, OptimizationPassStage
@@ -39,8 +42,10 @@ class ITERegionConverter(OptimizationPass):
             if not ite_assign_regions:
                 break
 
-            for region_head, region_tail, true_stmt, false_stmt in ite_assign_regions:
-                round_update |= self._convert_region_to_ternary_expr(region_head, region_tail, true_stmt, false_stmt)
+            for region_head, region_tail, true_block, true_stmt, false_block, false_stmt in ite_assign_regions:
+                round_update |= self._convert_region_to_ternary_expr(
+                    region_head, region_tail, true_block, true_stmt, false_block, false_stmt
+                )
 
             if not round_update:
                 break
@@ -52,6 +57,7 @@ class ITERegionConverter(OptimizationPass):
 
     def _find_ite_assignment_regions(self):
         # find all the if-stmt blocks in a graph with no single successor edges
+        blocks_by_end_addr = {(block.addr + block.original_size, block.idx): block for block in self._graph.nodes()}
         super_graph = to_ail_supergraph(remove_labels(self._graph))
         if_stmt_blocks = []
         for node in super_graph.nodes():
@@ -105,11 +111,7 @@ class ITERegionConverter(OptimizationPass):
 
             true_stmt = true_stmts[0]
             false_stmt = false_stmts[0]
-            if (
-                not isinstance(true_stmt, Assignment)
-                or not isinstance(false_stmt, Assignment)
-                or not true_stmt.dst.likes(false_stmt.dst)
-            ):
+            if not self._is_assigning_to_vvar(true_stmt) or not self._is_assigning_to_vvar(false_stmt):
                 continue
 
             # must contain a single common predecessor
@@ -125,19 +127,72 @@ class ITERegionConverter(OptimizationPass):
                 continue
             common_successor = true_successors[0]
 
+            # find the corresponding blocks for true_child and false_child in the original graph
+            # this is because the phi expressions only records source addresses of the original blocks, not the
+            # addresses of super blocks
+            true_child_original = blocks_by_end_addr.get(
+                (true_child.addr + true_child.original_size, true_child.idx), true_child
+            )
+            false_child_original = blocks_by_end_addr.get(
+                (false_child.addr + false_child.original_size, false_child.idx), false_child
+            )
+
+            # the common successor must have a phi assignment with source variables being assigned to in true_stmt and
+            # false_stmt
+            if not self._has_qualified_phi_assignments(
+                common_successor, true_child_original, true_stmt, false_child_original, false_stmt
+            ):
+                continue
+
             # lastly, normalize the region we will be editing
-            region_head = super_to_normal_node.get(if_stmt_block, None)
+            region_head = super_to_normal_node.get(if_stmt_block)
             tail_blocks = list(self.blocks_by_addr.get(common_successor.addr, []))
             region_tail = tail_blocks[0] if tail_blocks else None
             if region_head is None or region_tail is None:
                 continue
 
             # we have now found a valid ITE-like expression case
-            ite_candidates.append((region_head, region_tail, true_stmt, false_stmt))
+            ite_candidates.append((region_head, region_tail, true_child, true_stmt, false_child, false_stmt))
 
         return ite_candidates
 
-    def _convert_region_to_ternary_expr(self, region_head, region_tail, true_stmt, false_stmt):
+    @staticmethod
+    def _has_qualified_phi_assignments(
+        block: Block, block0: Block, stmt0: Assignment | Call, block1: Block, stmt1: Assignment | Call
+    ):
+        vvar0 = stmt0.dst if isinstance(stmt0, Assignment) else stmt0.ret_expr
+        vvar1 = stmt1.dst if isinstance(stmt1, Assignment) else stmt1.ret_expr
+
+        addr0 = block0.addr, block0.idx
+        addr1 = block1.addr, block1.idx
+
+        found_phi_assignment = False
+        has_unexpected_phi_assignment = False
+        for stmt in block.statements:
+            if not is_phi_assignment(stmt):
+                continue
+            src_vars = {src: vvar.varid if vvar is not None else None for src, vvar in stmt.src.src_and_vvars}
+            if src_vars.get(addr0) == vvar0.varid and src_vars.get(addr1) == vvar1.varid:
+                # this is the phi assignment that assigns stmt0.dst and stmt1.dst to a new variable
+                found_phi_assignment = True
+            else:
+                if addr0 in src_vars and addr1 in src_vars and src_vars[addr0] == src_vars[addr1]:
+                    # for all other phi assignments, the source variable out of the two origins must be the same one
+                    pass
+                else:
+                    has_unexpected_phi_assignment = True
+
+        return found_phi_assignment and not has_unexpected_phi_assignment
+
+    def _convert_region_to_ternary_expr(
+        self,
+        region_head,
+        region_tail,
+        true_block,
+        true_stmt: Assignment | Call,
+        false_block,
+        false_stmt: Assignment | Call,
+    ):
         if region_head not in self._graph or region_tail not in self._graph:
             return False
 
@@ -147,18 +202,32 @@ class ITERegionConverter(OptimizationPass):
 
         new_region_head = region_head.copy()
         conditional_jump: ConditionalJump = region_head.statements[-1]
-        addr_obj = true_stmt.src if "ins_addr" in true_stmt.src.tags else true_stmt
+
+        true_stmt_src = true_stmt.src if isinstance(true_stmt, Assignment) else true_stmt
+        true_stmt_dst = true_stmt.dst if isinstance(true_stmt, Assignment) else true_stmt.ret_expr
+        false_stmt_src = false_stmt.src if isinstance(false_stmt, Assignment) else false_stmt
+
+        addr_obj = true_stmt_src if "ins_addr" in true_stmt_src.tags else true_stmt
         ternary_expr = ITE(
             None,
             conditional_jump.condition,
-            true_stmt.src,
-            false_stmt.src,
+            false_stmt_src,
+            true_stmt_src,
             ins_addr=addr_obj.ins_addr,
             vex_block_addr=addr_obj.vex_block_addr,
             vex_stmt_idx=addr_obj.vex_stmt_idx,
         )
-        new_assignment = true_stmt.copy()
-        new_assignment.src = ternary_expr
+        dst = VirtualVariable(
+            true_stmt_dst.idx,
+            self.vvar_id_start,
+            true_stmt_dst.bits,
+            true_stmt_dst.category,
+            oident=true_stmt_dst.oident,
+            **true_stmt_dst.tags,
+        )
+        self.vvar_id_start += 1
+        src = ternary_expr
+        new_assignment = Assignment(true_stmt.idx, dst, src, **true_stmt.tags)
         new_region_head.statements[-1] = new_assignment
 
         # add a goto statement to the region tail so it can be transformed into a break or other types of control-flow
@@ -179,11 +248,52 @@ class ITERegionConverter(OptimizationPass):
 
             self._remove_block(node)
 
+        #
+        # Update phi assignments in region tail
+        #
+
+        stmts = []
+        for stmt in region_tail.statements:
+            if not is_phi_assignment(stmt):
+                stmts.append(stmt)
+                continue
+            new_src_and_vvars = []
+            for src, vvar in stmt.src.src_and_vvars:
+                if src not in {(true_block.addr, true_block.idx), (false_block.addr, false_block.idx)}:
+                    new_src_and_vvars.append((src, vvar))
+            new_vvar = new_assignment.dst.copy()
+            new_src_and_vvars.append(((region_head.addr, region_head.idx), new_vvar))
+
+            new_phi = Phi(
+                stmt.src.idx,
+                stmt.src.bits,
+                new_src_and_vvars,
+                **stmt.src.tags,
+            )
+            new_phi_assignment = Assignment(
+                stmt.idx,
+                stmt.dst,
+                new_phi,
+                **stmt.tags,
+            )
+            stmts.append(new_phi_assignment)
+        new_region_tail = Block(region_tail.addr, region_tail.original_size, statements=stmts, idx=region_tail.idx)
+
         #
         # update head and tail
         #
 
         self._update_block(region_head, new_region_head)
-        self._graph.add_edge(new_region_head, region_tail)
+        self._update_block(region_tail, new_region_tail)
+        self._graph.add_edge(new_region_head, new_region_tail)
 
         return True
+
+    @staticmethod
+    def _is_assigning_to_vvar(stmt: Statement) -> bool:
+        return (
+            isinstance(stmt, Assignment)
+            and isinstance(stmt.dst, VirtualVariable)
+            or isinstance(stmt, Call)
+            and isinstance(stmt.ret_expr, VirtualVariable)
+        )