feat: multithread. WIP

examples/multithread.py

@@ -0,0 +1,46 @@
+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
+
+plt_map = { e.plt[x]: x for x in e.plt }
+
+def guess_arg(x):
+    printable_chars = bytes(string.printable, 'ascii')
+
+    # Is pointer?
+    if x > 0x100000000:
+        try:
+            data = p.read(x, 0x20)
+            if all([c in printable_chars for c in data[:4]]):
+                return str(data[:data.index(0)])
+        except:
+            pass
+
+    return hex(x)
+
+def lib_hook(p):
+    name = plt_map[p.regs.rip]
+    print(f"[thread {p.tid}] {name}(" + ", ".join([
+        f"rdi={guess_arg(p.regs.rdi)}",
+        f"rsi={guess_arg(p.regs.rsi)}",
+        f"rdx={guess_arg(p.regs.rdx)}",
+        f"rcx={guess_arg(p.regs.rcx)}",
+    ]) + ")")
+
+def thread_entry(p):
+    print(f"thread_entry for {p.tid}")
+
+for x in e.plt:
+    p.hook(e.plt[x], lib_hook)
+
+p.set_thread_entry(thread_entry)
+
+
+p.run()

lib/pyda/process.py

@@ -39,7 +39,7 @@ def after_call_hook(p):
         self.hook(addr, call_hook)
 
     def set_thread_entry(self, callback):
-        self._p.set_thread_init_hook(callback)
+        self._p.set_thread_init_hook(lambda p: callback(self))
 
     def read(self, addr, size):
         return self._p.read(addr, size)

pyda_core/pyda_core.c

@@ -18,11 +18,12 @@ pyda_process* pyda_mk_process() {
     ABORT_IF_NODYNAMORIO;
 
     pyda_process *proc = dr_global_alloc(sizeof(pyda_process));
-    proc->refcount = 2;
+    proc->refcount = 1;
     proc->dirty_hooks = 0;
     proc->main_thread = pyda_mk_thread(proc);
     proc->callbacks = NULL;
     proc->thread_init_hook = NULL;
+    proc->py_obj = NULL;
     return proc;
 }
 
