Fix address overflow in trampoline

Allocate memory for storing the trampoline code via mmap() with MAP_32BIT flag
to make sure that it stays withing 2GB range.

Also, add missing calls to subhook_free() in the C test program (C++ calls it
implicitly already via destructor).

5 files changed, 49 insertions, 19 deletions

diff --git a/subhook_private.h b/subhook_private.h
index 9d54781..16ad8fd 100644
--- a/subhook_private.h
+++ b/subhook_private.h
@@ -48,6 +48,8 @@ struct subhook_struct {
   size_t trampoline_len;
 };
 
-void *subhook_unprotect(void *address, size_t size);
+int subhook_unprotect(void *address, size_t size);
+void *subhook_alloc_code(size_t size);
+int subhok_free_code(void *address, size_t size);
 
 #endif /* SUBHOOK_PRIVATE_H */
diff --git a/subhook_unix.c b/subhook_unix.c
index 31f927e..3bde083 100644
--- a/subhook_unix.c
+++ b/subhook_unix.c
@@ -28,15 +28,26 @@
 #include <unistd.h>
 #include <sys/mman.h>
 
-void *subhook_unprotect(void *address, size_t size) {
+#define SUBHOOK_CODE_PROTECT_FLAGS (PROT_READ | PROT_WRITE | PROT_EXEC)
+
+int subhook_unprotect(void *address, size_t size) {
   long pagesize;
 
   pagesize = sysconf(_SC_PAGESIZE);
   address = (void *)((long)address & ~(pagesize - 1));
 
-  if (mprotect(address, size, PROT_READ | PROT_WRITE | PROT_EXEC) == 0) {
-    return address;
-  } else {
-    return NULL;
-  }
+  return mprotect(address, size, SUBHOOK_CODE_PROTECT_FLAGS);
+}
+
+void *subhook_alloc_code(size_t size) {
+  return mmap(NULL,
+              size,
+              SUBHOOK_CODE_PROTECT_FLAGS,
+              MAP_PRIVATE | MAP_ANONYMOUS | MAP_32BIT,
+              -1,
+              0);
+}
+
+int subhok_free_code(void *address, size_t size) {
+  return munmap(address, size);
 }
diff --git a/subhook_windows.c b/subhook_windows.c
index b1f0be6..b934651 100644
--- a/subhook_windows.c
+++ b/subhook_windows.c
@@ -27,12 +27,25 @@
 #include <stddef.h>
 #include <windows.h>
 
-void *subhook_unprotect(void *address, size_t size) {
-  DWORD old;
+#define SUBHOOK_CODE_PROTECT_FLAGS PAGE_EXECUTE_READWRITE
 
-  if (VirtualProtect(address, size, PAGE_EXECUTE_READWRITE, &old) != 0) {
-    return address;
-  } else {
-    return NULL;
-  }
+int subhook_unprotect(void *address, size_t size) {
+  DWORD old_flags;
+  BOOL result = VirtualProtect(address,
+                               size,
+                               SUBHOOK_CODE_PROTECT_FLAGS,
+                               &old_flags);
+  return !result;
+}
+
+void *subhook_alloc_code(size_t size) {
+  return VirtualAlloc(NULL,
+                      size,
+                      MEM_COMMIT | MEM_RESERVE,
+                      SUBHOOK_CODE_PROTECT_FLAGS);
+}
+
+int subhok_free_code(void *address, size_t size) {
+  (void)size;
+  return !VirtualFree(address, 0, MEM_RELEASE);
 }
diff --git a/subhook_x86.c b/subhook_x86.c
index 3a6a8c4..4759ece 100644
--- a/subhook_x86.c
+++ b/subhook_x86.c
@@ -482,13 +482,13 @@ SUBHOOK_EXPORT subhook_t SUBHOOK_API subhook_new(void *src,
 
   memcpy(hook->code, hook->src, hook->jmp_size);
 
-  hook->trampoline = calloc(1, hook->trampoline_size);
-  if (hook->trampoline == NULL) {
+  result = subhook_unprotect(hook->src, hook->jmp_size);
+  if (result != 0) {
     goto error_exit;
   }
 
-  if (subhook_unprotect(hook->src, hook->jmp_size) == NULL
-    || subhook_unprotect(hook->trampoline, hook->trampoline_size) == NULL) {
+  hook->trampoline = subhook_alloc_code(hook->trampoline_size);
+  if (hook->trampoline == NULL) {
     goto error_exit;
   }
 
@@ -521,7 +521,7 @@ SUBHOOK_EXPORT void SUBHOOK_API subhook_free(subhook_t hook) {
   if (hook == NULL) {
     return;
   }
-  free(hook->trampoline);
+  subhok_free_code(hook->trampoline, hook->trampoline_size);
   free(hook->code);
   free(hook);
 }
diff --git a/tests/test.c b/tests/test.c
index bbf75b5..d8b39fd 100644
--- a/tests/test.c
+++ b/tests/test.c
@@ -58,6 +58,8 @@ int main() {
   }
   foo();
 
+  subhook_free(foo_hook);
+
   puts("Testing trampoline");
 
   subhook_t foo_hook_tr = subhook_new((void *)foo,
@@ -74,5 +76,7 @@ int main() {
   }
   foo();
 
+  subhook_free(foo_hook_tr);
+
   return EXIT_SUCCESS;
 }