From d5faee56e9fb0b7798754692ed4f1ed5579c4759 Mon Sep 17 00:00:00 2001
From: Kubo Takehiro <kubo@jiubao.org>
Date: Sat, 16 Dec 2023 21:46:39 +0900
Subject: [PATCH] Check the page size of the system at runtime.

This commit fixes the EINVAL(Invalid argument) error when mprotect() is called with
an address which isn't at a page boundary.
closes #51
---
 src/arch_arm64.c        |  3 +++
 src/funchook_internal.h | 17 ++++-------------
 src/os_unix.c           |  5 ++++-
 src/os_windows.c        | 12 +++++++++---
 4 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/src/arch_arm64.c b/src/arch_arm64.c
index 076c01f..78766c1 100644
--- a/src/arch_arm64.c
+++ b/src/arch_arm64.c
@@ -37,6 +37,9 @@
 #include "funchook_internal.h"
 #include "disasm.h"
 
+#undef PAGE_SIZE
+#define PAGE_SIZE 4096
+
 // imm26 at bit 25~0
 #define IMM26_MASK 0x03FFFFFF
 #define IMM26_OFFSET(ins) ((int64_t)(int32_t)((ins) << 6) >> 4)
diff --git a/src/funchook_internal.h b/src/funchook_internal.h
index 1a555f0..d213a47 100644
--- a/src/funchook_internal.h
+++ b/src/funchook_internal.h
@@ -101,17 +101,6 @@
 
 #define TRAMPOLINE_BYTE_SIZE (TRAMPOLINE_SIZE * sizeof(insn_t))
 
-/* This must be same with sysconf(_SC_PAGE_SIZE) on Unix
- * or the dwPageSize member of the SYSTEM_INFO structure on Windows.
- */
-#undef PAGE_SIZE
-#define PAGE_SIZE 0x1000 /* 4k */
-
-/* This must be same with the dwAllocationGranularity
- * member of the SYSTEM_INFO structure on Windows.
- */
-#define ALLOCATION_UNIT 0x10000 /* 64k */
-
 struct funchook_arg_handle {
     const size_t *base_pointer;
     uint32_t flags;
@@ -158,8 +147,10 @@ void funchook_set_error_message(funchook_t *funchook, const char *fmt, ...) ATTR
 void *funchook_hook_caller(size_t transit_addr, const size_t *base_pointer);
 
 /* Functions in os_unix.c & os_windows.c */
-extern const size_t page_size;
-extern const size_t allocation_unit; /* windows only */
+#define page_size funchook_page_size
+#define allocation_unit funchook_allocation_unit
+extern size_t page_size;
+extern size_t allocation_unit; /* Windows only */
 
 funchook_t *funchook_alloc(void);
 int funchook_free(funchook_t *funchook);
diff --git a/src/os_unix.c b/src/os_unix.c
index 1f3f577..65150d2 100644
--- a/src/os_unix.c
+++ b/src/os_unix.c
@@ -51,7 +51,7 @@
 #define MAP_ANONYMOUS MAP_ANON
 #endif
 
-const size_t page_size = PAGE_SIZE;
+size_t page_size;
 
 const char *funchook_strerror(int errnum, char *buf, size_t buflen)
 {
@@ -69,6 +69,9 @@ const char *funchook_strerror(int errnum, char *buf, size_t buflen)
 
 funchook_t *funchook_alloc(void)
 {
+    if (page_size == 0) {
+        page_size = sysconf(_SC_PAGE_SIZE);
+    }
     size_t size = ROUND_UP(funchook_size, page_size);
     void *mem = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
     if (mem == (void*)-1) {
diff --git a/src/os_windows.c b/src/os_windows.c
index 77978e5..db2d12d 100644
--- a/src/os_windows.c
+++ b/src/os_windows.c
@@ -42,10 +42,9 @@ typedef struct page_info {
     char used[1];
 } page_list_t;
 
-const size_t page_size = PAGE_SIZE; /* 4K */
-const size_t allocation_unit = ALLOCATION_UNIT; /* 64K */
+size_t page_size;
+size_t allocation_unit;
 
-static size_t max_num_pages = ALLOCATION_UNIT / PAGE_SIZE - 1; /* 15 */
 static page_list_t page_list = {
     &page_list,
     &page_list,
@@ -71,6 +70,12 @@ static const char *to_errmsg(DWORD err, char *buf, size_t bufsiz)
 
 funchook_t *funchook_alloc(void)
 {
+    if (page_size == 0) {
+        SYSTEM_INFO si;
+        GetSystemInfo(&si);
+        page_size = si.dwPageSize;
+        allocation_unit = si.dwAllocationGranularity;
+    }
     size_t size = ROUND_UP(funchook_size, page_size);
     return VirtualAlloc(NULL, size, MEM_COMMIT, PAGE_READWRITE);
 }
@@ -159,6 +164,7 @@ int funchook_page_alloc(funchook_t *funchook, funchook_page_t **page_out, uint8_
 {
     page_list_t *pl;
     funchook_page_t *page = NULL;
+    size_t max_num_pages = allocation_unit / page_size - 1;
     size_t i;
 
     for (pl = page_list.next; pl != &page_list; pl = pl->next) {