python: support threading apis in Pyda scripts

closes #40

pyda_core/pyda_core_py.c

@@ -241,6 +241,14 @@ pyda_core_process(PyObject *self, PyObject *args, PyObject *kwargs) {
     return (PyObject*)result;
 }
 
+static int check_valid_thread(pyda_thread *t) {
+    if (!t) {
+        PyErr_SetString(PyExc_RuntimeError, "Threads created with Python threading APIs cannot use Pyda APIs");
+        return 1;
+    }
+    return 0;
+}
+
 static int check_python_thread(pyda_thread *t) {
     if (pyda_thread_getspecific(g_pyda_tls_is_python_thread_idx) != (void*)1) {
         PyErr_SetString(InvalidStateError, ".run()/.run_until() cannot be called from hooks.");
@@ -250,6 +258,7 @@ static int check_python_thread(pyda_thread *t) {
 }
 
 static int check_exited(pyda_thread *t) {
+    if (check_valid_thread(t)) return 1;
     if (t->app_exited) {
         PyErr_SetString(InvalidStateError, "Thread has already exited; cannot be resumed");
         return 1;
@@ -343,6 +352,7 @@ PydaProcess_capture_io(PyObject* self, PyObject *noarg) {
 static PyObject *
 PydaProcess_backtrace(PyObject* self, PyObject *noarg) {
     pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
+    if (check_exited(t)) return NULL;
 
     char *s = malloc(4096);
     pyda_get_backtrace(t, s, 4096);
@@ -384,6 +394,7 @@ PydaProcess_run_until_pc(PyObject* self, PyObject *args) {
 static PyObject *
 PydaProcess_exited(PyObject* self, PyObject *noarg) {
     pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
+    if (check_valid_thread(t)) return NULL;
     if (t->app_exited) {
         Py_INCREF(Py_True);
         return Py_True;
@@ -412,7 +423,9 @@ pyda_list_modules(PyObject* self, PyObject *noarg) {
 static PyObject *
 pyda_get_current_thread_id(PyObject* self, PyObject *noarg) {
 #ifdef PYDA_DYNAMORIO_CLIENT
-    int tid = ((pyda_thread*)pyda_thread_getspecific(g_pyda_tls_idx))->tid;
+    pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
+    if (check_valid_thread(t)) return NULL;
+    int tid = t->tid;
     return PyLong_FromLong(tid);
 #endif // PYDA_DYNAMORIO_CLIENT
 
@@ -430,6 +443,7 @@ static PyObject *
 PydaProcess_get_register(PyObject *self, PyObject *args) {
     PydaProcess *p = (PydaProcess*)self;
     pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
+    if (check_exited(t)) return NULL;
 
     unsigned long long reg_id;
 
@@ -488,6 +502,7 @@ static PyObject *
 PydaProcess_set_register(PyObject *self, PyObject *args) {
     PydaProcess *p = (PydaProcess*)self;
     pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
+    if (check_exited(t)) return NULL;
 
     unsigned long long reg_id;
     PyObject *val;

pyda_core/pyda_patch_python.c

@@ -18,6 +18,8 @@ static redirect_import_t python_redirect_imports[] = {
     { "pthread_cond_signal", (app_pc)pyda_cond_signal },
     { "pthread_mutex_init", (app_pc)pyda_mutex_init },
     { "pthread_self", (app_pc)pyda_thread_self },
+    { "pthread_create", (app_pc)pyda_thread_create },
+    { "pthread_detach", (app_pc)pyda_thread_detach },
     { "dlopen", (app_pc)pyda_dlopen },
     { "dlsym", (app_pc)pyda_dlsym },
 };

pyda_core/pyda_threads.c

@@ -1,10 +1,12 @@
+// We use this so that we have dr_set_tls_field
+#define STATIC_DRMGR_ONLY
+#include "pyda_threads.h"
 
 #include "dr_api.h"
 #include "dr_tools.h"
 #include "drmgr.h"
 #include "privload.h"
 #include "Python.h"
-#include <pthread.h>
 #include "util.h"
 
 // These are used by python as shims to dynamorio-safe pthread functions
@@ -104,12 +106,61 @@ int pyda_mutex_init(pthread_mutex_t *mutex, const pthread_mutexattr_t *attr) {
     return res;
 }
 
-int pyda_thread_self() {
+void* pyda_thread_self() {
     // XXX: We *could* try to return our pyda-specific tid -- but there
     // are technically two threads with that tid!! (Python and App).
     // If we returned the same ID for two python threads,
     // it seems likely it would break things.
     //
     // Instead, we are just going to return the dynamorio thread id
-    return dr_get_thread_id(dr_get_current_drcontext());
-}
+    return (void*)(uintptr_t)dr_get_thread_id(dr_get_current_drcontext());
+}
+
+extern void __ctype_init();
+void* python_thread_init(void *pyda_thread) {
+    __ctype_init();
+
+    void *drcontext = dr_get_current_drcontext();
+    void *tls = dr_thread_alloc(drcontext, sizeof(void*) * 130);
+    memset(tls, 0, sizeof(void*) * 130);
+    dr_set_tls_field(drcontext, (void *)tls);
+
+    dr_client_thread_set_suspendable(false);
+    pyda_thread_setspecific(g_pyda_tls_idx, (void*)pyda_thread);
+    pyda_thread_setspecific(g_pyda_tls_is_python_thread_idx, (void*)1);
+    return tls;
+}
+
+struct thread_start {
+    void *(*start_routine) (void *);
+    void *arg;
+    // void *pyda_thread;
+};
+
+static void client_thread_init(void *arg) {
+    struct thread_start *ts = (struct thread_start*)arg;
+    void *tls = python_thread_init(NULL);
+    ts->start_routine(ts->arg);
+    DEBUG_PRINTF("start_routine returned\n");
+    dr_client_thread_set_suspendable(true);
+    dr_thread_free(dr_get_current_drcontext(), tls, sizeof(void*) * 130);
+    dr_global_free(ts, sizeof(struct thread_start));
+}
+
+int pyda_thread_create(pthread_t *thread, const pthread_attr_t *attr, void *(*start_routine) (void *), void *arg) {
+    DEBUG_PRINTF("pthread_create %p %p %p %p\n", thread, attr, start_routine, arg);
+
+    struct thread_start *ts = dr_global_alloc(sizeof(struct thread_start));
+    ts->start_routine = start_routine;
+    ts->arg = arg;
+    // ts->pyda_thread = pyda_thread_getspecific(g_pyda_tls_idx);
+    dr_create_client_thread(client_thread_init, ts);
+    *thread = (pthread_t)0x13371337;
+    return 0;
+}
+
+int pyda_thread_detach(pthread_t thread) {
+    // nop
+    DEBUG_PRINTF("pthread_detach %p\n", thread);
+    return 0;
+}

pyda_core/pyda_threads.h

@@ -1,3 +1,4 @@
+#include <pthread.h>
 extern int g_pyda_tls_idx;
 extern int g_pyda_tls_is_python_thread_idx;
 int pyda_thread_setspecific(pthread_key_t key, void *val);
@@ -15,3 +16,7 @@ int pyda_mutex_init(pthread_mutex_t *mutex, const pthread_mutexattr_t *attr);
 void* pyda_dlopen(const char *filename, int flag);
 void* pyda_dlsym(void *handle, const char *symbol);
 void* pyda_thread_self();
+int pyda_thread_create(pthread_t *thread, const pthread_attr_t *attr, void *(*start_routine) (void *), void *arg);
+int pyda_thread_detach(pthread_t thread);
+
+void* python_thread_init(void *pyda_thread);

pyda_core/tool.c

@@ -36,6 +36,7 @@ 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);
+static void event_attach_post(void);
 
 
 extern int is_dynamorio_running;
@@ -81,6 +82,7 @@ dr_client_main(client_id_t id, int argc, const char *argv[])
     drmgr_register_pre_syscall_event(pre_syscall_event);
     drmgr_register_post_syscall_event(post_syscall_event);
     dr_register_filter_syscall_event(filter_syscall_event);
+    dr_register_post_attach_event(event_attach_post);
 
     drmgr_register_signal_event(signal_event);
 
@@ -309,6 +311,13 @@ static dr_signal_action_t signal_event(void *drcontext, dr_siginfo_t *siginfo) {
     return DR_SIGNAL_DELIVER;
 }
 
+static void event_attach_post() {
+    DEBUG_PRINTF("event_attach_post\n");
+    sleep(10);
+    // dr_fprintf(STDERR, "event_attach_post\n");
+    // dr_abort();
+}
+
 static void thread_entrypoint_break() {
     DEBUG_PRINTF("entrypoint (break)\n");
 
@@ -335,20 +344,6 @@ static void thread_entrypoint_break() {
 
 void drmgr_thread_init_event(void*);
 
-static void* python_thread_init(pyda_thread *t) {
-    __ctype_init();
-
-    void *drcontext = dr_get_current_drcontext();
-    void *tls = dr_thread_alloc(drcontext, sizeof(void*) * 130);
-    memset(tls, 0, sizeof(void*) * 130);
-    dr_set_tls_field(drcontext, (void *)tls);
-
-    dr_client_thread_set_suspendable(false);
-    pyda_thread_setspecific(g_pyda_tls_idx, (void*)t);
-    pyda_thread_setspecific(g_pyda_tls_is_python_thread_idx, (void*)1);
-    return tls;
-}
-
 void python_main_thread(void *arg) {
     pyda_thread *t = arg;
     void *drcontext = dr_get_current_drcontext();

tests/run_tests.py

@@ -215,6 +215,15 @@ def no_warnings_or_errors(stdout: bytes, stderr: bytes) -> bool:
             lambda o, e: o.count(b"pass\n") == 1,
         ]
     )),
+
+    ("test_python_threading", "simple.c", "test_python_threading.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_python_threading.py

@@ -0,0 +1,22 @@
+from pyda import *
+from pwnlib.elf.elf import ELF
+from pwnlib.util.packing import u64
+import time
+from threading import Thread
+
+p = process()
+
+e = ELF(p.exe_path)
+e.address = p.maps[p.exe_path].base
+
+def thread():
+    print("thread start")
+    time.sleep(1)
+    print("thread end")
+
+t = Thread(target=thread)
+t.start()
+
+time.sleep(3)
+p.run()
+print("pass")