Parallel gadget scan with --jobs (#23, part 2)

README.md

@@ -35,7 +35,7 @@ Now, you can install dependencies in [requirements.txt](requirements.txt):
 ```
 usage: rop3.py [-h] [-v] [--depth <bytes>] [--all] [--rop | --no-rop] [--retf | --no-retf] [--jop | --no-jop] [--allow-undeterministic-gadgets] [--allow-complex-memory-ops] [--verbose]
                [--binary <file> [<file> ...]] [--badchar <hex> [<hex> ...]] [--badchar-bytes <hex> [<hex> ...]] [--keep-canary-address] [--base <hex> [<hex> ...]] [--arch <name>] [--symbols]
-               [--output {text,json,csv}] [--op <op>] [--dst <reg>] [--src <reg>] [--ropchain <file>] [--exhaustive | --no-exhaustive] [--interactive] [--cache] [--cache-dir <dir>]
+               [--output {text,json,csv}] [--op <op>] [--dst <reg>] [--src <reg>] [--ropchain <file>] [--exhaustive | --no-exhaustive] [--interactive] [--jobs <n>] [--cache] [--cache-dir <dir>]
 
 This tool allows you to search for gadgets, operations, and ROP chains using a backtracking algorithm in a tree-like structure
 
@@ -73,10 +73,15 @@ options:
   --exhaustive, --no-exhaustive
                         exhaustive search for ROP chains
   --interactive         scan the binary once and drop into an interactive prompt
+  --jobs <n>            number of worker processes for the gadget scan (default: 1)
   --cache               cache discovered gadgets on disk and reuse them on repeated runs over the same file and options
   --cache-dir <dir>     directory for the gadget cache (default: $XDG_CACHE_HOME/rop3)
 ```
 
+### Parallel scan
+
+`--jobs N` distributes the gadget scan over `N` worker processes. Each executable section is split into chunks scanned independently, then the results are merged and deduplicated, so the output is identical to a serial run. The speedup is sublinear (the merge, deduplication and sort run in the parent, and there is per-process start-up cost), so it is worth it mainly for large binaries and/or a high `--depth`; on small inputs the process overhead dominates and `--jobs 1` (the default) is faster.
+
 ### Gadget cache
 
 With `--cache`, the gadgets discovered for a binary are stored on disk and reused on later runs over the same file and options, skipping the scan. The cache key binds the file content hash and every option that affects the result, so a changed binary or option misses cleanly. This is especially handy for large binaries and for the interactive mode.

rop3/api.py

@@ -39,7 +39,7 @@ def __init__(self, binaries, *, depth=gadfinder.DEPTH, rop=True, jop=False,
                  retf=False, all=False, allow_undeterministic=False,
                  allow_complex_mem=False, avoid_canary=True, base=None,
                  badchars=None, badchar_bytes=None, arch=None, symbols=False,
-                 cache=False, cache_dir=None):
+                 cache=False, cache_dir=None, jobs=1):
         self.binaries = [binaries] if isinstance(binaries, str) else list(binaries)
         self.base = base
         self.badchars = badchars
@@ -63,7 +63,8 @@ def __init__(self, binaries, *, depth=gadfinder.DEPTH, rop=True, jop=False,
         if avoid_canary:
             flags |= gadfinder.AVOID_CANARY
 
-        self._finder = GadFinder(depth, flags, cache=cache, cache_dir=cache_dir)
+        self._finder = GadFinder(depth, flags, cache=cache, cache_dir=cache_dir,
+                                 jobs=jobs)
         self._gadgets = None
 
     @classmethod
@@ -78,7 +79,8 @@ def from_args(cls, args):
         self.arch = args.arch
         self.symbols = args.symbols
         self._finder = GadFinder(args.depth, args.flags,
-                                 cache=args.cache, cache_dir=args.cache_dir)
+                                 cache=args.cache, cache_dir=args.cache_dir,
+                                 jobs=args.jobs)
         self._gadgets = None
         return self
 

rop3/args.py

@@ -51,6 +51,7 @@ def __init__(self):
         self.argparser.add_argument('--ropchain', type=str, metavar='<file>', help='plain text file with a ROP chain')
         self.argparser.add_argument('--exhaustive', action=argparse.BooleanOptionalAction, help="exhaustive search for ROP chains", default=False)
         self.argparser.add_argument('--interactive', action='store_true', default=False, help='scan the binary once and drop into an interactive prompt')
+        self.argparser.add_argument('--jobs', type=int, metavar='<n>', default=1, help='number of worker processes for the gadget scan (default: 1)')
         self.argparser.add_argument('--cache', action='store_true', default=False, help='cache discovered gadgets on disk and reuse them on repeated runs over the same file and options')
         self.argparser.add_argument('--cache-dir', type=str, metavar='<dir>', default=None, help='directory for the gadget cache (default: $XDG_CACHE_HOME/rop3)')
 
@@ -123,6 +124,9 @@ def _check_args(self, args):
                 if value < 0x00 or value > 0xff:
                     debug.error(f'{badchar}: bad char must be one byte (range 0x00-0xff)')
 
