From 707fe12e3db018240a44c8c2c6b0a7e321401660 Mon Sep 17 00:00:00 2001
From: anthraxx <levente@leventepolyak.net>
Date: Tue, 30 Mar 2021 22:34:14 +0200
Subject: [PATCH] chore(ghidra): use memoize feature to cache r2pipe handle

We need to cache and reset the r2pipe handle based on memoize conditions
as the PIE base could potentially change.
---
 pwndbg/ghidra.py  |  4 ++--
 pwndbg/memoize.py | 24 ++++++++++++++++++++++++
 pwndbg/radare2.py | 27 +++++++++++++++++++++------
 3 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/pwndbg/ghidra.py b/pwndbg/ghidra.py
index 25b252673..05465299a 100644
--- a/pwndbg/ghidra.py
+++ b/pwndbg/ghidra.py
@@ -17,9 +17,8 @@ def decompile(func=None):
 
     Raises Exception if any fatal error occures.
     """
-    filename = gdb.current_progspace().filename
     try:
-        r2 = pwndbg.radare2.r2pipe(filename)
+        r2 = pwndbg.radare2.r2pipe()
         # LD         list supported decompilers (e cmd.pdc=?)
         # Outputs for example:: pdc\npdg
         if not "pdg" in r2.cmd("LD").split("\n"):
@@ -67,6 +66,7 @@ def decompile(func=None):
         try: # try to read the source filename from debug information
             src_filename = gdb.selected_frame().find_sal().symtab.fullname()
         except: # if non, take the original filename and maybe append .c (just assuming is was c)
+            filename = gdb.current_progspace().filename
             src_filename = filename+".c" if os.path.basename(filename).find(".") < 0 else filename
         source = H.syntax_highlight(source, src_filename)
 
diff --git a/pwndbg/memoize.py b/pwndbg/memoize.py
index 130855f5d..77cbea6ae 100644
--- a/pwndbg/memoize.py
+++ b/pwndbg/memoize.py
@@ -10,6 +10,8 @@ import collections
 import functools
 import sys
 
+import gdb
+
 import pwndbg.events
 
 debug = False
@@ -178,6 +180,27 @@ class while_running(memoize):
     _reset = __reset_while_running
 
 
+class reset_on_new_base_address(memoize):
+    caches   = []
+    kind     = 'new_base_address'
+    filename = None
+    base     = None
+
+    @staticmethod
+    @pwndbg.events.start
+    @pwndbg.events.new_objfile
+    def __reset_on_base():
+        filename = gdb.current_progspace().filename
+        base = pwndbg.elf.exe().address if pwndbg.elf.exe() else None
+        if reset_on_new_base_address.base != base or reset_on_new_base_address.filename != filename:
+            reset_on_new_base_address.filename = filename
+            reset_on_new_base_address.base = base
+            for obj in reset_on_new_base_address.caches:
+                obj.clear()
+
+    _reset = __reset_on_base
+
+
 def reset():
     forever._reset()
     reset_on_stop._reset()
@@ -185,4 +208,5 @@ def reset():
     reset_on_objfile._reset()
     reset_on_start._reset()
     reset_on_cont._reset()
+    reset_on_new_base_address._reset()
     while_running._reset()
diff --git a/pwndbg/radare2.py b/pwndbg/radare2.py
index ae51f668b..7ecd1c7f0 100644
--- a/pwndbg/radare2.py
+++ b/pwndbg/radare2.py
@@ -1,17 +1,32 @@
+import gdb
+
 import pwndbg.elf
+import pwndbg.memoize
+
+
+@pwndbg.memoize.reset_on_new_base_address
+def r2pipe():
+    """
+    Spawn and return a r2pipe handle for the current process file.
+
+    This function requires a radare2 installation plus the r2pipe python
+    library. The base address is automatically set for PIE when loading the
+    binary.
+    After opening the handle, the binary is automatically analyzed.
 
-radare2 = {}
+    Raises ImportError if r2pipe python library is not available.
+    Raises Exception if anything goes fatally wrong.
 
+    Returns a r2pipe.open handle.
+    """
+    filename = gdb.current_progspace().filename
+    if not filename:
+        raise Exception('Could not find objfile to create a r2pipe for')
 
-def r2pipe(filename):
-    r2 = radare2.get(filename)
-    if r2:
-        return r2
     import r2pipe
     flags = ['-e', 'io.cache=true']
     if pwndbg.elf.get_elf_info(filename).is_pie and pwndbg.elf.exe():
         flags.extend(['-B', hex(pwndbg.elf.exe().address)])
     r2 = r2pipe.open(filename, flags=flags)
-    radare2[filename] = r2
     r2.cmd("aaaa")
     return r2