Quality cleanup: no shared-gadget mutation, faster chain build, style

requirements.txt

@@ -1,5 +1,5 @@
 capstone >= 4.0.1
 pefile >= 2019.4.18
-pyelftools
-pyyaml
-macholib
+pyelftools >= 0.27
+pyyaml >= 5.4
+macholib >= 1.14

rop3/arch.py

@@ -103,6 +103,10 @@ def initialize(self, arch: Architecture):
             return
         self._arch = arch
 
+    def reset(self) -> None:
+        ''' Clear the current architecture (mainly for tests / library use) '''
+        self._arch = None
+
     def is_initialized(self) -> bool:
         return self._arch is not None
 

rop3/gadfinder.py

@@ -15,6 +15,20 @@
 along with rop3. If not, see <https://www.gnu.org/licenses/>.
 '''
 
+import os
+import re
+import capstone
+
+import rop3.utils as utils
+import rop3.debug as debug
+import rop3.binary
+import rop3.operation as operation
+from rop3.arch import arch_singleton
+from rop3.ropchain import RopChain
+import rop3.parser as parser
+
+from .gadget import Gadget
+
 ''' Default depth engine '''
 DEPTH = 5
 
@@ -33,20 +47,6 @@
     terminators for gets() and alike) and 0xff (EOF). See issue #5. '''
 CANARY_BYTES = (0x00, 0x0a, 0x0d, 0xff)
 
-import os
-import re
-import capstone
-
-import rop3.utils as utils
-import rop3.debug as debug
-import rop3.binary
-import rop3.operation as operation
-from rop3.arch import arch_singleton
-from rop3.ropchain import RopChain
-import rop3.parser as parser
-
-from .gadget import Gadget
-
 class GadFinder:
     '''
     Class to search gadgets in a binary

rop3/operation.py

@@ -16,6 +16,7 @@
 '''
 
 import capstone
+import dataclasses
 
 from rop3.arch import arch_singleton
 
@@ -44,11 +45,17 @@ def filter_gadgets(self, gadgets) -> list[Gadget]:
         for gadget in gadgets:
             (equal, set_, dst, src) = self.template.is_equal(gadget.decodes)
             if equal:
-                gadget.op = self.template.name
-                gadget.dst = self.dst if self.dst else dst
-                gadget.src = self.src if self.src else src
-                gadget.calculate_side_effects()
-                ret.append(gadget)
+                ''' Annotate a copy so the shared input gadget is not mutated
+                    (the same object may be filtered for several operations).
+                    replace() also gives the copy fresh side-effect sets. '''
+                matched = dataclasses.replace(
+                    gadget,
+                    op=self.template.name,
+                    dst=self.dst if self.dst else dst,
+                    src=self.src if self.src else src,
+                )
+                matched.calculate_side_effects()
+                ret.append(matched)
 
         return ret
 
@@ -120,7 +127,7 @@ def is_equal(self, decodes):
             return (False, dst, src)
 
         for i, item in enumerate(self.items):
-            if type(item) != Instruction:
+            if not isinstance(item, Instruction):
                 return (False, dst, src)
 
             (equal, ins_dst, ins_src) = item.is_equal(decodes[i])

rop3/parsers/yaml_parser.py

@@ -100,7 +100,7 @@ def _parse_op(self, op, content):
                     for operand in ('op1', 'op2'):
                         if operand in item:
                             current_op = item[operand]
-                            if type(current_op) == dict:
+                            if isinstance(current_op, dict):
                                 raise NotImplementedError
                             else:
                                 # Resolve aliases if necessary

rop3/ropchain.py

