WIP: handle signals as exceptions (#8)

lib/pyda/base.py

@@ -1,5 +1,5 @@
 import pyda_core
-from pyda_core import MemoryError, ThreadExitError, InvalidStateError
+from pyda_core import MemoryError, ThreadExitError, InvalidStateError, FatalSignalError
 from .process import Process, Map
 from . import arch
 import sys

lib/pyda/process.py

@@ -107,7 +107,7 @@ def write(self, addr, data):
 
     def __getattr__(self, name):
         # TODO: Move these into CPython extension?
-        if name in "regs":
+        if name == "regs":
             return ProcessRegisters(self._p)
         elif name == "mem":
             return ProcessMemory(self)
@@ -118,6 +118,12 @@ def __getattr__(self, name):
 
         raise AttributeError(f"Invalid attribute '{name}'. Did you mean 'regs.{name}'?")
 
+    def __setattr__(self, name, value):
+        if not name.startswith("_") and name not in ["timeout", "buffer", "closed"]:
+            raise AttributeError(f"Cannot set attribute '{name}'")
+
+        super().__setattr__(name, value)
+
     def run(self):
         self._has_run = True
         self._p.run()
@@ -151,7 +157,8 @@ def call(*args):
             orig_rip = self.regs.rip
 
             # Push orig_rip as the return address
-            self.regs.rsp -= 16
+            self.regs.rsp &= ~0xf
+            self.regs.rsp -= 8
             self.write(self.regs.rsp, orig_rip.to_bytes(8, "little"))
 
             set_regs_for_call_linux_x86(self, args)

lib/pyda/tube.py

@@ -16,8 +16,8 @@ def __init__(self, stdin_fd, stdout_fd, **kwargs):
         else:
             self._captured = True
 
-        self.stdin_fd = stdin_fd
-        self.stdout_fd = stdout_fd
+        self._stdin_fd = stdin_fd
+        self._stdout_fd = stdout_fd
 
     # Overwritten for better usability
     def recvall(self, timeout = None):
@@ -40,7 +40,7 @@ def recv_raw(self, numb, *a):
 
         while True:
             try:
-                data = os.read(self.stdout_fd, numb)
+                data = os.read(self._stdout_fd, numb)
                 break
             except IOError as e:
                 if e.errno == errno.EAGAIN:
@@ -79,7 +79,7 @@ def send_raw(self, data):
         ptr = 0
         while ptr < len(data):
             try:
-                count = os.write(self.stdin_fd, data[ptr:])
+                count = os.write(self._stdin_fd, data[ptr:])
                 ptr += count
             except IOError as e:
                 eof_numbers = (errno.EPIPE, errno.ECONNRESET, errno.ECONNREFUSED)
@@ -108,9 +108,9 @@ def can_recv_raw(self, timeout):
 
         try:
             if timeout is None:
-                return select.select([self.stdout_fd], [], []) == ([self.stdout_fd], [], [])
+                return select.select([self._stdout_fd], [], []) == ([self._stdout_fd], [], [])
 
-            return select.select([self.stdout_fd], [], [], timeout) == ([self.stdout_fd], [], [])
+            return select.select([self._stdout_fd], [], [], timeout) == ([self._stdout_fd], [], [])
         except ValueError:
             # Not sure why this isn't caught when testing self.proc.stdout.closed,
             # but it's not.

pyda_core/pyda_core.c

@@ -37,6 +37,7 @@ pyda_process* pyda_mk_process() {
     proc->dirty_hooks = 0;
     drvector_init(&proc->threads, 0, true, NULL);
     drvector_init(&proc->thread_run_untils, 0, true, NULL);
+    drvector_init(&proc->hook_delete_queue, 0, true, NULL);
 
     proc->main_thread = pyda_mk_thread(proc);
     hashtable_init_ex(&proc->callbacks, 4, HASH_INTPTR, false, false, free_hook, NULL, NULL);
@@ -211,6 +212,7 @@ pyda_thread* pyda_mk_thread(pyda_process *proc) {
     thread->errored = 0;
     thread->python_blocked_on_io = 0;
     thread->run_until = 0;
+    thread->signal = 0;
     drvector_init(&thread->context_stack, 0, true, free_context);
 
     drvector_append(&proc->threads, thread);
@@ -374,6 +376,7 @@ void pyda_add_hook(pyda_process *t, uint64_t addr, PyObject *callback) {
 void pyda_remove_hook(pyda_process *p, uint64_t addr) {
     hashtable_remove(&p->callbacks, (void*)addr);
     p->dirty_hooks = 1;
+    drvector_append(&p->hook_delete_queue, (void*)addr);
 }
 
 void pyda_set_thread_init_hook(pyda_process *p, PyObject *callback) {
@@ -432,6 +435,14 @@ int pyda_flush_hooks() {
         hashtable_apply_to_all_payloads(&p->callbacks, flush_hook);
         p->dirty_hooks = 0;
         flushed = 1;
+
+        // Flush deleted hooks
+        drvector_lock(&p->hook_delete_queue);
+        for (int i=0; i<p->hook_delete_queue.entries; i++) {
+            dr_flush_region((void*)p->hook_delete_queue.array[i], 1);
+        }
+        p->hook_delete_queue.entries = 0;
+        drvector_unlock(&p->hook_delete_queue);
     }
 
     return flushed;

pyda_core/pyda_core.h

@@ -51,6 +51,7 @@ struct pyda_process_s {
     hashtable_t callbacks;
     drvector_t threads;
     drvector_t thread_run_untils; // vec of pcs
+    drvector_t hook_delete_queue;
 #endif
 
 };
@@ -78,6 +79,8 @@ struct pyda_thread_s {
     uint64_t run_until;
     int dirty_run_until;
 
+    int signal; // 0 if no signal, otherwise the signal number
+
 #ifdef PYDA_DYNAMORIO_CLIENT
     dr_mcontext_t cur_context;
     drvector_t context_stack;

pyda_core/pyda_core_py.c

@@ -66,6 +66,7 @@ static struct PyModuleDef pyda_module = {
 static PyObject *MemoryError;
 static PyObject *ThreadExitError;
 static PyObject *InvalidStateError;
+static PyObject *FatalSignalError;
 
 static void register_exception(PyObject *mod, PyObject **target, const char *fullname, const char *name) {
     *target = PyErr_NewException(fullname, NULL, NULL);
@@ -83,6 +84,7 @@ PyInit_pyda_core(void) {
     register_exception(m, &MemoryError, "pyda.MemoryError", "MemoryError");
     register_exception(m, &ThreadExitError, "pyda.ThreadExitError", "ThreadExitError");
     register_exception(m, &InvalidStateError, "pyda.InvalidStateError", "InvalidStateError");
+    register_exception(m, &FatalSignalError, "pyda.FatalSignalError", "FatalSignalError");
 
 #ifdef X86
     PyModule_AddIntConstant(m, "REG_RAX", DR_REG_RAX);
@@ -254,6 +256,13 @@ PydaProcess_run(PyObject* self, PyObject *noarg) {
 #endif // PYDA_DYNAMORIO_CLIENT
     Py_END_ALLOW_THREADS
 
+    if (t->signal) {
+        PyObject *tuple = PyTuple_New(1);
+        PyTuple_SetItem(tuple, 0, PyLong_FromLong(t->signal));
+        PyErr_SetObject(FatalSignalError, tuple);
+        return NULL;
+    }
+
     Py_INCREF(Py_None);
     return Py_None;
 }
@@ -281,6 +290,13 @@ PydaProcess_run_until_io(PyObject* self, PyObject *noarg) {
         return NULL;
     }
 
+    if (t->signal) {
+        PyObject *tuple = PyTuple_New(1);
+        PyTuple_SetItem(tuple, 0, PyLong_FromLong(t->signal));
+        PyErr_SetObject(FatalSignalError, tuple);
+        return NULL;
+    }
+
     Py_INCREF(Py_None);
     return Py_None;
 }
@@ -333,6 +349,13 @@ PydaProcess_run_until_pc(PyObject* self, PyObject *args) {
         return NULL;
     }
 
+    if (t->signal) {
+        PyObject *tuple = PyTuple_New(1);
+        PyTuple_SetItem(tuple, 0, PyLong_FromLong(t->signal));
+        PyErr_SetObject(FatalSignalError, tuple);
+        return NULL;
+    }
+
     Py_INCREF(Py_None);
     return Py_None;
 }

pyda_core/tool.c

@@ -8,6 +8,8 @@
 #include "Python.h"
 #include "util.h"
 
+#include <signal.h>
+
 #include "pyda_core_py.h"
 #include "pyda_core.h"
 #include "pyda_threads.h"
@@ -33,6 +35,8 @@ event_insert(void *drcontext, void *tag, instrlist_t *bb, instr_t *instr,
 static bool filter_syscall_event(void *drcontext, int sysnum);
 static bool pre_syscall_event(void *drcontext, int sysnum);
 static void post_syscall_event(void *drcontext, int sysnum);
+static dr_signal_action_t signal_event(void *drcontext, dr_siginfo_t *siginfo);
+
 
 extern int is_dynamorio_running;
 
@@ -78,11 +82,14 @@ dr_client_main(client_id_t id, int argc, const char *argv[])
     drmgr_register_post_syscall_event(post_syscall_event);
     dr_register_filter_syscall_event(filter_syscall_event);
 
+    drmgr_register_signal_event(signal_event);
+
     pthread_cond_init(&python_thread_init1, 0);
 
     g_pyda_tls_idx = drmgr_register_tls_field();
     g_pyda_tls_is_python_thread_idx = drmgr_register_tls_field();
 }
+
 void module_load_event(void *drcontext, const module_data_t *mod, bool loaded) {
     DEBUG_PRINTF("module_load_event: %s\n", mod->full_path);
 }
@@ -249,6 +256,44 @@ static void post_syscall_event(void *drcontext, int sysnum) {
     pyda_hook_syscall(sysnum, 0);
 }
 
+static dr_signal_action_t signal_event(void *drcontext, dr_siginfo_t *siginfo) {
+    pyda_thread *t = drmgr_get_tls_field(drcontext, g_pyda_tls_idx);
+
+    int sig = siginfo->sig;
+
+    // We only care about signals that indicate crashes. We only care if the python thread
+    // is still running (We need to have someone to raise the exception to!)
+    // 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;
+    }
+
+    return DR_SIGNAL_DELIVER;
+}
+
 static void thread_entrypoint_break() {
     DEBUG_PRINTF("entrypoint (break)\n");
 
@@ -335,7 +380,8 @@ void python_main_thread(void *arg) {
     PyEval_SaveThread(); // release GIL
 
     if (!t->app_exited) {
-        dr_fprintf(STDERR, "[Pyda] ERROR: Did you forget to call p.run()?\n");
+        if (t->yield_count == 0)
+            dr_fprintf(STDERR, "[Pyda] ERROR: Did you forget to call p.run()?\n");
         pyda_yield(t); // unblock (note: blocking)
         DEBUG_PRINTF("Implicit pyda_yield finished\n");
     }
@@ -384,6 +430,8 @@ void python_aux_thread(void *arg) {
     dr_client_thread_set_suspendable(true);
     DEBUG_PRINTF("python_aux_thread 4\n");
 
+    t->python_exited = 1;
+
     if (!t->app_exited) {
         pyda_yield(t);
         DEBUG_PRINTF("Implicit pyda_yield finished\n");