Repository: pwndbg/pwndbg
Pull request: Fix caching for pwndbg.gdblib.elf functions
Fixes #2237

Changes two things:
1. Functions in `pwndbg.gdblib.elf` that used to cache on only `start` now cache on both `start` and `objfile`, since changing object files should reload ELF information.
2. Implemented priority event handlers that run before other handlers, so that cache clears will happen before other handlers on that event.


Problem statement:
Incorrect caching on `start` command
### Description

`start` caches the entry breakpoint for too long (I think this is an issue with the caching on `pwndbg.gdblib.elf.entry()` specifically), so if you enable ASLR then `start` twice, the second time won't reset the breakpoint.

### Steps to reproduce

Pick a PIE executable, then run `set disable-randomization off`, then `start`, then `start` again. If you run `memoize` to disable caching, the issue no longer persists.

### My setup

```
Gdb:      14.2
Python:   3.12.4 (main, Jun  7 2024, 06:33:07) [GCC 14.1.1 20240522]
Pwndbg:   2024.02.14 build: 180102b
Capstone: 5.0.1280
Unicorn:  2.0.1
```

Hints from issue discussion:
The `start` command does not cache anything, but it creates a temporary breakpoint on one of the symbols:
https://github.com/pwndbg/pwndbg/blob/6cb96632a4e492a6c943a2a85ea0503b15d35ee4/pwndbg/commands/start.py#L59-L78

It takes the symbol addresses from `pwndbg.gdblib.symbo.address`:
https://github.com/pwndbg/pwndbg/blob/6cb96632a4e492a6c943a2a85ea0503b15d35ee4/pwndbg/gdblib/symbol.py#L139-L164

And this one is cached until a new `objfile` is loaded. I bet this is the thing that is cached and doesn't work properly here.
@Aplet123 Can you check if this works properly when you disable cache for `pwndbg.gdblib.symbol.address` (just comment out its `cache_until("objfile")`?

If so, we should probably do a: `cache_until("start", "objfile")` instead? Does that fix the issue?

We should probably also add a comment like `# We need to cache until "start" as ASLR may change the symbol address` in there.

If so, please send a PR with the fix :). 

(It would also be nice to add a test case for it)
Btw I have looked for other potential bugs like this, and this could be a candidate:
https://github.com/pwndbg/pwndbg/blob/6cb96632a4e492a6c943a2a85ea0503b15d35ee4/pwndbg/gdblib/elf.py#L152-L153

But it isn't a bug because the `vaddr` which is the base address to rebase to is an argument and our cache caches the results for given arguments, so this should work fine, huh :).
The binary I was testing on was stripped, so it went to the fallback case of breaking on entrypoint due to not finding any of the symbols:
https://github.com/pwndbg/pwndbg/blob/6cb96632a4e492a6c943a2a85ea0503b15d35ee4/pwndbg/commands/start.py#L23-L29

Commenting out the cache on `pwndbg.gdblib.elf.entry()` fixed the bug in that case. When testing on a binary that has symbols, it actually works for me as-is, because the breakpoint is set on the symbol, not on the cached address.

~~I think there's a bigger issue at hand though, as I sometimes don't get a start event when running the program using the `start` command. I need to do some more testing to see why this happens.~~ I may have just been hallucinating? I can't reproduce it anymore.
I think there's a bit of an event ordering issue, where the cache is cleared after everything else for the event is run, which means that event handlers can receive a stale cache when running.

Commit message: Fix formatting in files

