From 00adfbbb5e06564f7d351624e09dae898f485c57 Mon Sep 17 00:00:00 2001
From: Gulshan Singh <gsingh2011@gmail.com>
Date: Mon, 6 May 2024 08:58:04 -0700
Subject: [PATCH] Only look for readable address in retaddr command (#2143)

* Only look for readable address in retaddr command

* Rename stack.py to retaddr.py

* Add pwndbg.gdblib.stack.callstack and use it in retaddr

* Add callstack gdb test

* Add QEMU callstack test
---
 pwndbg/commands/__init__.py                   |  2 +-
 pwndbg/commands/{stack.py => retaddr.py}      | 13 +++---------
 pwndbg/gdblib/stack.py                        | 17 ++++++++++++++++
 tests/gdb-tests/tests/test_callstack.py       | 20 +++++++++++++++++++
 .../qemu-tests/tests/system/test_callstack.py | 11 ++++++++++
 5 files changed, 52 insertions(+), 11 deletions(-)
 rename pwndbg/commands/{stack.py => retaddr.py} (78%)
 create mode 100644 tests/gdb-tests/tests/test_callstack.py
 create mode 100644 tests/qemu-tests/tests/system/test_callstack.py

diff --git a/pwndbg/commands/__init__.py b/pwndbg/commands/__init__.py
index 363d8a392..b65ab9ef6 100644
--- a/pwndbg/commands/__init__.py
+++ b/pwndbg/commands/__init__.py
@@ -697,6 +697,7 @@ def load_commands() -> None:
     import pwndbg.commands.procinfo
     import pwndbg.commands.radare2
     import pwndbg.commands.reload
+    import pwndbg.commands.retaddr
     import pwndbg.commands.rizin
     import pwndbg.commands.rop
     import pwndbg.commands.ropper
@@ -706,7 +707,6 @@ def load_commands() -> None:
     import pwndbg.commands.sigreturn
     import pwndbg.commands.slab
     import pwndbg.commands.spray
-    import pwndbg.commands.stack
     import pwndbg.commands.start
     import pwndbg.commands.telescope
     import pwndbg.commands.tips
diff --git a/pwndbg/commands/stack.py b/pwndbg/commands/retaddr.py
similarity index 78%
rename from pwndbg/commands/stack.py
rename to pwndbg/commands/retaddr.py
index cc44d5fc6..bcb351d50 100644
--- a/pwndbg/commands/stack.py
+++ b/pwndbg/commands/retaddr.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import gdb
-
 import pwndbg.chain
 import pwndbg.commands
 import pwndbg.gdblib.arch
@@ -15,17 +13,12 @@ from pwndbg.commands import CommandCategory
 )
 @pwndbg.commands.OnlyWhenRunning
 def retaddr() -> None:
+    addresses = pwndbg.gdblib.stack.callstack()
+
     sp = pwndbg.gdblib.regs.sp
     stack = pwndbg.gdblib.vmmap.find(sp)
 
-    # Enumerate all return addresses
-    frame = gdb.newest_frame()
-    addresses = []
-    while frame:
-        addresses.append(int(frame.pc()))
-        frame = frame.older()
-
-    # Find all of them on the stack
+    # Find all return addresses on the stack
     start = stack.vaddr
     stop = start + stack.memsz
     while addresses and start < sp < stop:
diff --git a/pwndbg/gdblib/stack.py b/pwndbg/gdblib/stack.py
index 9172bc8c5..22f9b5726 100644
--- a/pwndbg/gdblib/stack.py
+++ b/pwndbg/gdblib/stack.py
@@ -7,6 +7,8 @@ binaries do things to remap the stack (e.g. pwnies' postit).
 
 from __future__ import annotations
 
+from typing import List
+
 import gdb
 
 import pwndbg.gdblib.abi
@@ -153,3 +155,18 @@ def _fetch_via_exploration() -> dict[int, pwndbg.lib.memory.Page]:
     curr_thread.switch()
 
     return stacks
+
+
+def callstack() -> List[int]:
+    """
+    Return the address of the return address for the current frame.
+    """
+    frame = gdb.newest_frame()
+    addresses = []
+    while frame:
+        addr = int(frame.pc())
+        if pwndbg.gdblib.memory.is_readable_address(addr):
+            addresses.append(addr)
+        frame = frame.older()
+
+    return addresses
diff --git a/tests/gdb-tests/tests/test_callstack.py b/tests/gdb-tests/tests/test_callstack.py
new file mode 100644
index 000000000..bc877d1f7
--- /dev/null
+++ b/tests/gdb-tests/tests/test_callstack.py
@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import gdb
+
+import pwndbg.gdblib.memory
+import pwndbg.gdblib.stack
+import tests
+
+REFERENCE_BINARY = tests.binaries.get("reference-binary.out")
+
+
+def test_callstack_readable(start_binary):
+    start_binary(REFERENCE_BINARY)
+    gdb.execute("b break_here")
+    gdb.execute("r")
+
+    addresses = pwndbg.gdblib.stack.callstack()
+
+    assert len(addresses) > 0
+    assert all(pwndbg.gdblib.memory.is_readable_address(address) for address in addresses)
diff --git a/tests/qemu-tests/tests/system/test_callstack.py b/tests/qemu-tests/tests/system/test_callstack.py
new file mode 100644
index 000000000..3efdebed1
--- /dev/null
+++ b/tests/qemu-tests/tests/system/test_callstack.py
@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import pwndbg.gdblib.memory
+import pwndbg.gdblib.stack
+
+
+def test_callstack_readable():
+    addresses = pwndbg.gdblib.stack.callstack()
+
+    assert len(addresses) > 0
+    assert all(pwndbg.gdblib.memory.is_readable_address(address) for address in addresses)