+        if args.jobs is not None and args.jobs < 1:
+            debug.error(f'--jobs must be >= 1 (got {args.jobs})')
+
         if args.ropchain:
             ropchain_filename = os.path.abspath(args.ropchain)
             if not os.path.isfile(ropchain_filename):

rop3/gadfinder.py

@@ -17,15 +17,18 @@
 
 import os
 import re
+import math
 import bisect
 import capstone
+import multiprocessing
 
 from rop3.cache import GadgetCache
 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.archs.x86_arch import X86_Architecture, X64_Architecture
 from rop3.ropchain import RopChain
 import rop3.parser as parser
 
@@ -53,10 +56,12 @@ class GadFinder:
     '''
     Class to search gadgets in a binary
     '''
-    def __init__(self, depth=DEPTH, flags=DEFAULT, cache=False, cache_dir=None):
+    def __init__(self, depth=DEPTH, flags=DEFAULT, cache=False, cache_dir=None,
+                 jobs=1):
         self.depth = depth
         self.flags = flags
         self._cache = GadgetCache(cache_dir) if cache else None
+        self._jobs = max(1, int(jobs)) if jobs else 1
 
     def find(self, filenames: list[str], base=None, badchars=None,
              badchar_bytes=None, arch=None, symbols=False) -> list[Gadget]:
