Add flag to support allocation of trampolines within 32-bit offset range of target code on 64-bit Windows

Plus minor formatting fixes and more documentation for flags.

Fixes #9

6 files changed, 100 insertions, 31 deletions

diff --git a/subhook.h b/subhook.h
index 5633f39..ef9d8b8 100644
--- a/subhook.h
+++ b/subhook.h
@@ -32,7 +32,7 @@
 #if defined _M_IX86 || defined __i386__
   #define SUBHOOK_X86
   #define SUBHOOK_BITS 32
-#elif defined _M_AMD64 || __amd64__
+#elif defined _M_AMD64 || defined __amd64__
   #define SUBHOOK_X86_64
   #define SUBHOOK_BITS 64
 #else
@@ -44,7 +44,7 @@
 #elif defined __linux__ \
    || defined __FreeBSD__ || defined __OpenBSD__ || defined __NetBSD__
   #define SUBHOOK_UNIX
-  #elif defined __APPLE__
+#elif defined __APPLE__
   #define SUBHOOK_APPLE
   #define SUBHOOK_UNIX
 #else
@@ -93,8 +93,38 @@
 #endif
 
 typedef enum subhook_flags {
-  /* Use the 64-bit jump method on x86-64 (requires more space). */
-  SUBHOOK_64BIT_OFFSET = 1
+  /*
+   * Use the 64-bit jump method on x86-64. Unlike the classical 32-bit JMP,
+   * this approach ensures that the destination code can be reached from any
+   * point in the 64-bit address space, even if the source and destination are
+   * more than 4GB away from each other (meaning we are not limited to using
+   * JMP 32-bit offsets).
+   *
+   * Keep in mind that it requires overwriting a few more leading instructions
+   * inside the target code, thus it may not work with extremely short
+   * functions (14 bytes vs 5 bytes).
+   *
+   * Credits to @Ozymandias117 and @RomanHargrave on GitHub for implementing
+   * this in subhook.
+   */
+  SUBHOOK_64BIT_OFFSET = 0x01,
+  /*
+   * Generate a trampoline for jumping back to the original code faster (without
+   * removing the hook each time).
+   *
+   * In some scenarios, trampolines cannot be created. See "Known limitations"
+   * in the README file.
+   */
+  SUBHOOK_TRAMPOLINE = 0x02,
+  /*
+   * Windows x64 only: Try to allocate a trampoline buffer within +/- 2GB range
+   * of the original function to overcome a possible issue with relocating memory
+   * referencing instructions, particularly those which use RIP-relative
+   * addresses (i.e. with 32-bit offsets).
+   *
+   * Caution: this feature may slow down your code.
+   */
+  SUBHOOK_TRAMPOLINE_ALLOC_NEARBY = 0x04
 } subhook_flags_t;
 
 struct subhook_struct;
