Add segv test

lib/pyda/process.py

@@ -162,10 +162,28 @@ def call(*args):
             self.write(self.regs.rsp, orig_rip.to_bytes(8, "little"))
 
             set_regs_for_call_linux_x86(self, args)
-            ## END ARCH-SPECIFIC SETUP
+            target_rsp = self.regs.rsp + 8
 
-            self.run_from_to(addr, orig_rip)
-            self._p.pop_state()
+            self.regs.rip = addr
+
+            # This is a bit hacky, but basically
+            # we don't actually know that orig_rip is outside
+            # of the function, we just know it's a reasonably
+            # safe address. You'll get unexpectedly bad perf
+            # if your original RIP is garbage
+
+            count = 0
+            try:
+                while self.regs.rsp != target_rsp:
+                    self.run_until(orig_rip)
+                    count += 1
+                ## END ARCH-SPECIFIC SETUP
+
+            finally:
+                if count > 1:
+                    self.warning(f"WARN: Callable should be used from a safe RIP not within the callee.")
+
+                self._p.pop_state()
 
         return call
 

pyda_core/tool.c

@@ -266,29 +266,33 @@ static dr_signal_action_t signal_event(void *drcontext, dr_siginfo_t *siginfo) {
     // Perhaps unexpectedly, we also only care if the process has not blocked the signal.
     // This prevents us from handling signals when the application has blocked them (e.g.,
     // because it is holding the GIL. We will still handle them before the app gets them.)
-    if ((sig == SIGSEGV || sig == SIGBUS || sig == SIGILL) && !t->python_exited && !siginfo->blocked) {
-        memcpy(&t->cur_context, siginfo->mcontext, sizeof(dr_mcontext_t));
-        t->signal = sig;
-
-        // Clear any previous run_until hooks: they are now invalid
-        // since we are throwing.
-        if (t->run_until)
-            pyda_clear_run_until(t);
-
-        // Raise an exception in Python +
-        // Wait for Python to yield back to us
-        pyda_break(t);
-
-        // Flushing is actually allowed in signal event handlers.
-        // This updates run_until handlers, updated hooks, etc.
-        pyda_flush_hooks();
-
-        // Copy the state back to the siginfo
-        memcpy(siginfo->mcontext, &t->cur_context, sizeof(dr_mcontext_t));
-
-        t->signal = 0;
-
-        return DR_SIGNAL_REDIRECT;
+    if ((sig == SIGSEGV || sig == SIGBUS || sig == SIGILL) && !siginfo->blocked) {
+        if (!t->python_exited) {
+            memcpy(&t->cur_context, siginfo->mcontext, sizeof(dr_mcontext_t));
+            t->signal = sig;
+
+            // Clear any previous run_until hooks: they are now invalid
+            // since we are throwing.
+            if (t->run_until)
+                pyda_clear_run_until(t);
+
+            // Raise an exception in Python +
+            // Wait for Python to yield back to us
+            pyda_break(t);
+
+            // Flushing is actually allowed in signal event handlers.
+            // This updates run_until handlers, updated hooks, etc.
+            pyda_flush_hooks();
+
+            // Copy the state back to the siginfo
+            memcpy(siginfo->mcontext, &t->cur_context, sizeof(dr_mcontext_t));
+
+            t->signal = 0;
+
+            return DR_SIGNAL_REDIRECT;
+        } else {
+            dr_fprintf(STDERR, "[Pyda] ERROR: Signal %d received after Python exited/died. Add p.run() to receive the signal as an exception.\n", sig);
+        }
     }
 
     return DR_SIGNAL_DELIVER;

tests/run_tests.py

@@ -205,6 +205,15 @@ def no_warnings_or_errors(stdout: bytes, stderr: bytes) -> bool:
             lambda o, e: o.count(b"pass\n") == 1,
         ]
     )),
+
+    ("test_segv", "test_segv.c", "test_segv.py", RunOpts(), ExpectedResult(
+        retcode=0,
+        checkers=[
+            output_checker,
+            no_warnings_or_errors,
+            lambda o, e: o.count(b"pass\n") == 1,
+        ]
+    )),
 ]
 
 def main():

tests/test_segv.c

@@ -0,0 +1,11 @@
+#include <stdio.h>
+
+void segfault() {
+    *(volatile int*)0 = 0;
+}
+
+int main() {
+    printf("Hello, World!\n");
+    segfault();
+    return 0;
+}

tests/test_segv.py

@@ -0,0 +1,51 @@
+from pyda import *
+from pwnlib.elf.elf import ELF
+from pwnlib.util.packing import u64
+import string
+import sys, time
+
+p = process()
+
+e = ELF(p.exe_path)
+e.address = p.maps[p.exe_path].base
+
+excepted = False
+try:
+    p.run()
+except FatalSignalError as err:
+    assert err.args[0] is not None
+    excepted = True
+
+assert excepted
+
+print("Exception 1")
+
+# Now, after the exception, redirect execution to main again
+p.regs.rip = e.symbols["main"]
+excepted = False
+try:
+    p.run()
+except FatalSignalError as err:
+    assert err.args[0] is not None
+    excepted = True
+
+assert excepted
+
+print("Exception 2")
+
+# Finally, let's try calling out to the function as a callable
+excepted = False
+try:
+    p.callable(e.symbols["segfault"])()
+except FatalSignalError as err:
+    assert err.args[0] is not None
+    excepted = True
+
+assert excepted
+
+print("Exception 3")
+
+p.regs.rsp += 0x18
+p.regs.rip = e.symbols["main"] + 0x26
+p.run()
+print("pass")