Patch to review:
diff --git a/pwndbg/gdblib/elf.py b/pwndbg/gdblib/elf.py
index 1444c8f0ffa..cf588c27f4e 100644
--- a/pwndbg/gdblib/elf.py
+++ b/pwndbg/gdblib/elf.py
@@ -240,7 +240,7 @@ def dump_relocations_by_section_name(
 
 
 @pwndbg.gdblib.proc.OnlyWhenRunning
-@pwndbg.lib.cache.cache_until("start")
+@pwndbg.lib.cache.cache_until("start", "objfile")
 def exe() -> Ehdr | None:
     """
     Return a loaded ELF header object pointing to the Ehdr of the
@@ -253,7 +253,7 @@ def exe() -> Ehdr | None:
 
 
 @pwndbg.gdblib.proc.OnlyWhenRunning
-@pwndbg.lib.cache.cache_until("start")
+@pwndbg.lib.cache.cache_until("start", "objfile")
 def entry() -> int:
     """
     Return the address of the entry point for the main executable.
@@ -293,7 +293,7 @@ def load(pointer: int) -> Ehdr | None:
 ehdr_type_loaded = 0
 
 
-@pwndbg.lib.cache.cache_until("start")
+@pwndbg.lib.cache.cache_until("start", "objfile")
 def reset_ehdr_type_loaded() -> None:
     global ehdr_type_loaded
     ehdr_type_loaded = 0
diff --git a/pwndbg/gdblib/events.py b/pwndbg/gdblib/events.py
index cdec7a44a72..d315be261b3 100644
--- a/pwndbg/gdblib/events.py
+++ b/pwndbg/gdblib/events.py
@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 import sys
-from functools import partial
 from functools import wraps
 from typing import Any
 from typing import Callable
@@ -89,6 +88,10 @@ def on_stop(self) -> None:
     gdb.events.register_changed: [],
 }
 
+# Registered events are wrapped and aren't directly connected to GDB
+# This is a map from event to the actual handler connected to GDB
+connected = {}
+
 
 # When performing remote debugging, gdbserver is very noisy about which
 # objects are loaded.  This greatly slows down the debugging session.
@@ -97,7 +100,9 @@ def on_stop(self) -> None:
 objfile_cache: Dict[str, Set[str]] = {}
 
 
-def connect(func: Callable[P, T], event_handler: Any, name: str = "") -> Callable[P, T]:
+def connect(
+    func: Callable[P, T], event_handler: Any, name: str = "", priority: bool = False
+) -> Callable[P, T]:
     if debug:
         print("Connecting", func.__name__, event_handler)
 
@@ -128,44 +133,55 @@ def caller(*a: P.args, **kw: P.kwargs) -> None:
             pwndbg.exception.handle()
             raise e
 
-    registered[event_handler].append(caller)
-    event_handler.connect(caller)
+    if priority:
+        registered[event_handler].insert(0, caller)
+    else:
+        registered[event_handler].append(caller)
+    if event_handler not in connected:
+
+        def handle(*a: P.args, **kw: P.kwargs) -> None:
+            for f in registered[event_handler]:
+                f(*a, **kw)
+
+        event_handler.connect(handle)
+        connected[event_handler] = handle
     return func
 
 
-def exit(func: Callable[[], T]) -> Callable[[], T]:
-    return connect(func, gdb.events.exited, "exit")
+def exit(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.exited, "exit", **kwargs)
 
 
-def cont(func: Callable[[], T]) -> Callable[[], T]:
-    return connect(func, gdb.events.cont, "cont")
+def cont(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.cont, "cont", **kwargs)
 
 
-def new_objfile(func: Callable[[], T]) -> Callable[[], T]:
-    return connect(func, gdb.events.new_objfile, "obj")
+def new_objfile(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.new_objfile, "obj", **kwargs)
 
 
-def stop(func: Callable[[], T]) -> Callable[[], T]:
-    return connect(func, gdb.events.stop, "stop")
+def stop(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.stop, "stop", **kwargs)
 
 
-def start(func: Callable[[], T]) -> Callable[[], T]:
-    return connect(func, gdb.events.start, "start")
+def start(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.start, "start", **kwargs)
 
 
-def thread(func: Callable[[], T]) -> Callable[[], T]:
-    return connect(func, gdb.events.new_thread, "thread")
+def thread(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.new_thread, "thread", **kwargs)
 
 
-before_prompt = partial(connect, event_handler=gdb.events.before_prompt, name="before_prompt")
+def before_prompt(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.before_prompt, "before_prompt", **kwargs)
 
 
-def reg_changed(func: Callable[[], T]) -> Callable[[], T]:
-    return connect(func, gdb.events.register_changed, "reg_changed")
+def reg_changed(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.register_changed, "reg_changed", **kwargs)
 
 
-def mem_changed(func: Callable[[], T]) -> Callable[[], T]:
-    return connect(func, gdb.events.memory_changed, "mem_changed")
+def mem_changed(func: Callable[[], T], **kwargs: Any) -> Callable[[], T]:
+    return connect(func, gdb.events.memory_changed, "mem_changed", **kwargs)
 
 
 def log_objfiles(ofile: gdb.NewObjFileEvent | None = None) -> None:
