fix function calling in i386

angrop/chain_builder/__init__.py

@@ -135,13 +135,13 @@ def execve(self, path=None, path_addr=None):
             return None
         return self._sys_caller.execve(path=path, path_addr=path_addr)
 
-    def shift(self, length, preserve_regs=None):
+    def shift(self, length, preserve_regs=None, next_pc_idx=-1):
         """
         build a rop chain to shift the stack to a specific value
         :param length: the length of sp you want to shift
         :param preserve_regs: set of registers to preserve, e.g. ('eax', 'ebx')
         """
-        return self._shifter.shift(length, preserve_regs=preserve_regs)
+        return self._shifter.shift(length, preserve_regs=preserve_regs, next_pc_idx=next_pc_idx)
 
     def retsled(self, size, preserve_regs=None):
         """

angrop/chain_builder/func_caller.py

@@ -66,7 +66,7 @@ def _func_call(self, func_gadget, cc, args, extra_regs=None, preserve_regs=None,
         # 1. handle stack arguments
         # 2. handle function return address to maintain the control flow
         if stack_arguments:
-            cleaner = self.chain_builder.shift((len(stack_arguments)+1)*arch_bytes) # +1 for itself
+            cleaner = self.chain_builder.shift((len(stack_arguments)+1)*arch_bytes, next_pc_idx=-1)
             chain.add_gadget(cleaner._gadgets[0])
             for arg in stack_arguments:
                 chain.add_value(arg)

angrop/chain_builder/shifter.py

@@ -1,6 +1,7 @@
 import logging
 from collections import defaultdict
 
+from .. import rop_utils
 from .builder import Builder
 from ..rop_chain import RopChain
 from ..errors import RopException
@@ -64,7 +65,12 @@ def same_effect(g1, g2):
             return False
         return True
 
-    def shift(self, length, preserve_regs=None):
+    def shift(self, length, preserve_regs=None, next_pc_idx=-1):
+        """
+        length:         how many bytes to shift
+        preserve_regs:  what registers not to clobber
+        next_pc_idx:    where is the next pc, e.g for ret, it is -1
+        """
         preserve_regs = set(preserve_regs) if preserve_regs else set()
         arch_bytes = self.project.arch.bytes
 
@@ -75,14 +81,31 @@ def shift(self, length, preserve_regs=None):
             raise RopException("Encounter a shifting request that requires chaining multiple shifting gadgets " +
                                "together which is not support atm. Plz create an issue on GitHub " +
                                "so we can add the support!")
+        g_cnt = length // arch_bytes
+        next_pc_idx = (next_pc_idx % g_cnt + g_cnt) % g_cnt # support negative indexing
         for g in self.shift_gadgets[length]:
             if preserve_regs.intersection(g.changed_regs):
                 continue
+            if next_pc_idx == g_cnt-1:
+                if g.transit_type != 'ret':
+                    continue
+            else:
+                if g.transit_type != 'pop_pc':
+                    continue
+                if g.pc_offset != next_pc_idx*arch_bytes:
+                    continue
             try:
                 chain = RopChain(self.project, self.chain_builder)
                 chain.add_gadget(g)
-                for _ in range(g.stack_change//arch_bytes-1):
-                    chain.add_value(self._get_fill_val())
+                for idx in range(g_cnt):
+                    if idx != next_pc_idx:
+                        chain.add_value(self._get_fill_val())
+                    else:
+                        next_pc_val = rop_utils.cast_rop_value(
+                            chain._blank_state.solver.BVS("next_pc", self.project.arch.bits),
+                            self.project,
+                        )
+                        chain.add_value(next_pc_val)
                 if self.verify_shift(chain, length, preserve_regs):
                     return chain
             except RopException:

angrop/rop_chain.py

@@ -11,12 +11,12 @@ class RopChain:
     """
     cls_timeout = CHAIN_TIMEOUT_DEFAULT
 
-    def __init__(self, project, rop, state=None, badbytes=None):
+    def __init__(self, project, builder, state=None, badbytes=None):
         """
         """
         self._p = project
         self._pie = self._p.loader.main_object.pic
-        self._rop = rop
+        self._builder = builder
 
         self._gadgets = []
         self._values = []
@@ -151,15 +151,17 @@ def _concretize_chain_values(self, constraints=None, timeout=None, preserve_next
         """
         concretize chain values with a timeout
         """
-        if self.next_pc_idx() is not None:
+        if self.next_pc_idx() is not None and not preserve_next_pc:
             try:
-                return (
-                    self + self._rop.chain_builder.shift(self._p.arch.bytes)
-                )._concretize_chain_values(
-                    constraints=constraints,
-                    timeout=timeout,
-                    preserve_next_pc=preserve_next_pc,
-                )
+                # the following line is the final touch for chains ending with retn-style
+                # gadget to make sure that the next_pc is at the end of the chain
+                chain = self + self._builder.chain_builder.shift(self._p.arch.bytes)
+                values = chain._concretize_chain_values(
+                                    constraints=constraints,
+                                    timeout=timeout,
+                                    preserve_next_pc=preserve_next_pc,
+                                )
+                return values
             except RopException:
                 pass
         if timeout is None:
@@ -288,7 +290,7 @@ def exec(self, max_steps=None, timeout=None):
                                                          allow_simprocedures=True)
 
     def copy(self):
-        cp = RopChain(self._p, self._rop)
+        cp = RopChain(self._p, self._builder)
         cp._gadgets = list(self._gadgets)
         cp._values = list(self._values)
         cp.payload_len = self.payload_len

tests/test_chainbuilder.py

@@ -4,6 +4,7 @@
 import angrop # pylint: disable=unused-import
 import claripy
 from angrop.rop_value import RopValue
+from angrop.errors import RopException
 
 BIN_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "binaries")
 CACHE_DIR = os.path.join(BIN_DIR, 'tests_data', 'angrop_gadgets_cache')
@@ -394,6 +395,30 @@ def test_pop_pc_syscall_chain():
     assert state.regs.rdi.concrete_value == 0x41414141
     assert 0 not in state.posix.fd
 
+def test_retn_i386_call_chain():
+    cache_path = os.path.join(CACHE_DIR, "bronze_ropchain")
+    proj = angr.Project(os.path.join(BIN_DIR, "tests", "i386", "bronze_ropchain"), auto_load_libs=False)
+    rop = proj.analyses.ROP()
+
+    if os.path.exists(cache_path):
+        rop.load_gadgets(cache_path)
+    else:
+        rop.find_gadgets()
+        rop.save_gadgets(cache_path)
+
+    # force to use 'retn 0xc' to clean up function arguments
+    g = rop.analyze_gadget(0x809d9fb)
+    rop._chain_builder._shifter.shift_gadgets = {g.stack_change: [g]}
+
+    rop.func_call('write', [1, 0x80AC5E8, 17], needs_return=False)
+
+    chain = None
+    try:
+        chain = rop.func_call('write', [1, 0x80AC5E8, 17])
+    except RopException:
+        pass
+    assert chain is None
+
 def test_aarch64_basic_reg_setting():
     proj = angr.load_shellcode(
         """