Fix x86_64 call replay trampolines

Author: YinMo19

src/arch/x86_64/trampoline.zig

@@ -42,13 +42,19 @@ const TrampolineEmitter = struct {
     }
 
     fn emitAbsolutePush(self: *TrampolineEmitter, value: u64) HookError!void {
-        // `push qword ptr [rip+0]` followed by an inline literal. This is used
-        // to synthesize the architectural "return address push" side effect of
-        // a displaced `call`.
-        const opcode = [_]u8{ 0xFF, 0x35, 0x00, 0x00, 0x00, 0x00 };
+        // x86_64 has no `push imm64`. We therefore synthesize one without
+        // permanently clobbering architectural state:
+        //   push rax
+        //   movabs rax, imm64
+        //   xchg qword ptr [rsp], rax
+        //
+        // After the final `xchg`, the stack top contains `value` and `rax`
+        // regains its original pre-push contents.
         const literal = std.mem.toBytes(std.mem.nativeToLittle(u64, value));
-        try self.emit(opcode[0..]);
+        try self.emit(&.{0x50});
+        try self.emit(&.{ 0x48, 0xB8 });
         try self.emit(literal[0..]);
+        try self.emit(&.{ 0x48, 0x87, 0x04, 0x24 });
     }
 
     fn emitStackPointerIndirectJump(

tests/api_integration.zig

@@ -5,6 +5,8 @@ const zighook = @import("zighook");
 
 extern fn demo_add_target(a: i32, b: i32) callconv(.c) i32;
 extern fn demo_add_patchpoint() callconv(.c) void;
+extern fn demo_direct_call_target(a: i32, b: i32) callconv(.c) i32;
+extern fn demo_direct_call_patchpoint() callconv(.c) void;
 extern fn demo_stack_call_target(a: i32, b: i32) callconv(.c) i32;
 extern fn demo_stack_call_patchpoint() callconv(.c) void;
 extern fn demo_prepatched_target() callconv(.c) i32;
@@ -67,6 +69,19 @@ test "x86_64 instrument replays stack-pointer indirect calls" {
     try std.testing.expect(saved.slice().len > 0);
 }
 
+test "x86_64 instrument replays direct calls" {
+    if (builtin.cpu.arch != .x86_64) return;
+
+    const patchpoint_addr: u64 = @intFromPtr(&demo_direct_call_patchpoint);
+
+    defer zighook.unhook(patchpoint_addr) catch {};
+
+    _ = try zighook.instrument(patchpoint_addr, signalReplay42);
+    try std.testing.expectEqual(@as(i32, 42), demo_direct_call_target(1, 2));
+    const saved = zighook.original_instruction(patchpoint_addr) orelse return error.TestUnexpectedResult;
+    try std.testing.expect(saved.slice().len > 0);
+}
+
 test "instrument_no_original skips the displaced instruction on both backends" {
     const patchpoint_addr: u64 = @intFromPtr(&demo_add_patchpoint);
 

tests/support/runtime_targets_x86_64.S

@@ -15,6 +15,21 @@ SYM(demo_add_patchpoint):
     leal (%rdi,%rsi), %eax
     ret
 
+// Replay test for ordinary direct calls. The patchpoint is the `call rel32`
+// itself so the execute-original trampoline must synthesize the pushed return
+// address before transferring control to the final callee.
+.globl SYM(demo_direct_call_target)
+.globl SYM(demo_direct_call_patchpoint)
+SYM(demo_direct_call_target):
+SYM(demo_direct_call_patchpoint):
+    call SYM(demo_direct_call_impl)
+    ret
+
+.globl SYM(demo_direct_call_impl)
+SYM(demo_direct_call_impl):
+    leal (%rdi,%rsi), %eax
+    ret
+
 // Replay test for stack-pointer-based indirect calls. The patchpoint is the
 // indirect call through the top-of-stack function pointer slot.
 .globl SYM(demo_stack_call_target)