@@ -182,28 +198,24 @@ def log_objfiles(ofile: gdb.NewObjFileEvent | None = None) -> None:
 gdb.events.new_objfile.connect(log_objfiles)
 
 
+# invoke all registered handlers of a certain event type
+def invoke_event(event: Any) -> None:
+    if event in registered:
+        for f in registered[event]:
+            f()
+
+
 def after_reload(start: bool = True) -> None:
     if gdb.selected_inferior().pid:
-        if gdb.events.stop in registered:
-            for f in registered[gdb.events.stop]:
-                f()
-        if gdb.events.start in registered:
-            for f in registered[gdb.events.start]:
-                if start:
-                    f()
-        if gdb.events.new_objfile in registered:
-            for f in registered[gdb.events.new_objfile]:
-                f()
-        if gdb.events.before_prompt in registered:
-            for f in registered[gdb.events.before_prompt]:
-                f()
+        invoke_event(gdb.events.stop)
+        invoke_event(gdb.events.start)
+        invoke_event(gdb.events.new_objfile)
+        invoke_event(gdb.events.before_prompt)
 
 
 def on_reload() -> None:
-    for event, functions in registered.items():
-        for function in functions:
-            event.disconnect(function)
-        registered[event] = []
+    for functions in registered.values():
+        functions.clear()
 
 
 @new_objfile
diff --git a/pwndbg/lib/cache.py b/pwndbg/lib/cache.py
index cb56d665fdc..a447e3d879b 100644
--- a/pwndbg/lib/cache.py
+++ b/pwndbg/lib/cache.py
@@ -75,7 +75,7 @@ def connect_event_hooks(self, event_hooks: Tuple[Any, ...]) -> None:
         by `stop`, `mem_changed` and `reg_changed` events.
         """
         for event_hook in event_hooks:
-            event_hook(self.clear)
+            event_hook(self.clear, priority=True)
 
     def clear(self) -> None:
         for cache in self.caches:
diff --git a/tests/gdb-tests/tests/test_cache.py b/tests/gdb-tests/tests/test_cache.py
index 9d6d4c7fcd8..1a3a515a6ea 100644
--- a/tests/gdb-tests/tests/test_cache.py
+++ b/tests/gdb-tests/tests/test_cache.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import pwndbg.gdblib.events
 import tests
 from pwndbg.lib import cache
 
@@ -68,3 +69,28 @@ def foo(arg0, *args, **kwargs):
 
     assert foo("abc") == (5, "abc", (), {}) and x == 5
     assert foo(100, 200) == (6, 100, (200,), {}) and x == 6
+
+
+def test_cache_clear_has_priority(start_binary):
+    actions = []
+
+    @pwndbg.gdblib.events.stop
+    def on_stop():
+        actions.append("on_stop")
+        # test to make sure event handlers don't have a stale cache
+        foo()
+
+    @cache.cache_until("stop")
+    def foo():
+        actions.append("foo")
+
+    foo()
+    foo()
+    assert actions == ["foo"]
+
+    start_binary(BINARY)
+    assert actions == ["foo", "on_stop", "foo"]
+
+    foo()
+    foo()
+    assert actions == ["foo", "on_stop", "foo"]


Code review comments:
diff_hunk: @@ -97,7 +100,9 @@ def on_stop(self) -> None:
 objfile_cache: Dict[str, Set[str]] = {}
 
 
-def connect(func: Callable[P, T], event_handler: Any, name: str = "") -> Callable[P, T]:
+def connect(
+    func: Callable[P, T], event_handler: Any, name: str = "", priority: bool = False
path: pwndbg/gdblib/events.py
text: @user2:
@user1 thoughts on this? I believe we need it so that clearing cache is triggered as very first event, but I'm not sure I like a boolean value for it. I mean, this eventually means calling multiple 'priority' events will end up with the same issue as we had before.

@user1:
Yea, `priority` is usually an enum or an int.

Also, my general thought about adding priority to these events is that we need to really think about the right way to do this, otherwise this is going to turn into a mess and get completely rewritten at some point when it becomes hard to maintain, or doesn't solve the problem well in the long run. And I don't have the time to think through it. But I guess if it's fixing an actual bug and we don't have an alternative solution it makes sense to merge it? Up to you.

@author:
I can switch it to an enum priority. Are there any other concerns about the implementation?

@user2:
Let's merge this as of now. I don't think we need to make it more complex with int/enum for now :).

@user2:
Oh, we already made a switch. I guess that is fine.