rebalance add/sub in ropvalue ast so we can track symbolic user ropva…

…lue correctly

angrop/chain_builder/builder.py

@@ -103,6 +103,45 @@ def _get_ptr_to_null(self):
                     return addr
         return None
 
+    def _ast_contains_stack_data(self, ast):
+        vs = ast.variables
+        return len(vs) == 1 and list(vs)[0].startswith('symbolic_stack_')
+
+    def _rebalance_ast(self, lhs, rhs):
+        """
+        we know that lhs (stack content with modification) == rhs (user ropvalue)
+        since user ropvalue may be symbolic, we need to present the stack content using the user ropvalue and store it
+        on stack so that users can eval on their own ropvalue and get the correct solves
+        TODO: currently, we only support add/sub
+        """
+        assert self._ast_contains_stack_data(lhs)
+        while lhs.depth != 1:
+            match lhs.op:
+                case "__add__" | "__sub__":
+                    arg0 = lhs.args[0]
+                    arg1 = lhs.args[1]
+                    flag = self._ast_contains_stack_data(arg0)
+                    op = lhs.op
+                    if flag:
+                        lhs = arg0
+                        other = arg1
+                    else:
+                        lhs = arg1
+                        other = arg0
+                    if op == "__add__":
+                        rhs -= other
+                    elif flag:
+                        rhs += other
+                    else:
+                        rhs = other - rhs
+                case "Reverse":
+                    lhs = lhs.args[0]
+                    rhs = claripy.Reverse(rhs)
+                case _:
+                    raise ValueError(f"{lhs.op} cannot be rebalanced at the moment. plz create an issue!")
+        assert self._ast_contains_stack_data(lhs)
+        return lhs, rhs
+
     @rop_utils.timeout(8)
     def _build_reg_setting_chain(
         self, gadgets, modifiable_memory_range, register_dict, stack_change
@@ -185,6 +224,14 @@ def map_stack_var(ast, value):
                     raise RopException("Register set to incorrect value")
             else:
                 state.solver.add(var == val.data)
+                lhs, rhs = self._rebalance_ast(var, val.data)
+                rhs = claripy.Reverse(rhs)
+                ropvalue = val.copy()
+                if val.rebase:
+                    ropvalue._value = rhs - ropvalue._code_base
+                else:
+                    ropvalue._value = rhs
+                map_stack_var(lhs, ropvalue)
 
         # Constrain memory access addresses.
         for action in state.history.actions:

angrop/gadget_finder/__init__.py

@@ -125,7 +125,10 @@ def _initialize_gadget_analyzer(self):
 
     def analyze_gadget(self, addr):
         g = self.gadget_analyzer.analyze_gadget(addr)
-        if g:
+        if isinstance(g, list):
+            for x in g:
+                x.project = self.project
+        elif g:
             g.project = self.project
         return g
 
@@ -156,7 +159,7 @@ def get_duplicates(self):
         cache = self._cache
         return {k:v for k,v in cache.items() if len(v) >= 2}
 
-    def find_gadgets(self, processes=16, show_progress=True):
+    def find_gadgets(self, processes=4, show_progress=True):
         self._cache = {}
 
         initargs = (self.gadget_analyzer,)

angrop/rop.py

@@ -96,9 +96,15 @@ def _screen_gadgets(self):
 
     def analyze_gadget(self, addr):
         g = self.gadget_finder.analyze_gadget(addr)
-        if g:
+        if g is None:
+            return g
+
+        if isinstance(g, list):
+            self._all_gadgets += g
+        else:
             self._all_gadgets.append(g)
-            self._screen_gadgets()
+
+        self._screen_gadgets()
         return g
 
     def analyze_gadget_list(self, addr_list, processes=4, show_progress=True):