i#3544 RV64: Implement return address handling for post wrappers (#6736)

Substitute value of return address register to enable post wrappers. Partially enables drwrap-test for RISC-V: 1. Add drwrap-test-callconv test to pipeline 2. Partially supported drwrap-drreg-test: check for pre/post wrappers passes, but checks for clean_call and restore registers functionality are failing. Issue: #3544
DynamoRIO · Aug 16, 2024 · 30c3a9b · 30c3a9b
1 parent 3ae837e
commit 30c3a9b
Show file tree

Hide file tree

Showing 7 changed files with 65 additions and 17 deletions.
diff --git a/core/arch/asm_defines.asm b/core/arch/asm_defines.asm
@@ -418,6 +418,9 @@ ASSUME fs:_DATA @N@\
 # define REG_R29  x29
 # define REG_R30  x30
 # define REG_R31  x31
+# define REG_A0   x10
+# define REG_A1   x11
+# define REG_A2   x12
 #else /* Intel X86 */
 # ifdef X64
 #  define REG_XAX rax

diff --git a/ext/drwrap/drwrap.c b/ext/drwrap/drwrap.c
@@ -1940,12 +1940,9 @@ drwrap_ensure_postcall(void *drcontext, per_thread_t *pt, wrap_entry_t *wrap,
         pt->retaddr[pt->wrap_level] = wrapcxt->retaddr; /* Original, not load tgt. */
 #ifdef X86
         set_retaddr_on_stack(wrapcxt->mc->xsp, (app_pc)replace_retaddr_sentinel);
-#elif defined(RISCV64)
-        /* FIXME i#3544: Not implemented */
-        ASSERT(false, "Not implemented");
 #else
         drwrap_get_mcontext_internal(wrapcxt, DR_MC_CONTROL);
-        wrapcxt->mc->lr = (reg_t)replace_retaddr_sentinel;
+        wrapcxt->mc->IF_RISCV64_ELSE(ra, lr) = (reg_t)replace_retaddr_sentinel;
         wrapcxt->mc_modified = true;
 #endif
         return;
@@ -1986,7 +1983,7 @@ drwrap_ensure_postcall(void *drcontext, per_thread_t *pt, wrap_entry_t *wrap,
 
 /* called via clean call at the top of callee */
 static void
-drwrap_in_callee(void *arg1, reg_t xsp _IF_NOT_X86(reg_t lr))
+drwrap_in_callee(void *arg1, reg_t xsp _IF_NOT_X86(IF_RISCV64_ELSE(reg_t ra, reg_t lr)))
 {
     void *drcontext = dr_get_current_drcontext();
     per_thread_t *pt = (per_thread_t *)drmgr_get_tls_field(drcontext, tls_idx);
@@ -2006,6 +2003,8 @@ drwrap_in_callee(void *arg1, reg_t xsp _IF_NOT_X86(reg_t lr))
 #ifdef AARCHXX
     /* ditto */
     mc.lr = lr;
+#elif defined(RISCV64)
+    mc.ra = ra;
 #endif
     mc.flags = 0; /* if anything else is asked for, lazily initialize */
 
@@ -2020,7 +2019,8 @@ drwrap_in_callee(void *arg1, reg_t xsp _IF_NOT_X86(reg_t lr))
 
     NOTIFY(2, "%s: level %d function " PFX "\n", __FUNCTION__, pt->wrap_level + 1, pc);
 
-    app_pc retaddr = IF_X86_ELSE(get_retaddr_from_stack(xsp), (app_pc)lr);
+    app_pc retaddr =
+        IF_X86_ELSE(get_retaddr_from_stack(xsp), (app_pc)IF_AARCHXX_ELSE(lr, ra));
     if (TEST(DRWRAP_REPLACE_RETADDR, global_flags)) {
         /* In case of a tailcall, the return address has already been replaced by
          * the sentinel in the stack, we need to retrieve the return address from the
@@ -2467,12 +2467,12 @@ drwrap_event_bb_insert_where(void *drcontext, void *tag, instrlist_t *bb, instr_
                 ? (DR_CLEANCALL_NOSAVE_FLAGS | DR_CLEANCALL_NOSAVE_XMM_NONPARAM)
                 : 0;
             flags |= DR_CLEANCALL_READS_APP_CONTEXT | DR_CLEANCALL_WRITES_APP_CONTEXT;
-            /* FIXME i#3544: Adapt to a real RISC-V clean call implementation */
             dr_insert_clean_call_ex(
                 drcontext, bb, where, (void *)drwrap_in_callee, flags,
-                IF_AARCHXX_ELSE(3, 2), OPND_CREATE_INTPTR((ptr_int_t)arg1),
+                IF_AARCHXX_OR_RISCV64_ELSE(3, 2), OPND_CREATE_INTPTR((ptr_int_t)arg1),
                 /* pass in xsp to avoid dr_get_mcontext */
-                opnd_create_reg(DR_REG_XSP) _IF_AARCHXX(opnd_create_reg(DR_REG_LR)));
+                opnd_create_reg(DR_REG_XSP) _IF_AARCHXX_OR_RISCV64(
+                    opnd_create_reg(IF_AARCHXX_ELSE(DR_REG_LR, DR_REG_RA))));
         }
         dr_recurlock_unlock(wrap_lock);
     }

diff --git a/ext/drwrap/drwrap.h b/ext/drwrap/drwrap.h
@@ -421,8 +421,9 @@ typedef enum {
      * violation may be exposed to the client's dr_register_kernel_xfer_event() callback
      * if it inspects the mcontext PC on the stack; drwrap_get_retaddr_if_sentinel()
      * may be used to mitigate such cases. Use #DRWRAP_REPLACE_RETADDR at your own risk.
+     * Currently is not supported for RISC-V.
      */
-    DRWRAP_REPLACE_RETADDR = 0x04,
+    DRWRAP_REPLACE_RETADDR = 0x04, /* TODO i#3544: support flag for RISC-V 64 */
 } drwrap_wrap_flags_t;
 
 /* offset of drwrap_callconv_t in drwrap_wrap_flags_t */
@@ -459,6 +460,9 @@ typedef enum {
 #    ifdef AARCH64
     /** Default calling convention for the platform. */
     DRWRAP_CALLCONV_DEFAULT = DRWRAP_CALLCONV_AARCH64,
+#    elif defined(RISCV64)
+    /** Default calling convention for the platform. */
+    DRWRAP_CALLCONV_DEFAULT = DRWRAP_CALLCONV_RISCV_LP64,
 #    elif defined(UNIX) /* x64 */
     /** Default calling convention for the platform. */
     DRWRAP_CALLCONV_DEFAULT = DRWRAP_CALLCONV_AMD64,

diff --git a/suite/tests/CMakeLists.txt b/suite/tests/CMakeLists.txt
@@ -6307,6 +6307,7 @@ if (RISCV64)
     code_api|client.app_args
     code_api|client.blackbox
     code_api|client.crashmsg
+    code_api|client.drwrap-test-callconv
     code_api|client.execfault
     code_api|client.mangle_suspend
     code_api|client.null_instrument

diff --git a/suite/tests/client-interface/drwrap-drreg-test.appdll.c b/suite/tests/client-interface/drwrap-drreg-test.appdll.c
@@ -117,6 +117,16 @@ GLOBAL_LABEL(FUNCNAME:)
         nop
         add      r0, r1, r2
         bx       lr
+# elif defined(RISCV64)
+        mov       REG_A1, HEX(4)
+        /* The clean call writes to x7, replacing this value. */
+        mov       REG_A2, HEX(42)
+        /* The clean call is inserted after 3 nops. */
+        nop
+        nop
+        nop
+        add       REG_A0, REG_A1, REG_A2
+        ret
 # endif
         END_FUNC(FUNCNAME)
 #  undef FUNCNAME
@@ -166,6 +176,18 @@ GLOBAL_LABEL(FUNCNAME:)
         nop
         add      r0, r1, r2
         bx       lr
+# elif defined(RISCV64)
+        mov      REG_A0, ARG1 /* Used to skip clean call. */
+        mov      REG_A1, HEX(4)   /* Read in clean call. */
+        mov      REG_A2, HEX(42)
+        /* Aflags has special x86 behavior; we do not test it here. */
+        /* The clean call is inserted after 4 nops. */
+        nop
+        nop
+        nop
+        nop
+        add      REG_A0, REG_A1, REG_A2
+        ret
 # endif
         END_FUNC(FUNCNAME)
 #  undef FUNCNAME

diff --git a/suite/tests/client-interface/drwrap-drreg-test.dll.c b/suite/tests/client-interface/drwrap-drreg-test.dll.c
@@ -104,8 +104,9 @@ clean_call_rw(void)
     mc.flags = DR_MC_CONTROL | DR_MC_INTEGER;
     bool ok = dr_get_mcontext(drcontext, &mc);
     CHECK(ok, "dr_get_mcontext failed");
-    CHECK(IF_X86_ELSE(mc.xdx, mc.r1) == 4, "app reg val not restored for clean call");
-    IF_X86_ELSE(mc.xcx, mc.r2) = 3;
+    CHECK(IF_X86_ELSE(mc.xdx, IF_AARCHXX_ELSE(mc.r1, mc.a1)) == 4,
+          "app reg val not restored for clean call");
+    IF_X86_ELSE(mc.xcx, IF_AARCHXX_ELSE(mc.r2, mc.a2)) = 3;
     ok = dr_set_mcontext(drcontext, &mc);
     CHECK(ok, "dr_set_mcontext failed");
 }
@@ -121,9 +122,9 @@ clean_call_check_rw(reg_t reg1, reg_t reg2)
     mc.flags = DR_MC_CONTROL | DR_MC_INTEGER;
     bool ok = dr_get_mcontext(drcontext, &mc);
     CHECK(ok, "dr_get_mcontext failed");
-    CHECK(IF_X86_ELSE(mc.xdx, mc.r1) == SENTINEL,
+    CHECK(IF_X86_ELSE(mc.xdx, IF_AARCHXX_ELSE(mc.r1, mc.a1)) == SENTINEL,
           "tool val1 in mc not restored after call");
-    CHECK(IF_X86_ELSE(mc.xdi, mc.r4) == SENTINEL,
+    CHECK(IF_X86_ELSE(mc.xdi, IF_AARCHXX_ELSE(mc.r4, mc.a4)) == SENTINEL,
           "tool val2 in mc not restored after call");
 }
 
@@ -136,7 +137,8 @@ clean_call_multipath(void)
     mc.flags = DR_MC_CONTROL | DR_MC_INTEGER;
     bool ok = dr_get_mcontext(drcontext, &mc);
     CHECK(ok, "dr_get_mcontext failed");
-    CHECK(IF_X86_ELSE(mc.xdx, mc.r1) == 4, "app reg val not restored for clean call");
+    CHECK(IF_X86_ELSE(mc.xdx, IF_AARCHXX_ELSE(mc.r1, mc.a1)) == 4,
+          "app reg val not restored for clean call");
 #ifdef X86
     /* This tests the drreg_statelessly_restore_app_value() respill which only
      * happens with aflags in xax.
@@ -169,6 +171,10 @@ clobber_key_regs(void *drcontext, instrlist_t *bb, instr_t *inst)
     drreg_set_vector_entry(&allowed, DR_REG_XAX, true);
     drreg_set_vector_entry(&allowed, DR_REG_XDI, true);
     drreg_set_vector_entry(&allowed, DR_REG_XSI, true);
+#elif defined(RISCV64)
+    drreg_set_vector_entry(&allowed, DR_REG_A0, true);
+    drreg_set_vector_entry(&allowed, DR_REG_A1, true);
+    drreg_set_vector_entry(&allowed, DR_REG_A2, true);
 #else
     drreg_set_vector_entry(&allowed, DR_REG_R0, true);
     drreg_set_vector_entry(&allowed, DR_REG_R1, true);
@@ -218,6 +224,9 @@ insert_rw_call(void *drcontext, instrlist_t *bb, instr_t *inst)
 #ifdef X86
     drreg_set_vector_entry(&allowed, DR_REG_XDX, true);
     drreg_set_vector_entry(&allowed, DR_REG_XDI, true);
+#elif defined(RISCV64)
+    drreg_set_vector_entry(&allowed, DR_REG_A1, true);
+    drreg_set_vector_entry(&allowed, DR_REG_A4, true);
 #else
     drreg_set_vector_entry(&allowed, DR_REG_R1, true);
     drreg_set_vector_entry(&allowed, DR_REG_R4, true);
@@ -267,6 +276,8 @@ insert_multipath_call(void *drcontext, instrlist_t *bb, instr_t *inst)
     /* Clobber the reg we check in clean_call_multipath(). */
 #ifdef X86
     drreg_set_vector_entry(&allowed, DR_REG_XDX, true);
+#elif defined(RISCV64)
+    drreg_set_vector_entry(&allowed, DR_REG_A1, true);
 #else
     drreg_set_vector_entry(&allowed, DR_REG_R1, true);
 #endif
@@ -283,13 +294,20 @@ insert_multipath_call(void *drcontext, instrlist_t *bb, instr_t *inst)
 
     instr_t *skip_call = INSTR_CREATE_label(drcontext);
     /* The app executes twice and sets rcx/r0 to 0 for one of them. */
+#if defined(RISCV64)
+    instrlist_meta_preinsert(bb, inst,
+                             INSTR_CREATE_beq(drcontext, opnd_create_instr(skip_call),
+                                              opnd_create_reg(DR_REG_A0),
+                                              OPND_CREATE_INT32(0)));
+#else
     instrlist_meta_preinsert(
         bb, inst,
         XINST_CREATE_cmp(drcontext, opnd_create_reg(IF_X86_ELSE(DR_REG_XCX, DR_REG_R0)),
                          OPND_CREATE_INT32(0)));
     instrlist_meta_preinsert(
         bb, inst,
         XINST_CREATE_jump_cond(drcontext, DR_PRED_EQ, opnd_create_instr(skip_call)));
+#endif
     dr_insert_clean_call_ex(drcontext, bb, inst, clean_call_multipath,
                             DR_CLEANCALL_READS_APP_CONTEXT | DR_CLEANCALL_MULTIPATH, 0);
     instrlist_meta_preinsert(bb, inst, skip_call);
@@ -331,7 +349,7 @@ event_app_instruction(void *drcontext, void *tag, instrlist_t *bb, instr_t *inst
     /* Look for nop;nop;nop in reg_val_test() and 4 nops in multipath_test(). */
     if (instr_is_app(inst)) {
         int *nop_count = (int *)user_data;
-        if (instr_get_opcode(inst) == OP_nop IF_ARM(|| instr_is_mov_nop(inst))) {
+        if (instr_is_nop(inst) IF_ARM(|| instr_is_mov_nop(inst))) {
             ++(*nop_count);
         } else {
             if (*nop_count == 3) {

diff --git a/suite/tests/client-interface/drwrap-test-callconv.dll.cpp b/suite/tests/client-interface/drwrap-test-callconv.dll.cpp
@@ -48,7 +48,7 @@
 #    define PLATFORM_HAS_THISCALL 1
 #endif
 
-#if !defined(ARM) && !defined(X64)
+#if !defined(ARM) && !defined(X64) && !defined(RISCV64)
 #    define PLATFORM_HAS_FASTCALL 1
 #endif