fix/badbyte

angrop/chain_builder/builder.py

@@ -256,7 +256,8 @@ def _rebalance_ast(self, lhs, rhs, mode='stack'):
                 if ast.op == 'BVS' and ast.args[0].startswith('symbolic_stack'):
                     target_ast = ast
                     break
-            assert target_ast is not None
+            if target_ast is None:
+                raise RopException("rebalancing non-stack value")
 
             solver = claripy.Solver()
             solver.add(lhs == rhs)

angrop/rop_chain.py

@@ -219,8 +219,10 @@ def copy(self):
         cp._values = list(self._values)
         cp.payload_len = self.payload_len
         cp._blank_state = self._blank_state.copy()
-        cp.badbytes = self.badbytes.copy()
+        cp.badbytes = self.badbytes
 
+        cp._pivoted = self._pivoted
+        cp._init_sp = self._init_sp
         return cp
 
     #### Solver Layer ####

angrop/rop_effect.py

@@ -230,14 +230,18 @@ def import_effect(self, gadget):
     def copy_effect(self, cp):
         cp.stack_change = self.stack_change
         cp.changed_regs = set(self.changed_regs)
-        cp.reg_pops = set(self.reg_pops)
         cp.concrete_regs = dict(self.concrete_regs)
         cp.reg_dependencies = dict(self.reg_dependencies)
         cp.reg_controllers = dict(self.reg_controllers)
+        cp.reg_pops = set(self.reg_pops)
         cp.reg_moves = list(self.reg_moves)
+
         cp.mem_reads = list(self.mem_reads)
         cp.mem_writes = list(self.mem_writes)
         cp.mem_changes = list(self.mem_changes)
         cp.bbl_addrs = list(self.bbl_addrs)
         cp.isn_count = self.isn_count
+
+        cp.pop_equal_set = set(self.pop_equal_set)
+        cp.branch_dependencies = self.has_conditional_branch
         return cp

angrop/rop_gadget.py

@@ -104,8 +104,6 @@ def copy(self):
         out.pc_offset = self.pc_offset
         out.pc_reg = self.pc_reg
         out.pc_target = self.pc_target
-        out.branch_dependencies = set(self.branch_dependencies)
-        out.has_conditional_branch = self.has_conditional_branch
         return out
 
     def __getstate__(self):

tests/test_chainbuilder.py

@@ -1248,6 +1248,25 @@ def test_push_pop_move():
     chain = rop.move_regs(rdi='rax')
     assert chain
 
+def test_rebalance_ast_failsafe():
+    proj = angr.load_shellcode(
+            """
+            pop ebx; pop esi; pop edi; xor eax, eax; ret
+            lea eax, [edx + 0xf]; pop edi; ret
+            """,
+            "i386",
+            load_address=0,
+            auto_load_libs=False
+            )
+    rop = proj.analyses.ROP()
+    rop.find_gadgets_single_threaded()
+
+    rop.set_badbytes([0x00])
+    try:
+        rop.set_regs(eax=0, edi=0xdeadbeef)
+    except RopException:
+        pass
+
 def run_all():
     functions = globals()
     all_functions = {x:y for x, y in functions.items() if x.startswith('test_')}