From 39dff2cdfeb62869b032e23e0dac3820f1365db3 Mon Sep 17 00:00:00 2001
From: cursey <cursey@users.noreply.github.com>
Date: Wed, 4 Oct 2023 17:25:43 -0700
Subject: [PATCH] Feat/MidHook Improvements (#39)

* MidHook: Add IP to context

* MidHook: Capture XMM registers

* Context: Update doc comment
---
 include/safetyhook/context.hpp | 19 +++++++++++---
 src/mid_hook.cpp               | 47 ++++++++++++++++++++++++---------
 src/mid_hook.x86.asm           | 28 +++++++++++++++++---
 src/mid_hook.x86_64.asm        | 48 +++++++++++++++++++++++++++++-----
 unittest/mid_hook.cpp          | 30 ++++++++++++++++++++-
 5 files changed, 144 insertions(+), 28 deletions(-)

diff --git a/include/safetyhook/context.hpp b/include/safetyhook/context.hpp
index 231aefb..c763f08 100644
--- a/include/safetyhook/context.hpp
+++ b/include/safetyhook/context.hpp
@@ -6,20 +6,31 @@
 #include <cstdint>
 
 namespace safetyhook {
+union Xmm {
+    uint8_t u8[16];
+    uint16_t u16[8];
+    uint32_t u32[4];
+    uint64_t u64[2];
+    float f32[4];
+    double f64[2];
+};
+
 /// @brief Context structure for 64-bit MidHook.
 /// @details This structure is used to pass the context of the hooked function to the destination allowing full access
 /// to the 64-bit registers at the moment the hook is called.
-/// @note The structure only provides access to integer registers.
+/// @note rip will point to a trampoline containing the replaced instruction(s).
 struct Context64 {
-    uintptr_t rflags, r15, r14, r13, r12, r11, r10, r9, r8, rdi, rsi, rdx, rcx, rbx, rax, rbp, rsp;
+    Xmm xmm0, xmm1, xmm2, xmm3, xmm4, xmm5, xmm6, xmm7, xmm8, xmm9, xmm10, xmm11, xmm12, xmm13, xmm14, xmm15;
+    uintptr_t rflags, r15, r14, r13, r12, r11, r10, r9, r8, rdi, rsi, rdx, rcx, rbx, rax, rbp, rsp, rip;
 };
 
 /// @brief Context structure for 32-bit MidHook.
 /// @details This structure is used to pass the context of the hooked function to the destination allowing full access
 /// to the 32-bit registers at the moment the hook is called.
-/// @note The structure only provides access to integer registers.
+/// @note eip will point to a trampoline containing the replaced instruction(s).
 struct Context32 {
-    uintptr_t eflags, edi, esi, edx, ecx, ebx, eax, ebp, esp;
+    Xmm xmm0, xmm1, xmm2, xmm3, xmm4, xmm5, xmm6, xmm7;
+    uintptr_t eflags, edi, esi, edx, ecx, ebx, eax, ebp, esp, eip;
 };
 
 /// @brief Context structure for MidHook.
diff --git a/src/mid_hook.cpp b/src/mid_hook.cpp
index f0d974c..dd879c3 100644
--- a/src/mid_hook.cpp
+++ b/src/mid_hook.cpp
@@ -1,4 +1,5 @@
 #include <algorithm>
+#include <array>
 
 #include <safetyhook/allocator.hpp>
 #include <safetyhook/inline_hook.hpp>
@@ -9,16 +10,36 @@
 namespace safetyhook {
 
 #ifdef _M_X64
-constexpr uint8_t asm_data[] = {0x54, 0x55, 0x50, 0x53, 0x51, 0x52, 0x56, 0x57, 0x41, 0x50, 0x41, 0x51, 0x41, 0x52,
-    0x41, 0x53, 0x41, 0x54, 0x41, 0x55, 0x41, 0x56, 0x41, 0x57, 0x9C, 0x48, 0x8D, 0x0C, 0x24, 0x48, 0x89, 0xE3, 0x48,
-    0x83, 0xEC, 0x30, 0x48, 0x83, 0xE4, 0xF0, 0xFF, 0x15, 0x22, 0x00, 0x00, 0x00, 0x48, 0x89, 0xDC, 0x9D, 0x41, 0x5F,
-    0x41, 0x5E, 0x41, 0x5D, 0x41, 0x5C, 0x41, 0x5B, 0x41, 0x5A, 0x41, 0x59, 0x41, 0x58, 0x5F, 0x5E, 0x5A, 0x59, 0x5B,
-    0x58, 0x5D, 0x5C, 0xFF, 0x25, 0x08, 0x00, 0x00, 0x00, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90,
-    0x90, 0x90, 0x90, 0x90, 0x90, 0x90};
+constexpr auto asm_data = std::to_array<uint8_t>({0xff, 0x35, 0x5c, 0x01, 0x00, 0x00, 0x54, 0x55, 0x50, 0x53, 0x51,
+    0x52, 0x56, 0x57, 0x41, 0x50, 0x41, 0x51, 0x41, 0x52, 0x41, 0x53, 0x41, 0x54, 0x41, 0x55, 0x41, 0x56, 0x41, 0x57,
+    0x9c, 0x48, 0x81, 0xec, 0x00, 0x01, 0x00, 0x00, 0xf3, 0x44, 0x0f, 0x7f, 0xbc, 0x24, 0x10, 0xff, 0xff, 0xff, 0xf3,
+    0x44, 0x0f, 0x7f, 0xb4, 0x24, 0x20, 0xff, 0xff, 0xff, 0xf3, 0x44, 0x0f, 0x7f, 0xac, 0x24, 0x30, 0xff, 0xff, 0xff,
+    0xf3, 0x44, 0x0f, 0x7f, 0xa4, 0x24, 0x40, 0xff, 0xff, 0xff, 0xf3, 0x44, 0x0f, 0x7f, 0x9c, 0x24, 0x50, 0xff, 0xff,
+    0xff, 0xf3, 0x44, 0x0f, 0x7f, 0x94, 0x24, 0x60, 0xff, 0xff, 0xff, 0xf3, 0x44, 0x0f, 0x7f, 0x8c, 0x24, 0x70, 0xff,
+    0xff, 0xff, 0xf3, 0x44, 0x0f, 0x7f, 0x44, 0x24, 0x80, 0xf3, 0x0f, 0x7f, 0x7c, 0x24, 0x90, 0xf3, 0x0f, 0x7f, 0x74,
+    0x24, 0xa0, 0xf3, 0x0f, 0x7f, 0x6c, 0x24, 0xb0, 0xf3, 0x0f, 0x7f, 0x64, 0x24, 0xc0, 0xf3, 0x0f, 0x7f, 0x5c, 0x24,
+    0xd0, 0xf3, 0x0f, 0x7f, 0x54, 0x24, 0xe0, 0xf3, 0x0f, 0x7f, 0x4c, 0x24, 0xf0, 0xf3, 0x0f, 0x7f, 0x04, 0x24, 0x48,
+    0x8d, 0x0c, 0x24, 0x48, 0x89, 0xe3, 0x48, 0x83, 0xec, 0x30, 0x48, 0x83, 0xe4, 0xf0, 0xff, 0x15, 0xa3, 0x00, 0x00,
+    0x00, 0x48, 0x89, 0xdc, 0xf3, 0x0f, 0x6f, 0x04, 0x24, 0xf3, 0x0f, 0x6f, 0x4c, 0x24, 0x10, 0xf3, 0x0f, 0x6f, 0x54,
+    0x24, 0x20, 0xf3, 0x0f, 0x6f, 0x5c, 0x24, 0x30, 0xf3, 0x0f, 0x6f, 0x64, 0x24, 0x40, 0xf3, 0x0f, 0x6f, 0x6c, 0x24,
+    0x50, 0xf3, 0x0f, 0x6f, 0x74, 0x24, 0x60, 0xf3, 0x0f, 0x6f, 0x7c, 0x24, 0x70, 0xf3, 0x44, 0x0f, 0x6f, 0x84, 0x24,
+    0x80, 0x00, 0x00, 0x00, 0xf3, 0x44, 0x0f, 0x6f, 0x8c, 0x24, 0x90, 0x00, 0x00, 0x00, 0xf3, 0x44, 0x0f, 0x6f, 0x94,
+    0x24, 0xa0, 0x00, 0x00, 0x00, 0xf3, 0x44, 0x0f, 0x6f, 0x9c, 0x24, 0xb0, 0x00, 0x00, 0x00, 0xf3, 0x44, 0x0f, 0x6f,
+    0xa4, 0x24, 0xc0, 0x00, 0x00, 0x00, 0xf3, 0x44, 0x0f, 0x6f, 0xac, 0x24, 0xd0, 0x00, 0x00, 0x00, 0xf3, 0x44, 0x0f,
+    0x6f, 0xb4, 0x24, 0xe0, 0x00, 0x00, 0x00, 0xf3, 0x44, 0x0f, 0x6f, 0xbc, 0x24, 0xf0, 0x00, 0x00, 0x00, 0x48, 0x81,
+    0xc4, 0x00, 0x01, 0x00, 0x00, 0x9d, 0x41, 0x5f, 0x41, 0x5e, 0x41, 0x5d, 0x41, 0x5c, 0x41, 0x5b, 0x41, 0x5a, 0x41,
+    0x59, 0x41, 0x58, 0x5f, 0x5e, 0x5a, 0x59, 0x5b, 0x58, 0x5d, 0x5c, 0xc3, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00});
 #else
-constexpr uint8_t asm_data[] = {0x54, 0x55, 0x50, 0x53, 0x51, 0x52, 0x56, 0x57, 0x9C, 0x54, 0xFF, 0x15, 0x00, 0x00,
-    0x00, 0x00, 0x83, 0xC4, 0x04, 0x9D, 0x5F, 0x5E, 0x5A, 0x59, 0x5B, 0x58, 0x5D, 0x5C, 0xFF, 0x25, 0x00, 0x00, 0x00,
-    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
+constexpr auto asm_data = std::to_array<uint8_t>({0xff, 0x35, 0x91, 0x00, 0x00, 0x00, 0x54, 0x55, 0x50, 0x53, 0x51,
+    0x52, 0x56, 0x57, 0x9c, 0x81, 0xec, 0x80, 0x00, 0x00, 0x00, 0xf3, 0x0f, 0x7f, 0x7c, 0x24, 0x90, 0xf3, 0x0f, 0x7f,
+    0x74, 0x24, 0xa0, 0xf3, 0x0f, 0x7f, 0x6c, 0x24, 0xb0, 0xf3, 0x0f, 0x7f, 0x64, 0x24, 0xc0, 0xf3, 0x0f, 0x7f, 0x5c,
+    0x24, 0xd0, 0xf3, 0x0f, 0x7f, 0x54, 0x24, 0xe0, 0xf3, 0x0f, 0x7f, 0x4c, 0x24, 0xf0, 0xf3, 0x0f, 0x7f, 0x04, 0x24,
+    0x54, 0xff, 0x15, 0x8d, 0x00, 0x00, 0x00, 0x83, 0xc4, 0x04, 0xf3, 0x0f, 0x6f, 0x04, 0x24, 0xf3, 0x0f, 0x6f, 0x4c,
+    0x24, 0x10, 0xf3, 0x0f, 0x6f, 0x54, 0x24, 0x20, 0xf3, 0x0f, 0x6f, 0x5c, 0x24, 0x30, 0xf3, 0x0f, 0x6f, 0x64, 0x24,
+    0x40, 0xf3, 0x0f, 0x6f, 0x6c, 0x24, 0x50, 0xf3, 0x0f, 0x6f, 0x74, 0x24, 0x60, 0xf3, 0x0f, 0x6f, 0x7c, 0x24, 0x70,
+    0x81, 0xc4, 0x80, 0x00, 0x00, 0x00, 0x9d, 0x5f, 0x5e, 0x5a, 0x59, 0x5b, 0x58, 0x5d, 0x5c, 0xc3, 0x00, 0x00, 0x00,
+    0x00, 0x00, 0x00, 0x00, 0x00});
 #endif
 
 std::expected<MidHook, MidHook::Error> MidHook::create(void* target, MidHookFn destination) {
@@ -64,7 +85,7 @@ std::expected<void, MidHook::Error> MidHook::setup(
     m_target = target;
     m_destination = destination;
 
-    auto stub_allocation = allocator->allocate(sizeof(asm_data));
+    auto stub_allocation = allocator->allocate(asm_data.size());
 
     if (!stub_allocation) {
         return std::unexpected{Error::bad_allocation(stub_allocation.error())};
@@ -72,7 +93,7 @@ std::expected<void, MidHook::Error> MidHook::setup(
 
     m_stub = std::move(*stub_allocation);
 
-    std::copy_n(asm_data, sizeof(asm_data), m_stub.data());
+    std::copy(asm_data.begin(), asm_data.end(), m_stub.data());
 
 #ifdef _M_X64
     store(m_stub.data() + sizeof(asm_data) - 16, m_destination);
@@ -80,8 +101,8 @@ std::expected<void, MidHook::Error> MidHook::setup(
     store(m_stub.data() + sizeof(asm_data) - 8, m_destination);
 
     // 32-bit has some relocations we need to fix up as well.
-    store(m_stub.data() + 0xA + 2, m_stub.data() + sizeof(asm_data) - 8);
-    store(m_stub.data() + 0x1C + 2, m_stub.data() + sizeof(asm_data) - 4);
+    store(m_stub.data() + 0x02, m_stub.data() + m_stub.size() - 4);
+    store(m_stub.data() + 0x47, m_stub.data() + m_stub.size() - 8);
 #endif
 
     auto hook_result = InlineHook::create(allocator, m_target, m_stub.data());
diff --git a/src/mid_hook.x86.asm b/src/mid_hook.x86.asm
index e8b9fc0..b648d2f 100644
--- a/src/mid_hook.x86.asm
+++ b/src/mid_hook.x86.asm
@@ -1,4 +1,7 @@
+bits 32
+
 ; save context
+push dword [trampoline]
 push esp
 push ebp
 push eax
@@ -8,6 +11,15 @@ push edx
 push esi
 push edi
 pushfd
+sub esp, 128
+movdqu [esp-112], xmm7
+movdqu [esp-96], xmm6
+movdqu [esp-80], xmm5
+movdqu [esp-64], xmm4
+movdqu [esp-48], xmm3
+movdqu [esp-32], xmm2
+movdqu [esp-16], xmm1
+movdqu [esp], xmm0
 
 ; call destination
 push esp
@@ -15,6 +27,15 @@ call [destination]
 add esp, 4
 
 ; restore context
+movdqu xmm0, [esp]
+movdqu xmm1, [esp+16]
+movdqu xmm2, [esp+32]
+movdqu xmm3, [esp+48]
+movdqu xmm4, [esp+64]
+movdqu xmm5, [esp+80]
+movdqu xmm6, [esp+96]
+movdqu xmm7, [esp+112]
+add esp, 128
 popfd
 pop edi
 pop esi
@@ -24,10 +45,9 @@ pop ebx
 pop eax
 pop ebp
 pop esp
-
-jmp [trampoline]
+ret
 
 destination:
-.dd 0
+dd 0
 trampoline:
-.dd 0
\ No newline at end of file
+dd 0
\ No newline at end of file
diff --git a/src/mid_hook.x86_64.asm b/src/mid_hook.x86_64.asm
index 74e1afd..148e177 100644
--- a/src/mid_hook.x86_64.asm
+++ b/src/mid_hook.x86_64.asm
@@ -1,4 +1,7 @@
- ; save context
+bits 64
+
+; save context
+push qword [rel trampoline]
 push rsp
 push rbp
 push rax
@@ -16,6 +19,23 @@ push r13
 push r14
 push r15
 pushfq
+sub rsp, 256
+movdqu [rsp-240], xmm15
+movdqu [rsp-224], xmm14
+movdqu [rsp-208], xmm13
+movdqu [rsp-192], xmm12
+movdqu [rsp-176], xmm11
+movdqu [rsp-160], xmm10
+movdqu [rsp-144], xmm9
+movdqu [rsp-128], xmm8
+movdqu [rsp-112], xmm7
+movdqu [rsp-96], xmm6
+movdqu [rsp-80], xmm5
+movdqu [rsp-64], xmm4
+movdqu [rsp-48], xmm3
+movdqu [rsp-32], xmm2
+movdqu [rsp-16], xmm1
+movdqu [rsp], xmm0
 
 ; set destination parameter
 lea rcx, [rsp]
@@ -26,12 +46,29 @@ sub rsp, 48
 and rsp, -16
 
 ; call destination
-call [destination]
+call [rel destination]
 
 ; restore stack
 mov rsp, rbx
 
 ; restore context
+movdqu xmm0, [rsp]
+movdqu xmm1, [rsp+16]
+movdqu xmm2, [rsp+32]
+movdqu xmm3, [rsp+48]
+movdqu xmm4, [rsp+64]
+movdqu xmm5, [rsp+80]
+movdqu xmm6, [rsp+96]
+movdqu xmm7, [rsp+112]
+movdqu xmm8, [rsp+128]
+movdqu xmm9, [rsp+144]
+movdqu xmm10, [rsp+160]
+movdqu xmm11, [rsp+176]
+movdqu xmm12, [rsp+192]
+movdqu xmm13, [rsp+208]
+movdqu xmm14, [rsp+224]
+movdqu xmm15, [rsp+240]
+add rsp, 256
 popfq
 pop r15
 pop r14
@@ -49,10 +86,9 @@ pop rbx
 pop rax
 pop rbp
 pop rsp
-
-jmp [trampoline]
+ret
 
 destination:
-.dq 0
+dq 0
 trampoline:
-.dq 0
\ No newline at end of file
+dq 0
\ No newline at end of file
diff --git a/unittest/mid_hook.cpp b/unittest/mid_hook.cpp
index e611c11..675c49e 100644
--- a/unittest/mid_hook.cpp
+++ b/unittest/mid_hook.cpp
@@ -33,4 +33,32 @@ TEST_CASE("Mid hook to change a register", "[mid_hook]") {
     hook.reset();
 
     REQUIRE(Target::add_42(2) == 44);
-}
\ No newline at end of file
+}
+
+#ifdef _M_X64
+TEST_CASE("Mid hook to change an XMM register", "[mid_hook]") {
+    struct Target {
+        __declspec(noinline) static float __fastcall add_42(float a) { return a + 0.42f; }
+    };
+
+    REQUIRE(Target::add_42(0.0f) == 0.42f);
+
+    static SafetyHookMid hook;
+
+    struct Hook {
+        static void add_42(SafetyHookContext& ctx) { ctx.xmm0.f32[0] = 1337.0f - 0.42f; }
+    };
+
+    auto hook_result = SafetyHookMid::create(Target::add_42, Hook::add_42);
+
+    REQUIRE(hook_result);
+
+    hook = std::move(*hook_result);
+
+    REQUIRE(Target::add_42(1.0f) == 1337.0f);
+
+    hook.reset();
+
+    REQUIRE(Target::add_42(2.0f) == 2.42f);
+}
+#endif
\ No newline at end of file