@@ -149,7 +179,9 @@ namespace subhook {
 
 enum HookFlags {
   HookNoFlags = 0,
-  HookFlag64BitOffset = SUBHOOK_64BIT_OFFSET
+  HookFlag64BitOffset = SUBHOOK_64BIT_OFFSET,
+  HookFlagTrampoline = SUBHOOK_TRAMPOLINE,
+  HookFlagTrampolineAllocNearby = SUBHOOK_TRAMPOLINE_ALLOC_NEARBY
 };
 
 inline HookFlags operator|(HookFlags o1, HookFlags o2) {
diff --git a/subhook_private.h b/subhook_private.h
index 4b345af..e89aa95 100644
--- a/subhook_private.h
+++ b/subhook_private.h
@@ -49,7 +49,7 @@ struct subhook_struct {
 };
 
 int subhook_unprotect(void *address, size_t size);
-void *subhook_alloc_code(size_t size);
+void *subhook_alloc_code(void *target_address, size_t size);
 int subhook_free_code(void *address, size_t size);
 
 #endif /* SUBHOOK_PRIVATE_H */
diff --git a/subhook_windows.c b/subhook_windows.c
index 321207e..8c9b477 100644
--- a/subhook_windows.c
+++ b/subhook_windows.c
@@ -27,6 +27,8 @@
 #include <stddef.h>
 #include <windows.h>
 
+typedef __int64 QWORD;
+
 #define SUBHOOK_CODE_PROTECT_FLAGS PAGE_EXECUTE_READWRITE
 
 int subhook_unprotect(void *address, size_t size) {
@@ -38,7 +40,38 @@ int subhook_unprotect(void *address, size_t size) {
   return !result;
 }
 
-void *subhook_alloc_code(size_t size) {
+void *subhook_alloc_code(void *target_address, size_t size) {
+  SYSTEM_INFO sys_info;
+  DWORD page_size = 0x1000;
+  QWORD offset;
+
+#if defined _M_AMD64 || defined __amd64__
+  if (target_address != NULL) {
+    GetSystemInfo(&sys_info);
+    page_size = sys_info.dwPageSize;
+
+    QWORD pivot = (QWORD)target_address & ~((QWORD)page_size - 1);
+    void *result;
+
+    for (offset = page_size; offset <= ((QWORD)1 << 31); offset += page_size) {
+      result = VirtualAlloc((void *)(pivot - offset),
+                            size,
+                            MEM_COMMIT | MEM_RESERVE,
+                            SUBHOOK_CODE_PROTECT_FLAGS);
+      if (result != NULL) {
+        return result;
+      }
+      result = VirtualAlloc((void *)(pivot + offset),
+                            size,
+                            MEM_COMMIT | MEM_RESERVE,
+                            SUBHOOK_CODE_PROTECT_FLAGS);
+      if (result != NULL) {
+        return result;
+      }
+    }
+  }
+#endif
+
   return VirtualAlloc(NULL,
                       size,
                       MEM_COMMIT | MEM_RESERVE,
diff --git a/subhook_x86.c b/subhook_x86.c
index 0941be4..71c25f9 100644
--- a/subhook_x86.c
+++ b/subhook_x86.c
@@ -488,18 +488,22 @@ SUBHOOK_EXPORT subhook_t SUBHOOK_API subhook_new(void *src,
     goto error_exit;
   }
 
-  hook->trampoline = subhook_alloc_code(hook->trampoline_size);
-  if (hook->trampoline != NULL) {
-    error = subhook_make_trampoline(hook->trampoline,
-                                    hook->src,
-                                    hook->jmp_size,
-                                    &hook->trampoline_len,
-                                    hook->flags);
-    if (error != 0) {
-      subhook_free_code(hook->trampoline, hook->trampoline_size);
-      hook->trampoline = NULL;
-      hook->trampoline_size = 0;
-      hook->trampoline_len = 0;
+  if (hook->flags & SUBHOOK_TRAMPOLINE) {
+    hook->trampoline = subhook_alloc_code(
+      (hook->flags & SUBHOOK_TRAMPOLINE_ALLOC_NEARBY) ? hook->src : NULL,
+      hook->trampoline_size);
+    if (hook->trampoline != NULL) {
+      error = subhook_make_trampoline(hook->trampoline,
+                                      hook->src,
+                                      hook->jmp_size,
+                                      &hook->trampoline_len,
+                                      hook->flags);
+      if (error != 0) {
+        subhook_free_code(hook->trampoline, hook->trampoline_size);
+        hook->trampoline = NULL;
+        hook->trampoline_size = 0;
+        hook->trampoline_len = 0;
+      }
     }
   }
 
diff --git a/tests/test.c b/tests/test.c
index d8b39fd..a3ecb73 100644
--- a/tests/test.c
+++ b/tests/test.c
@@ -31,9 +31,8 @@ void foo_hooked_tr(void) {
 int main() {
   puts("Testing initial install");
 
-  subhook_t foo_hook = subhook_new((void *)foo,
-                                   (void *)foo_hooked,
-                                   SUBHOOK_64BIT_OFFSET);
+  subhook_t foo_hook = subhook_new(
+    foo, foo_hooked, SUBHOOK_64BIT_OFFSET | SUBHOOK_TRAMPOLINE);
   if (foo_hook == NULL || subhook_install(foo_hook) < 0) {
     puts("Install failed");
     return EXIT_FAILURE;
@@ -62,9 +61,8 @@ int main() {
 
   puts("Testing trampoline");
 
-  subhook_t foo_hook_tr = subhook_new((void *)foo,
-                                      (void *)foo_hooked_tr,
-                                      SUBHOOK_64BIT_OFFSET);
+  subhook_t foo_hook_tr = subhook_new(
+    foo, foo_hooked_tr, SUBHOOK_64BIT_OFFSET | SUBHOOK_TRAMPOLINE);
   if (subhook_install(foo_hook_tr) < 0) {
     puts("Install failed");
     return EXIT_FAILURE;
diff --git a/tests/test.cpp b/tests/test.cpp
index 7ba167a..f31d1fb 100644
--- a/tests/test.cpp
+++ b/tests/test.cpp
@@ -30,9 +30,10 @@ void foo_hooked_tr() {
 int main() {
   std::cout << "Testing initial install" << std::endl;
 
-  subhook::Hook foo_hook((void *)foo,
-                         (void *)foo_hooked,
-                         subhook::HookFlag64BitOffset);
+  subhook::Hook foo_hook(
+    (void *)foo,
+    (void *)foo_hooked,
+    subhook::HookFlag64BitOffset | subhook::HookFlagTrampoline);
   if (!foo_hook.Install()) {
     std::cout << "Install failed" << std::endl;
     return EXIT_FAILURE;
@@ -59,9 +60,10 @@ int main() {
 
   std::cout << "Testing trampoline" << std::endl;
 
-  subhook::Hook foo_hook_tr((void *)foo,
-                            (void *)foo_hooked_tr,
-                            subhook::HookFlag64BitOffset);
+  subhook::Hook foo_hook_tr(
+    (void *)foo,
+    (void *)foo_hooked_tr,
+    subhook::HookFlag64BitOffset | subhook::HookFlagTrampoline);
   if (!foo_hook_tr.Install()) {
     std::cout << "Install failed" << std::endl;
     return EXIT_FAILURE;