@@ -80,11 +85,15 @@ def find(self, filenames: list[str], base=None, badchars=None,
                     else:
                         existing.count += 1
                         ''' Among duplicates, keep the address with the fewest
-                            terminator canary bytes (see issue #5) '''
-                        if self._avoid_canary() and \
-                                self._addr_canary_score(gadget, avoid) < self._addr_canary_score(existing, avoid):
-                            gadget.count = existing.count
-                            seen[gadget] = gadget
+                            terminator canary bytes; break ties by the lower
+                            address so the result is deterministic regardless
+                            of scan order (serial or parallel). See issue #5. '''
+                        if self._avoid_canary():
+                            new_key = (self._addr_canary_score(gadget, avoid), gadget.vaddr)
+                            cur_key = (self._addr_canary_score(existing, avoid), existing.vaddr)
+                            if new_key < cur_key:
+                                gadget.count = existing.count
+                                seen[gadget] = gadget
                 unique = len(seen) - before
                 debug.info(f'{unique} unique gadgets ({total - unique} duplicates discarded)')
             return self._sort_gadgets(list(seen.values()))
@@ -163,6 +172,7 @@ def _search_gadgets(self, binary, badchars, badchar_bytes=None, symbol_table=Non
         stores the raw (vaddr, bytes) records; everything address/disassembly
         derived (decodes, symbol) is rebuilt here.
         '''
+        key = None
         if self._cache is not None:
             key = self._cache.key(
                 self._cache.file_hash(binary.raw_data),
@@ -174,6 +184,13 @@ def _search_gadgets(self, binary, badchars, badchar_bytes=None, symbol_table=Non
                 yield from self._reconstruct(binary, cached, symbol_table)
                 return
 
+        if self._jobs > 1:
+            records = self._scan_parallel(binary, badchars, badchar_bytes)
+            if self._cache is not None:
+                self._cache.store(key, records)
+            yield from self._reconstruct(binary, records, symbol_table)
+            return
+
         records = [] if self._cache is not None else None
         arch = arch_singleton.arch.arch
         mode = arch_singleton.arch.mode
@@ -187,6 +204,43 @@ def _search_gadgets(self, binary, badchars, badchar_bytes=None, symbol_table=Non
         if records is not None:
             self._cache.store(key, records)
 
+    def _scan_parallel(self, binary, badchars, badchar_bytes):
+        '''
+        Scan the executable sections across worker processes. Each section is
+        split into chunks; a chunk emits only the gadgets whose termination
+        falls inside its window (the slice extends `depth` bytes earlier so
+        gadgets straddling a boundary are still complete), so there are no
+        cross-chunk duplicates. Returns sorted [vaddr, hex] records.
+        '''
+        arch = arch_singleton.arch.arch
+        mode = arch_singleton.arch.mode
+        terminations = self._gad_terminations()
+
+        tasks = []
+        for section in binary.get_exec_sections():
+            opcodes = section['opcodes']
+            sec_vaddr = section['vaddr']
+            n = len(opcodes)
+            chunk = max(4096, math.ceil(n / (self._jobs * 4)))
+            for lo in range(0, n, chunk):
+                hi = min(lo + chunk, n)
+                start = max(0, lo - self.depth)
+                ''' Termination END offsets run in [0, n]; the final chunk owns
+                    the closing n as well, so make its window inclusive. '''
+                emit_hi = hi + 1 if hi == n else hi
+                tasks.append((
+                    arch, mode, self.depth, int(self.flags), terminations,
+                    badchars, badchar_bytes,
+                    opcodes[start:hi], start, sec_vaddr, lo, emit_hi,
+                ))
+
+        records = []
+        with multiprocessing.Pool(self._jobs) as pool:
+            for part in pool.imap_unordered(_scan_worker, tasks):
+                records.extend(part)
+        records.sort()   # deterministic order regardless of worker scheduling
+        return records
+
     def _scan(self, binary, badchars, badchar_bytes):
         ''' Single pass over the executable sections; yields the raw
             (vaddr, bytes, decodes) of every valid gadget (one disassembly). '''
@@ -317,3 +371,48 @@ def _is_valid_bytes(self, gadget_bytes, badchar_bytes):
         forbidden = {int(b, 0) for b in badchar_bytes}
         return not any(byte in forbidden for byte in gadget_bytes)
 
+
+def _arch_for(arch_const, mode):
+    ''' Rebuild the architecture object inside a worker process. '''
+    return X64_Architecture() if mode == capstone.CS_MODE_64 else X86_Architecture()
+
+
+def _scan_worker(task):
+    '''
+    Worker (runs in its own process): scan one section chunk and return the
+    raw [vaddr, hex] records for the gadgets whose termination lies in the
+    chunk's window. Decodes are not returned (capstone objects are not
+    picklable); the parent rebuilds them.
+    '''
+    (arch_const, mode, depth, flags, terminations, badchars, badchar_bytes,
+     slice_bytes, slice_start, sec_vaddr, emit_lo, emit_hi) = task
+
+    arch_singleton.reset()
+    arch_singleton.initialize(_arch_for(arch_const, mode))
+    finder = GadFinder(depth, flags)
+
+    md = capstone.Cs(arch_const, mode)
+    md.detail = True
+
+    out = []
+    for termination in terminations:
+        for match in re.finditer(termination['bytes'], slice_bytes):
+            ref_local = match.end()
+            ref_off = slice_start + ref_local        # offset within the section
+            ''' Only this chunk owns terminations in [emit_lo, emit_hi) '''
+            if not (emit_lo <= ref_off < emit_hi):
+                continue
+            for d in range(termination['size'], depth + 1):
+                start_local = ref_local - d
+                if start_local < 0:
+                    continue
+                vaddr = sec_vaddr + ref_off - d
+                if finder._is_valid_address(vaddr, badchars, mode):
+                    raw = slice_bytes[start_local:ref_local]
+                    if not finder._is_valid_bytes(raw, badchar_bytes):
+                        continue
+                    decodes = list(md.disasm(raw, vaddr))
+                    if finder._is_valid_gadget(decodes):
+                        out.append([vaddr, raw.hex()])
+    return out
+

rop3/utils.py

@@ -28,7 +28,9 @@
 PATCH = 0
 VERSION = f'{MAJOR}.{MINOR}.{PATCH}'
 
-TOOL_NAME = os.path.basename(os.path.realpath(__main__.__file__))
+# __main__ may lack __file__ when imported as a library or in a spawned
+# multiprocessing worker; fall back to a sensible default in that case.
+TOOL_NAME = os.path.basename(os.path.realpath(getattr(__main__, '__file__', 'rop3.py')))
 
 HEADER = '\
                           .d8888b.  \n\

tests/test_parallel.py

@@ -0,0 +1,60 @@
+'''
+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
+
+from rop3 import Rop3
+
+from conftest import build_minimal_elf, EM_X86_64, ET_DYN
+
+
+def _key(gadgets):
+    return sorted((g.vaddr, g.text_repr, g.count) for g in gadgets)
+
+
+@pytest.fixture
+def big_elf(tmp_path):
+    # A larger .text with many ret (0xc3) bytes so chunks actually split, and
+    # a ret as the very last byte to exercise the closing-window boundary.
+    text = (b'\x58\x5b\x59\x5a\xc3' * 4000) + b'\xc3'
+    data = build_minimal_elf(64, EM_X86_64, text, 0x1000, ET_DYN)
+    path = tmp_path / 'big.elf'
+    path.write_bytes(data)
+    return str(path)
+
+
+@pytest.mark.parametrize('jobs', [2, 4])
+def test_parallel_matches_serial(big_elf, jobs):
+    serial = Rop3(big_elf).gadgets()
+    parallel = Rop3(big_elf, jobs=jobs).gadgets()
+    assert _key(serial) == _key(parallel)
+
+
+def test_parallel_matches_serial_with_options(big_elf):
+    ''' Same equivalence with bad-char and depth options. '''
+    kwargs = dict(depth=8, badchar_bytes=['0x59'])
+    serial = Rop3(big_elf, **kwargs).gadgets()
+    parallel = Rop3(big_elf, jobs=4, **kwargs).gadgets()
+    assert _key(serial) == _key(parallel)
+
+
+def test_parallel_with_cache(big_elf, tmp_path):
+    ''' Parallel scan results are cacheable and reload identically. '''
+    cache_dir = str(tmp_path / 'cache')
+    cold = Rop3(big_elf, jobs=4, cache=True, cache_dir=cache_dir).gadgets()
+    warm = Rop3(big_elf, cache=True, cache_dir=cache_dir).gadgets()
+    assert _key(cold) == _key(warm)