@@ -146,28 +146,34 @@ def _build_per_comb_gadgets(
           - sorted by heuristic_basic_count (fewest side effects first)
           - pruned of subsumed gadgets (when prune_equivalent), exploiting sort order
         Returns an array indexed [comb_idx][step_idx].
+
+        Each operation's gadget list is sorted once up front, and the
+        filter+prune result is memoised per (step, req_dst, req_src): different
+        combinations frequently request the same concrete registers for a given
+        step, so this avoids recomputing the same filtered list repeatedly.
         """
-        result = []
-        working_gadgets = [list(gl) for gl in ops_gadgets]
+        sorted_gadgets = [sorted(gl, key=heuristic_basic_count) for gl in ops_gadgets]
+        cache: dict = {}
+
+        def build_step(i, req_dst, req_src):
+            key = (i,
+                   None if req_dst is None else str(req_dst),
+                   None if req_src is None else str(req_src))
+            if key not in cache:
+                filtered = [ gad for gad in sorted_gadgets[i] \
+                        if (req_dst is None or str(gad.dst) == str(req_dst)) \
+                        and (req_src is None or str(gad.src) == str(req_src)) ]
+                cache[key] = self._prune(filtered) if prune_equivalent else filtered
+            return cache[key]
 
+        result = []
         for comb in combinations:
             per_step = []
-            for i, gadget_list in enumerate(working_gadgets):
+            for i in range(len(sorted_gadgets)):
                 op = ropchain[i]
                 req_dst = comb.get(op.get('dst'))
                 req_src = comb.get(op.get('src'))
-
-                filtered = [ gad for gad in gadget_list \
-                        if (req_dst is None or str(gad.dst) == str(req_dst)) \
-                        and (req_src is None or str(gad.src) == str(req_src)) ]
-
-                filtered.sort(key=heuristic_basic_count)
-
-                if prune_equivalent:
-                    filtered = self._prune(filtered)
-
-                per_step.append(filtered)
-
+                per_step.append(build_step(i, req_dst, req_src))
             result.append(per_step)
 
         return result

tests/conftest.py

@@ -31,21 +31,21 @@ def reset_arch():
     The architecture is a process-global singleton; reset it around every
     test so cases are independent and can pick their own architecture.
     '''
-    arch_singleton._arch = None
+    arch_singleton.reset()
     yield
-    arch_singleton._arch = None
+    arch_singleton.reset()
 
 
 @pytest.fixture
 def x64():
-    arch_singleton._arch = None
+    arch_singleton.reset()
     arch_singleton.initialize(X64_Architecture())
     return arch_singleton.arch
 
 
 @pytest.fixture
 def x86():
-    arch_singleton._arch = None
+    arch_singleton.reset()
     arch_singleton.initialize(X86_Architecture())
     return arch_singleton.arch
 

tests/test_operation.py

@@ -48,6 +48,17 @@ def test_filter_gadgets_empty_input(x64):
     assert operation.Operation('lc').filter_gadgets([]) == []
 
 
+def test_filter_gadgets_does_not_mutate_input(x64):
+    ''' filter_gadgets must annotate copies, not the shared input gadgets. '''
+    g = make_gadget(b'\x58\xc3', 0x1000)             # pop rax ; ret
+    assert g.op is None and g.dst is None
+    matched = operation.Operation('lc', dst='rax').filter_gadgets([g])
+    assert matched and matched[0] is not g           # a copy was returned
+    assert matched[0].op == 'lc' and matched[0].dst == 'rax'
+    # original is untouched
+    assert g.op is None and g.dst is None and g.side_regs == set()
+
+
 def test_operand_parse_imm_supports_hex_and_negative(x64):
     ''' Regression: immediates parsed with int(x, 0). '''
     op = operation.Operand('rax')

tests/test_ropchain.py

@@ -0,0 +1,67 @@
+'''
+This file is part of rop3 (https://github.com/reverseame/rop3).
+
+rop3 is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+rop3 is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with rop3. If not, see <https://www.gnu.org/licenses/>.
+'''
+
+import pytest
+
+import rop3.ropchain as ropchain_mod
+from rop3.ropchain import RopChain
+
+from conftest import make_gadget
+
+
+def _op(op, dst=None, src=None):
+    return {'data': f'{op}({dst or ""},{src or ""})', 'op': op, 'dst': dst, 'src': src}
+
+
+def test_search_simple_concrete_chain(x64):
+    gadgets = [
+        make_gadget(b'\x58\xc3', 0x1000),   # pop rax ; ret
+        make_gadget(b'\x5b\xc3', 0x1010),   # pop rbx ; ret
+    ]
+    results = list(RopChain(None).search(gadgets, [_op('lc', dst='rax')]))
+    assert results
+    assert len(results[0]) == 1
+    assert results[0][0].text_repr == 'pop rax ; ret'
+
+
+def test_search_two_step_chain(x64):
+    gadgets = [
+        make_gadget(b'\x58\xc3', 0x1000),   # pop rax ; ret
+        make_gadget(b'\x5b\xc3', 0x1010),   # pop rbx ; ret
+    ]
+    chain = [_op('lc', dst='rax'), _op('lc', dst='rbx')]
+    results = list(RopChain(None).search(gadgets, chain))
+    assert results
+    texts = [g.text_repr for g in results[0]]
+    assert texts == ['pop rax ; ret', 'pop rbx ; ret']
+
+
+def test_search_raises_when_no_gadget(x64):
+    gadgets = [make_gadget(b'\x90\xc3', 0x1000)]   # nop ; ret (no lc)
+    with pytest.raises(ropchain_mod.RopChainNotFound):
+        list(RopChain(None).search(gadgets, [_op('lc', dst='rax')]))
+
+
+def test_search_generic_registers(x64):
+    gadgets = [
+        make_gadget(b'\x48\x89\xd8\xc3', 0x1000),   # mov rax, rbx ; ret
+        make_gadget(b'\x48\x89\xd1\xc3', 0x1010),   # mov rcx, rdx ; ret
+    ]
+    chain = [_op('mov', dst='REG1', src='REG2')]
+    results = list(RopChain(None).search(gadgets, chain))
+    assert results
+    assert all(len(r) == 1 for r in results)