From 820a193f3880a929f1c8a0dbba04206b7d8bee19 Mon Sep 17 00:00:00 2001
From: cursey <cursey@live.com>
Date: Tue, 13 Jun 2023 21:16:52 -0700
Subject: [PATCH] Utility: Add is_executable utility to replace IsBadCodePtr

---
 CMakeLists.txt                 |  1 +
 include/safetyhook/utility.hpp |  2 ++
 src/utility.cpp                | 50 ++++++++++++++++++++++++++++++++++
 src/vmt_hook.cpp               |  2 +-
 4 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 src/utility.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 5e00841..6a5300b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -93,6 +93,7 @@ set(safetyhook_SOURCES
 	"src/inline_hook.cpp"
 	"src/mid_hook.cpp"
 	"src/thread_freezer.cpp"
+	"src/utility.cpp"
 	"src/vmt_hook.cpp"
 	cmake.toml
 )
diff --git a/include/safetyhook/utility.hpp b/include/safetyhook/utility.hpp
index 3813e5c..4b83b10 100644
--- a/include/safetyhook/utility.hpp
+++ b/include/safetyhook/utility.hpp
@@ -11,4 +11,6 @@ template <typename T> constexpr void store(uint8_t* address, const T& value) {
 
 template <typename T>
 concept FnPtr = requires(T f) { std::is_pointer_v<T>&& std::is_function_v<std::remove_pointer_t<T>>; };
+
+bool is_executable(uint8_t* address);
 } // namespace safetyhook
\ No newline at end of file
diff --git a/src/utility.cpp b/src/utility.cpp
new file mode 100644
index 0000000..616204a
--- /dev/null
+++ b/src/utility.cpp
@@ -0,0 +1,50 @@
+#include <Windows.h>
+
+#include <safetyhook/utility.hpp>
+
+namespace safetyhook {
+bool is_page_executable(uint8_t* address) {
+    MEMORY_BASIC_INFORMATION mbi;
+
+    if (VirtualQuery(address, &mbi, sizeof(mbi)) == 0) {
+        return false;
+    }
+
+    const auto executable_protect = PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY;
+
+    return (mbi.Protect & executable_protect) != 0;
+}
+
+bool is_executable(uint8_t* address) {
+    LPVOID image_base_ptr;
+
+    if (RtlPcToFileHeader(address, &image_base_ptr) == nullptr) {
+        return is_page_executable(address);
+    }
+
+    // Just check if the section is executable.
+    const auto* image_base = reinterpret_cast<uint8_t*>(image_base_ptr);
+    const auto* dos_hdr = reinterpret_cast<const IMAGE_DOS_HEADER*>(image_base);
+
+    if (dos_hdr->e_magic != IMAGE_DOS_SIGNATURE) {
+        return is_page_executable(address);
+    }
+
+    const auto* nt_hdr = reinterpret_cast<const IMAGE_NT_HEADERS*>(image_base + dos_hdr->e_lfanew);
+
+    if (nt_hdr->Signature != IMAGE_NT_SIGNATURE) {
+        return is_page_executable(address);
+    }
+
+    const auto* section = IMAGE_FIRST_SECTION(nt_hdr);
+
+    for (auto i = 0; i < nt_hdr->FileHeader.NumberOfSections; ++i, ++section) {
+        if (address >= image_base + section->VirtualAddress &&
+            address < image_base + section->VirtualAddress + section->Misc.VirtualSize) {
+            return (section->Characteristics & IMAGE_SCN_MEM_EXECUTE) != 0;
+        }
+    }
+
+    return is_page_executable(address);
+}
+} // namespace safetyhook
\ No newline at end of file
diff --git a/src/vmt_hook.cpp b/src/vmt_hook.cpp
index 954ec51..cec6512 100644
--- a/src/vmt_hook.cpp
+++ b/src/vmt_hook.cpp
@@ -48,7 +48,7 @@ std::expected<VmtHook, VmtHook::Error> VmtHook::create(void* object) {
     // Count the number of virtual method pointers. We start at one to account for the RTTI pointer.
     auto num_vmt_entries = 1;
 
-    for (auto vm = original_vmt; !IsBadCodePtr(reinterpret_cast<FARPROC>(*vm)); ++vm) {
+    for (auto vm = original_vmt; is_executable(*vm); ++vm) {
         ++num_vmt_entries;
     }