Basic ARM64 support (#7)

ndrewh · Aug 20, 2024 · 0be7bf5 · 0be7bf5
1 parent fb8fcda
commit 0be7bf5
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 10 deletions.
diff --git a/examples/ltrace.py b/examples/ltrace.py
@@ -26,12 +26,12 @@ def guess_arg(x):
     return hex(x)
 
 def lib_hook(p):
-    name = plt_map[p.regs.rip]
+    name = plt_map[p.regs.pc]
     print(f"{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)}",
+        f"rdi={guess_arg(p.regs.arg0)}",
+        f"rsi={guess_arg(p.regs.arg1)}",
+        f"rdx={guess_arg(p.regs.arg2)}",
+        f"rcx={guess_arg(p.regs.arg3)}",
     ]) + ")")
 
 for x in e.plt:

diff --git a/pyda_core/pyda_core_py.c b/pyda_core/pyda_core_py.c
@@ -87,7 +87,7 @@ PyInit_pyda_core(void) {
     register_exception(m, &InvalidStateError, "pyda.InvalidStateError", "InvalidStateError");
     register_exception(m, &FatalSignalError, "pyda.FatalSignalError", "FatalSignalError");
 
-#ifdef X86
+#if defined(X86)
     PyModule_AddIntConstant(m, "REG_RAX", DR_REG_RAX);
     PyModule_AddIntConstant(m, "REG_RBX", DR_REG_RBX);
     PyModule_AddIntConstant(m, "REG_RCX", DR_REG_RCX);
@@ -115,7 +115,13 @@ PyInit_pyda_core(void) {
     PyModule_AddIntConstant(m, "REG_XMM5", DR_REG_XMM5);
     PyModule_AddIntConstant(m, "REG_XMM6", DR_REG_XMM6);
     PyModule_AddIntConstant(m, "REG_XMM7", DR_REG_XMM7);
-#elif ARM64
+    PyModule_AddIntConstant(m, "REG_ARG0", DR_REG_RDI);
+    PyModule_AddIntConstant(m, "REG_ARG1", DR_REG_RSI);
+    PyModule_AddIntConstant(m, "REG_ARG2", DR_REG_RDX);
+    PyModule_AddIntConstant(m, "REG_ARG3", DR_REG_RCX);
+    PyModule_AddIntConstant(m, "REG_ARG4", DR_REG_R8);
+    PyModule_AddIntConstant(m, "REG_ARG5", DR_REG_R9);
+#elif defined(AARCH64)
     PyModule_AddIntConstant(m, "REG_X0", DR_REG_X0);
     PyModule_AddIntConstant(m, "REG_X1", DR_REG_X1);
     PyModule_AddIntConstant(m, "REG_X2", DR_REG_X2);
@@ -149,6 +155,12 @@ PyInit_pyda_core(void) {
     PyModule_AddIntConstant(m, "REG_X30", DR_REG_X30);
     PyModule_AddIntConstant(m, "REG_SP", DR_REG_SP);
     PyModule_AddIntConstant(m, "REG_PC", PYDA_REG_PC);
+    PyModule_AddIntConstant(m, "REG_ARG0", DR_REG_R0);
+    PyModule_AddIntConstant(m, "REG_ARG1", DR_REG_R1);
+    PyModule_AddIntConstant(m, "REG_ARG2", DR_REG_R2);
+    PyModule_AddIntConstant(m, "REG_ARG3", DR_REG_R3);
+    PyModule_AddIntConstant(m, "REG_ARG4", DR_REG_R4);
+    PyModule_AddIntConstant(m, "REG_ARG5", DR_REG_R5);
 #endif
 
     return m;

diff --git a/pyda_core/pyda_unwind.c b/pyda_core/pyda_unwind.c
@@ -9,6 +9,8 @@ int pyda_get_backtrace (pyda_thread *t, char *buf, int size) {
 
   unw_getcontext(&uc);
   ucontext_t *uc2 = (ucontext_t *) &uc;
+
+#if defined(__x86_64__)
   uc2->uc_mcontext.gregs[REG_RIP] = (uintptr_t)t->cur_context.pc;
   uc2->uc_mcontext.gregs[REG_RSP] = t->cur_context.rsp;
   uc2->uc_mcontext.gregs[REG_RBP] = t->cur_context.rbp;
@@ -18,6 +20,21 @@ int pyda_get_backtrace (pyda_thread *t, char *buf, int size) {
   uc2->uc_mcontext.gregs[REG_RCX] = t->cur_context.rcx;
   uc2->uc_mcontext.gregs[REG_R8] = t->cur_context.r8;
   uc2->uc_mcontext.gregs[REG_R9] = t->cur_context.r8;
+#elif defined(AARCH64)
+  uc2->uc_mcontext.pc = (uintptr_t)t->cur_context.pc;
+  uc2->uc_mcontext.regs[0] = t->cur_context.r0;
+  uc2->uc_mcontext.regs[1] = t->cur_context.r1;
+  uc2->uc_mcontext.regs[2] = t->cur_context.r2;
+  uc2->uc_mcontext.regs[3] = t->cur_context.r3;
+  uc2->uc_mcontext.regs[4] = t->cur_context.r4;
+  uc2->uc_mcontext.regs[5] = t->cur_context.r5;
+  uc2->uc_mcontext.regs[6] = t->cur_context.r6;
+  uc2->uc_mcontext.regs[7] = t->cur_context.r7;
+  uc2->uc_mcontext.sp = t->cur_context.sp;
+  uc2->uc_mcontext.lr = t->cur_context.sp;
+#else
+#error "Unsupported architecture"
+#endif
 
   unw_init_local(&cursor, &uc);
 

diff --git a/pyda_core/tool.c b/pyda_core/tool.c
@@ -220,18 +220,26 @@ event_insert(void *drcontext, void *tag, instrlist_t *bb, instr_t *instr,
 
     // XXX: I don't think this is safe, since the thread that updates
     // the code cache may not be the executing thread.
+    //
+#if defined(X86)
+    bool save_fpstate = true;
+#elif defined(AARCH64)
+    bool save_fpstate = false;
+#else
+    #error "Unsupported arch"
+#endif
     if (instr_get_app_pc(instr) == t->proc->entrypoint) {
         DEBUG_PRINTF("** Found PC\n");
         dr_insert_clean_call(drcontext, bb, instrlist_first_app(bb), (void *)thread_entrypoint_break,
                          false /* save fpstate */, 0);
     } else if ((callback = pyda_get_callback(t->proc, instr_get_app_pc(instr)))) {
         DEBUG_PRINTF("installing hook at %p\n", instr_get_app_pc(instr));
         dr_insert_clean_call(drcontext, bb, instr, (void *)pyda_hook_cleancall,
-                         true /* save fpstate */, 1, OPND_CREATE_INTPTR(callback));
+                         save_fpstate /* save fpstate */, 1, OPND_CREATE_INTPTR(callback));
     } else if (pyda_check_run_until(t->proc, instr_get_app_pc(instr))) {
         DEBUG_PRINTF("installing run_until hook at %p\n", instr_get_app_pc(instr));
         dr_insert_clean_call(drcontext, bb, instr, (void *)pyda_hook_rununtil_reached,
-                         true /* save fpstate */, 1, OPND_CREATE_INTPTR(instr_get_app_pc(instr)));
+                         save_fpstate /* save fpstate */, 1, OPND_CREATE_INTPTR(instr_get_app_pc(instr)));
     }
     return DR_EMIT_DEFAULT;
 }
@@ -450,4 +458,4 @@ void python_aux_thread(void *arg) {
     DEBUG_PRINTF("python_aux_thread 5\n");
     dr_thread_free(drcontext, tls, sizeof(void*) * 130);
     DEBUG_PRINTF("python_aux_thread 6\n");
-}
+}