From e8c57245ea1263db025e266fbc36b21d680eaeb5 Mon Sep 17 00:00:00 2001
From: Kyle Zeng <zengyhkyle@gmail.com>
Date: Tue, 21 Jan 2025 11:41:30 -0700
Subject: [PATCH] make _add_gadget_value a separate function

---
 angrop/chain_builder/builder.py |  2 +-
 angrop/rop_chain.py             | 30 ++++++++++++++----------------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/angrop/chain_builder/builder.py b/angrop/chain_builder/builder.py
index d361326..66d9bda 100644
--- a/angrop/chain_builder/builder.py
+++ b/angrop/chain_builder/builder.py
@@ -222,7 +222,7 @@ def map_stack_var(ast, value):
             if sym_var in stack_var_to_value:
                 val = stack_var_to_value[sym_var]
                 if isinstance(val, RopGadget):
-                    chain.add_gadget(val, append_addr_only=True)
+                    chain._add_gadget_value(val)
                 else:
                     # HACK: Because angrop appears to have originally been written
                     # with assumptions around x86 ret gadgets, the target of the final jump
diff --git a/angrop/rop_chain.py b/angrop/rop_chain.py
index 1a7557b..f50c4c4 100644
--- a/angrop/rop_chain.py
+++ b/angrop/rop_chain.py
@@ -68,21 +68,10 @@ def add_value(self, value):
         self._values.append(value)
         self.payload_len += self._p.arch.bytes
 
-    def add_gadget(self, gadget, append_addr_only=False):
-        # angrop was originally written with the assumption that gadget addresses
-        # appear in the chain in the same order in which the gadgets are executed.
-        # This is not always true when there are gadgets that end with a jump to
-        # an address from a register instead of the stack.
-        # For example, if the ROP chain has three gadgets A, B, and C where gadget
-        # B ends with a jump to some register, gadget A would have to load the
-        # address of gadget C into the register before jumping to gadget B.
-        # Therefore, the address of gadget C might need to be placed before the
-        # address of gadget B.
-        # The append_addr_only argument and the set_gadgets method below were added
-        # to support chains like this without breaking the existing API.
-        if not append_addr_only:
-            self._gadgets.append(gadget)
-
+    def _add_gadget_value(self, gadget):
+        """
+        create a RopValue for the gadget's address and add it to chain._values
+        """
         value = gadget.addr
         if self._pie:
             value -= self._p.loader.main_object.mapped_base
@@ -90,11 +79,20 @@ def add_gadget(self, gadget, append_addr_only=False):
         if self._pie:
             value._rebase = True
 
-        if append_addr_only or (idx := self.next_pc_idx()) is None:
+        if (idx := self.next_pc_idx()) is None:
             self.add_value(value)
         else:
             self._values[idx] = value
 
+    def add_gadget(self, gadget):
+        # angrop was originally written with the assumption that gadget addresses
+        # appear in the chain in the same order in which the gadgets are executed.
+        # This is not always true when there are gadgets that end with a jump to
+        # an address from a register instead of the stack.
+        # For example, when we do `ret` in aarch64
+        self._add_gadget_value(gadget)
+        self._gadgets.append(gadget)
+
     def set_gadgets(self, gadgets: list[RopGadget]):
         self._gadgets = gadgets