feat: multithreading

examples/multithread.py → examples/ltrace_multithreaded.py

@@ -1,3 +1,6 @@
+# This is basically just ltrace.py but we also print the thread number
+# and print a message when threads are created.
+
 from pyda import *
 from pwnlib.elf.elf import ELF
 from pwnlib.util.packing import u64
@@ -37,10 +40,9 @@ def lib_hook(p):
 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)
 
+for x in e.plt:
+    p.hook(e.plt[x], lib_hook)
 
 p.run()

patches/cpython-3.10.12.patch

@@ -12,7 +12,7 @@ index 8d2221cfd8..1372cbc291 100644
 
  #if defined(FAULTHANDLER_USE_ALT_STACK) && defined(HAVE_LINUX_AUXVEC_H) && defined(HAVE_SYS_AUXV_H)
 diff --git a/Python/thread_pthread.h b/Python/thread_pthread.h
-index 35b9810aa3..4085158070 100644
+index 35b9810aa3..130528aec3 100644
 --- a/Python/thread_pthread.h
 +++ b/Python/thread_pthread.h
 @@ -100,7 +100,7 @@
@@ -24,3 +24,12 @@ index 35b9810aa3..4085158070 100644
  #endif
 
 
+@@ -392,7 +392,7 @@ PyThread_allocate_lock(void)
+     lock = (sem_t *)PyMem_RawMalloc(sizeof(sem_t));
+
+     if (lock) {
+-        status = sem_init(lock,0,1);
++        status = sem_init(lock,1,1);
+         CHECK_STATUS("sem_init");
+
+         if (error) {

pyda_core/pyda_core.c

@@ -18,25 +18,60 @@ pyda_process* pyda_mk_process() {
     ABORT_IF_NODYNAMORIO;
 
     pyda_process *proc = dr_global_alloc(sizeof(pyda_process));
-    proc->refcount = 1;
+    proc->refcount = 0; // xxx: will be incremented to 1 by first pyda_mk_thread
     proc->dirty_hooks = 0;
     proc->main_thread = pyda_mk_thread(proc);
     proc->callbacks = NULL;
     proc->thread_init_hook = NULL;
     proc->py_obj = NULL;
+
+    pthread_condattr_t condattr;
+    int ret;
+    if (ret = pthread_condattr_setpshared(&condattr, PTHREAD_PROCESS_SHARED)) {
+        dr_fprintf(STDERR, "pthread_condattr_setpshared failed: %d\n", ret);
+        dr_abort();
+    }
+    if (ret = pthread_cond_init(&proc->thread_exit_cond, &condattr)) {
+        dr_fprintf(STDERR, "pthread_cond_init failed %d\n", ret);
+        dr_abort();
+    }
+
+    pthread_mutexattr_t attr;
+    pthread_mutexattr_init(&attr);
+    pthread_mutexattr_settype(&attr, PTHREAD_MUTEX_NORMAL);
+    pthread_mutexattr_setpshared(&attr, PTHREAD_PROCESS_SHARED);
+    if (ret = pthread_mutex_init(&proc->refcount_mutex, &attr)) {
+        dr_fprintf(STDERR, "pthread_mutex_init failed %d\n", ret);
+        dr_abort();
+    }
     return proc;
 }
 
 pyda_thread* pyda_mk_thread(pyda_process *proc) {
     ABORT_IF_NODYNAMORIO;
 
     pyda_thread *thread = dr_global_alloc(sizeof(pyda_thread));
-    pthread_cond_init(&thread->resume_cond, 0);
-    pthread_cond_init(&thread->break_cond, 0);
+    pthread_condattr_t condattr;
+    pthread_condattr_init(&condattr);
+    int ret;
+    if (ret = pthread_condattr_setpshared(&condattr, PTHREAD_PROCESS_SHARED)) {
+        dr_fprintf(STDERR, "pthread_condattr_setpshared failed: %d\n", ret);
+        dr_abort();
+    }
+    if (ret = pthread_cond_init(&thread->resume_cond, &condattr)) {
+        dr_fprintf(STDERR, "pthread_cond_init failed %d\n", ret);
+        dr_abort();
+    }
+
+    if (pthread_cond_init(&thread->break_cond, &condattr)) {
+        dr_fprintf(STDERR, "pthread_cond_init failed\n");
+        dr_abort();
+    }
 
     pthread_mutexattr_t attr;
     pthread_mutexattr_init(&attr);
     pthread_mutexattr_settype(&attr, PTHREAD_MUTEX_NORMAL);
+    pthread_mutexattr_setpshared(&attr, PTHREAD_PROCESS_SHARED);
     pthread_mutex_init(&thread->mutex, &attr);
 
     // Start with it locked...
@@ -45,69 +80,120 @@ pyda_thread* pyda_mk_thread(pyda_process *proc) {
     thread->python_yielded = 0;
     thread->app_yielded = 0;
     thread->proc = proc;
-    thread->proc->refcount++; // TODO: lock
+
+    dr_atomic_add32_return_sum(&thread->proc->refcount, 1);
+
     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;
+    thread->errored = 0;
 
     // PyErr_SetString(PyExc_RuntimeError, "OK");
     return thread;
 }
 
 void pyda_process_destroy(pyda_process *p) {
+    // We must be holding the GIL lock so we can drop the refs
+    PyGILState_STATE gstate = PyGILState_Ensure();
+
+    DEBUG_PRINTF("pyda_process_destroy\n");
+    if (p->thread_init_hook)
+        Py_DECREF(p->thread_init_hook);
+
+    p->thread_init_hook = NULL;
+
     pyda_hook *cb = p->callbacks;
     while (cb) {
+        Py_DECREF(cb->py_func);
+        cb->py_func = NULL;
+
         void *del = cb;
         cb = cb->next;
+
         dr_global_free(del, sizeof(pyda_hook));
     }
     dr_global_free(p, sizeof(pyda_process));
+
+    PyGILState_Release(gstate);
 }
 
 void pyda_thread_destroy(pyda_thread *t) {
-    if (--t->proc->refcount == 0) {
+    pthread_mutex_lock(&t->proc->refcount_mutex);
+
+    int new_refcount = dr_atomic_add32_return_sum(&t->proc->refcount, -1);
+    if (new_refcount == 0) {
         pyda_process_destroy(t->proc);
+    } else {
+        pthread_cond_signal(&t->proc->thread_exit_cond);
+        pthread_mutex_unlock(&t->proc->refcount_mutex);
     }
+
     dr_global_free(t, sizeof(pyda_thread));
 }
 
+void pyda_thread_destroy_last(pyda_thread *t) {
+    // wait for this thread to be the final thread
+    pthread_mutex_lock(&t->proc->refcount_mutex);
+    while (t->proc->refcount > 1)
+        pthread_cond_wait(&t->proc->thread_exit_cond, &t->proc->refcount_mutex);
+
+    DEBUG_PRINTF("pyda_thread_destroy_last unblock\n")
+    pthread_mutex_unlock(&t->proc->refcount_mutex);
+    pyda_thread_destroy(t);
+}
+
+void pyda_yield_noblock(pyda_thread *t) {
+    t->python_yielded = 1;
+    pthread_mutex_lock(&t->mutex);
+    pthread_cond_signal(&t->resume_cond);
+    pthread_mutex_unlock(&t->mutex);
+}
+
 // 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");
 
     // here we wait for the executable to signal
     // dr_set_safe_for_sync(false);
 
     pthread_mutex_lock(&t->mutex);
+    pthread_cond_signal(&t->resume_cond);
+
     while (!t->app_yielded)
         pthread_cond_wait(&t->break_cond, &t->mutex);
 
-    t->python_yielded = 0;
+    t->app_yielded = 0;
     pthread_mutex_unlock(&t->mutex);
 
     // dr_set_safe_for_sync(true);
 }
 
+void pyda_break_noblock(pyda_thread *t) {
+    t->app_yielded = 1;
+    pthread_mutex_lock(&t->mutex);
+    pthread_cond_signal(&t->break_cond);
+    pthread_mutex_unlock(&t->mutex);
+}
+
 // yield from the executable back to python
 void pyda_break(pyda_thread *t) {
     t->app_yielded = 1;
-    pthread_cond_signal(&t->break_cond);
 
     // here we wait for the python to signal
     // dr_set_safe_for_sync(false);
 
     pthread_mutex_lock(&t->mutex);
+    pthread_cond_signal(&t->break_cond);
+
     while (!t->python_yielded)
         pthread_cond_wait(&t->resume_cond, &t->mutex);
 
-    t->app_yielded = 0;
+    t->python_yielded = 0;
     pthread_mutex_unlock(&t->mutex);
     // dr_set_safe_for_sync(true);
 }
@@ -118,6 +204,7 @@ void pyda_initial_break(pyda_thread *t) {
     while (!t->python_yielded)
         pthread_cond_wait(&t->resume_cond, &t->mutex);
 
+    t->python_yielded = 0;
     pthread_mutex_unlock(&t->mutex);
     // dr_set_safe_for_sync(true);
 }
@@ -130,6 +217,10 @@ PyObject *pyda_run_until(pyda_thread *proc, uint64_t addr) {
 void pyda_add_hook(pyda_process *t, uint64_t addr, PyObject *callback) {
     pyda_hook *cb = dr_global_alloc(sizeof(pyda_hook));
     cb->py_func = callback;
+
+    Py_INCREF(callback);
+    DEBUG_PRINTF("pyda_add_hook %p %p\n", cb, cb->py_func);
+
     cb->callback_type = 0;
     cb->next = t->callbacks;
     cb->addr = (void*)addr;
@@ -157,8 +248,7 @@ void pyda_remove_hook(pyda_process *p, uint64_t addr) {
 }
 
 void pyda_set_thread_init_hook(pyda_process *p, PyObject *callback) {
-    // TODO: hold some global lock here, just in case this gets called
-    // other than at startup
+    // NOTE: GIL is held
 
     if (p->thread_init_hook)
         Py_DECREF(p->thread_init_hook);
@@ -177,7 +267,7 @@ int pyda_flush_hooks() {
         if (cb->callback_type == 0) {
             DEBUG_PRINTF("dr_flush_region: %llx\n", (void*)cb->addr);
             dr_flush_region((void*)cb->addr, 1);
-            DEBUG_PRINTF("dr_flush_region end");
+            DEBUG_PRINTF("dr_flush_region end\n");
         }
         cb = cb->next;
     }
@@ -201,6 +291,8 @@ void pyda_hook_cleancall(pyda_hook *cb) {
         return;
     }
 
+    if (t->errored) return;
+
     PyGILState_STATE gstate;
     gstate = PyGILState_Ensure();
 
@@ -211,18 +303,28 @@ void pyda_hook_cleancall(pyda_hook *cb) {
     t->cur_context.pc = (app_pc)cb->addr;
     t->rip_updated_in_cleancall = 0;
 
+    DEBUG_PRINTF("cleancall %p %p %p\n", cb, cb->py_func, t);
+
     PyObject *result = PyObject_CallFunctionObjArgs(cb->py_func, t->proc->py_obj, NULL);
     if (result == NULL) {
+        dr_fprintf(STDERR, "\n[Pyda] ERROR: Hook call failed. Skipping future hooks on thread %d\n", t->tid);
+        dr_flush_file(STDERR);
+        t->errored = 1;
         PyErr_Print();
-        dr_fprintf(STDERR, "[Pyda] ERROR: Hook call failed. Aborting.\n");
-        dr_abort();
+        dr_fprintf(STDERR, "\n");
+        // dr_abort();
+    } else {
+        Py_DECREF(result);
     }
-    Py_DECREF(result);
+
+    DEBUG_PRINTF("cleancall ret %p %p %p\n", cb, cb->py_func, t);
 
     if (t->cur_context.pc == (app_pc)cb->addr && t->rip_updated_in_cleancall) {
         if (t->rip_updated_in_cleancall) {
-            fprintf(stderr, "Hook updated RIP to the same address. This is UB. Aborting.\n");
-            dr_abort();
+            dr_fprintf(STDERR, "\n[Pyda] Hook updated RIP to the same address. This is UB. Skipping future hooks.\n");
+            dr_flush_file(STDERR);
+            t->errored = 1;
+            // dr_abort();
         }
     }
 

pyda_core/pyda_core.h

@@ -37,6 +37,11 @@ struct pyda_process_s {
     pyda_thread *main_thread;
     PyObject *thread_init_hook;
     PyObject *py_obj;
+
+    pthread_cond_t thread_exit_cond;
+    pthread_mutex_t refcount_mutex;
+
+    void* entrypoint;
 };
 
 struct pyda_thread_s {
@@ -47,7 +52,6 @@ struct pyda_thread_s {
     pthread_mutex_t mutex;
 
     int python_yielded, app_yielded;
-    void* start_pc;
 
     pyda_process *proc;
     // PyObject *py_obj;
@@ -57,6 +61,8 @@ struct pyda_thread_s {
     int python_exited;
     int yield_count;
 
+    int errored;
+
 #ifdef PYDA_DYNAMORIO_CLIENT
     dr_mcontext_t cur_context;
 #endif
@@ -67,14 +73,17 @@ pyda_thread* pyda_mk_thread(pyda_process*);
 
 void pyda_process_destroy(pyda_process *p);
 void pyda_thread_destroy(pyda_thread *t);
+void pyda_thread_destroy_last(pyda_thread *t);
 
 PyObject *pyda_run_until(pyda_thread *, uint64_t addr);
 
 // yield from python to the executable
 void pyda_yield(pyda_thread *t);
+void pyda_yield_noblock(pyda_thread *t); // used when thread entry hook returns, we don't need to return to python.
 
 // break from the executable to python
 void pyda_break(pyda_thread *t);
+void pyda_break_noblock(pyda_thread *t); // used when app exits, no need to return to it.
 
 void pyda_initial_break(pyda_thread *t);
 void pyda_add_hook(pyda_process *p, uint64_t addr, PyObject *callback);

pyda_core/pyda_patch_python.c

@@ -16,6 +16,7 @@ static redirect_import_t python_redirect_imports[] = {
     { "pthread_cond_init", (app_pc)pyda_cond_init },
     { "pthread_cond_timedwait", (app_pc)pyda_cond_timedwait },
     { "pthread_cond_signal", (app_pc)pyda_cond_signal },
+    { "pthread_mutex_init", (app_pc)pyda_mutex_init },
     { "dlopen", (app_pc)pyda_dlopen },
     { "dlsym", (app_pc)pyda_dlsym },
 };