From 50540b049ee83ab47641d1d36c9465e1ff39983d Mon Sep 17 00:00:00 2001
From: Burnt-o <52075740+Burnt-o@users.noreply.github.com>
Date: Fri, 24 Jan 2025 23:57:42 +1030
Subject: [PATCH] Test only the pages spanning VirtualProtect

---
 src/os.windows.cpp | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/os.windows.cpp b/src/os.windows.cpp
index 06d47e3..28a9af3 100644
--- a/src/os.windows.cpp
+++ b/src/os.windows.cpp
@@ -259,12 +259,10 @@ void trap_threads(uint8_t* from, uint8_t* to, size_t len, const std::function<vo
     MEMORY_BASIC_INFORMATION find_me_mbi{};
     MEMORY_BASIC_INFORMATION from_mbi{};
     MEMORY_BASIC_INFORMATION to_mbi{};
-    MEMORY_BASIC_INFORMATION virtual_protect_mbi{};
 
     VirtualQuery(reinterpret_cast<void*>(find_me), &find_me_mbi, sizeof(find_me_mbi));
     VirtualQuery(from, &from_mbi, sizeof(from_mbi));
     VirtualQuery(to, &to_mbi, sizeof(to_mbi));
-    VirtualQuery(reinterpret_cast<void*>(VirtualProtect), &virtual_protect_mbi, sizeof(virtual_protect_mbi));
 
     auto new_protect = PAGE_READWRITE;
 
@@ -272,8 +270,11 @@ void trap_threads(uint8_t* from, uint8_t* to, size_t len, const std::function<vo
         new_protect = PAGE_EXECUTE_READWRITE;
     }
 
-    if (from_mbi.BaseAddress == virtual_protect_mbi.BaseAddress ||
-        to_mbi.BaseAddress == virtual_protect_mbi.BaseAddress) {
+    auto si = system_info();
+
+    // Check if the target shares a memory page with VirtualProtect
+    if (reinterpret_cast<uint8_t*>(&VirtualProtect) <= align_up(from + len, si.page_size) &&
+        align_down(from, si.page_size) <= reinterpret_cast<uint8_t*>(&VirtualProtect) + 0x20) {
         new_protect = PAGE_EXECUTE_READWRITE;
     }