@@ -44,11 +45,14 @@ pyda_thread* pyda_mk_thread(pyda_process *proc) {
     thread->python_yielded = 0;
     thread->app_yielded = 0;
     thread->proc = proc;
+    thread->proc->refcount++; // TODO: lock
+    thread->yield_count = 0;
 
     static volatile unsigned int tid = 0;
     thread->tid = dr_atomic_add32_return_sum(&tid, 1);
     thread->rip_updated_in_cleancall = 0;
     thread->skip_next_hook = 0;
+    thread->python_exited = 0;
 
     // PyErr_SetString(PyExc_RuntimeError, "OK");
     return thread;
@@ -74,6 +78,7 @@ void pyda_thread_destroy(pyda_thread *t) {
 // yield from python to the executable
 void pyda_yield(pyda_thread *t) {
     t->python_yielded = 1;
+    t->yield_count++;
     pthread_cond_signal(&t->resume_cond);
     DEBUG_PRINTF("pyda_yield\n");
 
@@ -206,7 +211,7 @@ void pyda_hook_cleancall(pyda_hook *cb) {
     t->cur_context.pc = (app_pc)cb->addr;
     t->rip_updated_in_cleancall = 0;
 
-    PyObject *result = PyObject_CallFunctionObjArgs(cb->py_func, t->py_obj, NULL);
+    PyObject *result = PyObject_CallFunctionObjArgs(cb->py_func, t->proc->py_obj, NULL);
     if (result == NULL) {
         PyErr_Print();
         dr_fprintf(STDERR, "[Pyda] ERROR: Hook call failed. Aborting.\n");

pyda_core/pyda_core.h

@@ -36,6 +36,7 @@ struct pyda_process_s {
 
     pyda_thread *main_thread;
     PyObject *thread_init_hook;
+    PyObject *py_obj;
 };
 
 struct pyda_thread_s {
@@ -49,10 +50,12 @@ struct pyda_thread_s {
     void* start_pc;
 
     pyda_process *proc;
-    PyObject *py_obj;
+    // PyObject *py_obj;
 
     int rip_updated_in_cleancall;
     int skip_next_hook;
+    int python_exited;
+    int yield_count;
 
 #ifdef PYDA_DYNAMORIO_CLIENT
     dr_mcontext_t cur_context;

pyda_core/pyda_core_py.c

@@ -121,12 +121,17 @@ pyda_core_process(PyObject *self, PyObject *args, PyObject *kwargs) {
         return NULL;
     }
 
+    if (t->proc->py_obj) {
+        PyErr_SetString(PyExc_RuntimeError, "You may only call process() once");
+        return NULL;
+    }
+
     PyType_Ready(&PydaProcess_Type);
     result = PyObject_NEW(PydaProcess, &PydaProcess_Type);
     if (result != NULL)
         result->main_thread = t;
 
-    result->main_thread->py_obj = (PyObject*)result;
+    t->proc->py_obj = (PyObject*)result;
 
     PyBuffer_Release(&bin_path);
     return (PyObject*)result;
@@ -335,7 +340,6 @@ PydaProcess_register_hook(PyObject *self, PyObject *args) {
 #ifdef PYDA_DYNAMORIO_CLIENT
     DEBUG_PRINTF("register_hook: %llx\n", addr);
 #endif // PYDA_DYNAMORIO_CLIENT
-    Py_INCREF(callback);
     pyda_add_hook(p->main_thread->proc, addr, callback);
 
     Py_INCREF(Py_None);

pyda_core/tool.c

@@ -92,6 +92,11 @@ void thread_init_event(void *drcontext) {
 
     drmgr_set_tls_field(drcontext, g_pyda_tls_idx, (void*)t);
 
+    if (global_proc->main_thread->python_exited) {
+        pyda_thread_destroy(t); // decrement refcount
+        return;
+    }
+
     // Every thread has its own corresponding python thread
     if (t == global_proc->main_thread) {
         python_init();
@@ -110,19 +115,20 @@ void thread_init_event(void *drcontext) {
         mc.flags = DR_MC_ALL;
         dr_get_mcontext(drcontext, &mc);
         t->start_pc = (void*)mc.rip;
-        DEBUG_PRINTF("start_pc: %p\n", t->start_pc);
-        dr_flush_region(t->start_pc, 1);
+        // DEBUG_PRINTF("start_pc: %p\n", t->start_pc);
+        // dr_flush_region(t->start_pc, 1);
     }
-    DEBUG_PRINTF("thread_init_event: %p\n", t->start_pc);
+    DEBUG_PRINTF("thread_init_event: %p (entry %p)\n", t, t->start_pc);
 }
 
 void thread_exit_event(void *drcontext) {
     DEBUG_PRINTF("thread_exit_event\n"); 
     pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
     DEBUG_PRINTF("thread_exit_event: %p\n", t);
-    dr_abort();
 
-    pyda_break(t);
+    if (!t->python_exited)
+        pyda_break(t);
+
     DEBUG_PRINTF("broke end\n", t);
     // TODO: exit event
     pyda_thread_destroy(t);
@@ -274,20 +280,23 @@ void python_main_thread(void *arg) {
     fclose(f);
 
 python_exit:
-
     DEBUG_PRINTF("Script exited...\n");
+    t->python_exited = 1;
 
     PyGILState_Release(gstate);
     pthread_mutex_unlock(&python_thread_init1_mutex);
 
     dr_client_thread_set_suspendable(true);
-    pyda_yield(t); // unblock
+
+    if (t->yield_count == 0) {
+        dr_fprintf(STDERR, "[PYDA] WARN: Did you forget to call p.run()?\n");
+        pyda_yield(t); // unblock (note: blocking)
+    }
 
     dr_thread_free(drcontext, tls, sizeof(void*) * 130);
-    DEBUG_PRINTF("Calling dr_exit\n");
 
     // pyda_thread_destroy(t);
-    drmgr_exit();
+    // drmgr_exit();
 }
 
 void python_aux_thread(void *arg) {
@@ -298,15 +307,29 @@ void python_aux_thread(void *arg) {
     PyGILState_STATE gstate;
     gstate = PyGILState_Ensure();
 
+    DEBUG_PRINTF("python_aux_thread thread_init_hook %p\n", t->proc->thread_init_hook);
+
     // We just call the thread init hook, if one exists
     if (t->proc->thread_init_hook) {
-        PyObject *result = PyObject_CallFunctionObjArgs(t->proc->thread_init_hook, t->proc->main_thread->py_obj, NULL);
+        DEBUG_PRINTF("Calling thread_init_hook\n");
+        PyObject *result = PyObject_CallFunctionObjArgs(t->proc->thread_init_hook, t->proc->py_obj, NULL);
+        if (result == NULL) {
+            PyErr_Print();
+            dr_fprintf(STDERR, "[Pyda] ERROR: Hook call failed. Aborting.\n");
+            dr_abort();
+        }
     }
 
     PyGILState_Release(gstate);
 
     dr_client_thread_set_suspendable(true);
-    pyda_yield(t); // unblock
+    DEBUG_PRINTF("python_aux_thread 4\n");
+
+    // pyda_yield(t); // unblock (note: blocking)
+    t->python_yielded = 1;
+    pthread_cond_signal(&t->resume_cond);
 
+    DEBUG_PRINTF("python_aux_thread 5\n");
     dr_thread_free(drcontext, tls, sizeof(void*) * 130);
+    DEBUG_PRINTF("python_aux_thread 6\n");
 }