From 964391109bce5e6d95d848f2daed72985c330e79 Mon Sep 17 00:00:00 2001
From: Andrew Haberlandt <ahaberla@andrew.cmu.edu>
Date: Thu, 14 Nov 2024 21:05:54 +0000
Subject: [PATCH] feat: basic support for fork() + tests

---
 pyda_core/pyda_core.c |  4 ++--
 pyda_core/tool.c      | 29 +++++++++++++++++++++++++++++
 tests/run_tests.py    | 10 +++++++++-
 tests/test_fork.c     | 30 ++++++++++++++++++++++++++++++
 4 files changed, 70 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_fork.c

diff --git a/pyda_core/pyda_core.c b/pyda_core/pyda_core.c
index 6680af4..359ce87 100644
--- a/pyda_core/pyda_core.c
+++ b/pyda_core/pyda_core.c
@@ -573,7 +573,7 @@ void pyda_hook_cleancall(pyda_hook *cb) {
 
     // If this also happens to be the run_until target for this thread,
     // we deal with that here (instead of inserting two hooks)
-    if (cb->addr == pyda_get_run_until(t)) {
+    if (!t->python_exited && cb->addr == pyda_get_run_until(t)) {
         // It is UB to modify PC in a hook that is also the run_until target
         if (t->rip_updated_in_python) {
             dr_fprintf(STDERR, "\n[Pyda] ERROR: Hook updated RIP, but run_until target is hit. This is UB. Continuing.");
@@ -642,7 +642,7 @@ int pyda_hook_syscall(int syscall_num, int is_pre) {
 
 void pyda_hook_rununtil_reached(void *pc) {
     pyda_thread *t = pyda_thread_getspecific(g_pyda_tls_idx);
-    if (t->errored) return;
+    if (t->errored || t->python_exited) return;
     if (t->skip_next_hook) {
         t->skip_next_hook = 0;
         return;
diff --git a/pyda_core/tool.c b/pyda_core/tool.c
index 6dfebe3..11d3ac3 100644
--- a/pyda_core/tool.c
+++ b/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 fork_event(void *drcontext);
 static void event_attach_post(void);
 
 
@@ -85,6 +86,7 @@ dr_client_main(client_id_t id, int argc, const char *argv[])
     dr_register_post_attach_event(event_attach_post);
 
     drmgr_register_signal_event(signal_event);
+    dr_register_fork_init_event(fork_event);
     dr_request_synchronized_exit();
 
     pthread_cond_init(&python_thread_init1, 0);
@@ -342,6 +344,33 @@ static dr_signal_action_t signal_event(void *drcontext, dr_siginfo_t *siginfo) {
 
     return DR_SIGNAL_DELIVER;
 }
+static void fork_event(void *drcontext) {
+    // This is called on the NEW fork, which doesn't have any parallel Python threads anymore.
+    // TODO: How do we make sure that important locks aren't held at fork time when we have multiple threads?
+    pyda_thread *t = drmgr_get_tls_field(drcontext, g_pyda_tls_idx);
+    DEBUG_PRINTF("[Pyda] fork_init\n");
+
+    pyda_process *p = t->proc;
+
+    // Flush deleted hooks
+    drvector_lock(&p->threads);
+    for (int i=0; i<p->threads.entries; i++) {
+        dr_flush_file(STDERR);
+        if (p->threads.array[i] != t) pyda_thread_destroy(p->threads.array[i]);
+    }
+    p->threads.entries = 1;
+    p->threads.array[0] = t;
+    drvector_unlock(&p->threads);
+
+    // For now, we just mark the current thread as exited, which just means we won't
+    // try to yield to it in signal handlers or if we reach the last run_until hook.
+    //
+    // In the future, we could setup a new parallel Python thread that enters some
+    // "fork handler" as the entrypoint or whatever
+
+    t->python_exited = 1;
+    p->main_thread = t;
+}
 
 static void event_attach_post() {
     DEBUG_PRINTF("event_attach_post on tid %d\n", dr_get_thread_id(dr_get_current_drcontext()));
diff --git a/tests/run_tests.py b/tests/run_tests.py
index a15863a..a71524d 100644
--- a/tests/run_tests.py
+++ b/tests/run_tests.py
@@ -241,6 +241,14 @@ def no_warnings_or_errors(stdout: bytes, stderr: bytes) -> bool:
             lambda o, e: o.count(b"pass\n") == 1,
         ]
     )),
+
+    ("test_fork", "test_fork.c", "../examples/ltrace_multithreaded.py", RunOpts(), ExpectedResult(
+        retcode=0,
+        checkers=[
+            output_checker,
+            no_warnings_or_errors,
+        ]
+    )),
 ]
 
 def main():
@@ -336,4 +344,4 @@ def run_test(c_file, python_file, run_opts, expected_result, test_name, debug, n
 
 
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    main()
diff --git a/tests/test_fork.c b/tests/test_fork.c
new file mode 100644
index 0000000..11e0f8f
--- /dev/null
+++ b/tests/test_fork.c
@@ -0,0 +1,30 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <pthread.h>
+
+void child_main() {
+    printf("child\n");
+}
+
+void parent_main() {
+    printf("parent\n");
+}
+
+int main(int argc, char** argv) {
+    printf("start\n");
+    int f = fork();
+    if (f == 0) {
+        // child
+        child_main();
+    } else if (f > 0) {
+        // parent
+        parent_main();
+
+        int status;
+        waitpid(f, &status, 0);
+        printf("child status %d\n", status);
+    } else {
+        printf("error\n");
+    }
+    printf("end\n");
+}