Allocator: Add INT3 codecave scanner

CMakeLists.txt

@@ -483,6 +483,64 @@ if(SAFETYHOOK_BUILD_TESTS) # build-tests
 		safetyhook::safetyhook
 	)
 
+endif()
+# Target: test5
+if(SAFETYHOOK_BUILD_TESTS) # build-tests
+	set(test5_SOURCES
+		"tests/test5.cpp"
+		cmake.toml
+	)
+
+	add_executable(test5)
+
+	target_sources(test5 PRIVATE ${test5_SOURCES})
+	source_group(TREE ${CMAKE_CURRENT_SOURCE_DIR} FILES ${test5_SOURCES})
+
+	target_compile_features(test5 PRIVATE
+		cxx_std_23
+	)
+
+	if(MSVC) # msvc
+		target_compile_options(test5 PRIVATE
+			"/WX"
+			"/permissive-"
+			"/W4"
+			"/w14640"
+			"/EHsc"
+		)
+	endif()
+
+	if((CMAKE_CXX_COMPILER_ID MATCHES "Clang" AND NOT CMAKE_CXX_COMPILER_FRONTEND_VARIANT MATCHES "^MSVC$") OR (CMAKE_C_COMPILER_ID MATCHES "Clang" AND NOT CMAKE_C_COMPILER_FRONTEND_VARIANT MATCHES "^MSVC$")) # clang
+		target_compile_options(test5 PRIVATE
+			-Werror
+			-Wall
+			-Wextra
+			-Wshadow
+			-Wnon-virtual-dtor
+			-pedantic
+		)
+	endif()
+
+	if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_C_COMPILER_ID STREQUAL "GNU") # gcc
+		target_compile_options(test5 PRIVATE
+			-Werror
+			-Wall
+			-Wextra
+			-Wshadow
+			-Wnon-virtual-dtor
+			-pedantic
+		)
+	endif()
+
+	target_link_libraries(test5 PRIVATE
+		safetyhook::safetyhook
+	)
+
+	get_directory_property(CMKR_VS_STARTUP_PROJECT DIRECTORY ${PROJECT_SOURCE_DIR} DEFINITION VS_STARTUP_PROJECT)
+	if(NOT CMKR_VS_STARTUP_PROJECT)
+		set_property(DIRECTORY ${PROJECT_SOURCE_DIR} PROPERTY VS_STARTUP_PROJECT test5)
+	endif()
+
 endif()
 # Target: unittest
 if(SAFETYHOOK_BUILD_TESTS) # build-tests

cmake.toml

@@ -133,6 +133,10 @@ sources = ["tests/test3.cpp"]
 type = "test-dll"
 sources = ["tests/test4.cpp"]
 
+[target.test5]
+type = "test"
+sources = ["tests/test5.cpp"]
+
 [target.unittest]
 type = "test"
 sources = ["unittest/*.cpp"]

src/allocator.cpp

@@ -125,8 +125,57 @@ std::expected<Allocation, Allocator::Error> Allocator::internal_allocate_near(
 
     GetSystemInfo(&si);
 
-    const auto allocation_size = align_up(size, si.dwAllocationGranularity);
-    const auto allocation_address = allocate_nearby_memory(desired_addresses, allocation_size, max_distance);
+    auto allocation_size = align_up(size, si.dwAllocationGranularity);
+    auto allocation_address = allocate_nearby_memory(desired_addresses, allocation_size, max_distance);
+
+    // And finally, look for a codecave within int3 padding.
+    // Not just putting this in allocate_nearby_memory because
+    // it passes an aligned size, which is not what we want.
+    if (!allocation_address && !desired_addresses.empty()) {
+        // Locate an address in desired_addresses that has executable permissions.
+        // TODO: We could potentially look through other regions not in the desired_addresses list.
+        MEMORY_BASIC_INFORMATION mbi{};
+        uint8_t* address = nullptr;
+        for (const auto& addr : desired_addresses) {
+            if (VirtualQuery(addr, &mbi, sizeof(mbi)) == 0) {
+                continue;
+            }
+
+            if (mbi.Protect & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY)) {
+                address = reinterpret_cast<uint8_t*>(addr);
+                break;
+            }
+        }
+
+        if (address == nullptr) {
+            return std::unexpected{allocation_address.error()};
+        }
+
+        const auto end = reinterpret_cast<uintptr_t>(mbi.BaseAddress) + mbi.RegionSize;
+
+        // Search for an int3 sled, starting from the target address.
+        for (auto ip = address; reinterpret_cast<uintptr_t>(ip) < end; ++ip) {
+            if (*ip == 0xCC) {
+                auto sled = ip;
+                uint32_t count = 0;
+
+                while (*sled == 0xCC) {
+                    ++sled;
+                    ++count;
+
+                    if (reinterpret_cast<uintptr_t>(sled) >= end) {
+                        break;
+                    }
+                }
+
+                if (count >= 10 && count >= size && in_range(ip, desired_addresses, max_distance)) {
+                    allocation_address = ip;
+                    allocation_size = count;
+                    break;
+                }
+            }
+        }
+    }
 
     if (!allocation_address) {
         return std::unexpected{allocation_address.error()};

src/inline_hook.cpp

@@ -241,6 +241,8 @@ std::expected<void, InlineHook::Error> InlineHook::e9_hook(const std::shared_ptr
 
     m_trampoline = std::move(*trampoline_allocation);
 
+    UnprotectMemory unprotect_trampoline{m_trampoline.data(), m_trampoline.size()};
+
     for (auto ip = m_target, tramp_ip = m_trampoline.data(); ip < m_target + m_original_bytes.size(); ip += ix.length) {
         if (!decode(&ix, ip)) {
             m_trampoline.free();
@@ -357,6 +359,8 @@ std::expected<void, InlineHook::Error> InlineHook::ff_hook(const std::shared_ptr
 
     m_trampoline = std::move(*trampoline_allocation);
 
+    UnprotectMemory unprotect_trampoline{m_trampoline.data(), m_trampoline.size()};
+
     std::copy(m_original_bytes.begin(), m_original_bytes.end(), m_trampoline.data());
 
     const auto trampoline_epilogue =

tests/test5.cpp

@@ -0,0 +1,195 @@
+// test5.cpp is test3.cpp but with 2GB of memory reserved around the target address.
+// This is to test if the allocator works correctly when there is no free memory near the target address.
+#include <iostream>
+#include <iomanip>
+#include <sstream>
+
+#include <safetyhook.hpp>
+
+#if __has_include(<Zydis/Zydis.h>)
+#include <Zydis/Zydis.h>
+#elif __has_include(<Zydis.h>)
+#include <Zydis.h>
+#else
+#error "Zydis not found"
+#endif
+
+auto operator ""_mb(unsigned long long mb) {
+    return mb * 1024 * 1024;
+}
+
+auto operator ""_gb(unsigned long long gb) {
+    return gb * 1024 * 1024 * 1024;
+}
+
+SafetyHookInline hook0, hook1, hook2, hook3;
+
+__declspec(noinline) void say_hi(const std::string& name) {
+    std::cout << "hello " << name << "\n";
+}
+
+void hook0_fn(const std::string& name) {
+    hook0.call<void, const std::string&>(name + " and bob");
+}
+
+void hook1_fn(const std::string& name) {
+    hook1.call<void, const std::string&>(name + " and alice");
+}
+
+void hook2_fn(const std::string& name) {
+    hook2.call<void, const std::string&>(name + " and eve");
+}
+
+void hook3_fn(const std::string& name) {
+    hook3.call<void, const std::string&>(name + " and carol");
+}
+
+// Intentionally takes up memory space +- 2GB around the target address.
+// so we can test if the allocator works correctly.
+void reserve_memory_2gb_around_target(uintptr_t target) {
+    // First we must obtain all currently allocated memory regions.
+    std::vector<std::pair<uintptr_t, uintptr_t>> regions{};
+    std::vector<std::pair<uintptr_t, uintptr_t>> free_regions{};
+
+    MEMORY_BASIC_INFORMATION mbi{};
+    uintptr_t address = 0;
+
+    while (VirtualQuery(reinterpret_cast<void*>(address), &mbi, sizeof(mbi)) == sizeof(mbi)) {
+        if (mbi.State == MEM_COMMIT) {
+            regions.emplace_back(reinterpret_cast<uintptr_t>(mbi.BaseAddress), reinterpret_cast<uintptr_t>(mbi.BaseAddress) + mbi.RegionSize);
+        } else if (mbi.State == MEM_FREE) {
+            free_regions.emplace_back(reinterpret_cast<uintptr_t>(mbi.BaseAddress), reinterpret_cast<uintptr_t>(mbi.BaseAddress) + mbi.RegionSize);
+        }
+
+        address += mbi.RegionSize;
+    }
+
+    std::vector<std::pair<uintptr_t, uintptr_t>> new_regions{};
+
+    // Go through the free regions and reserve memory if
+    // the distance to the target is less than 2GB.
+    for (auto& region : free_regions) {
+        if (region.first > target && region.second > target && region.first - target <= 2_gb) {
+            new_regions.emplace_back(region);
+        } else if (region.first < target && region.second < target && target - region.second <= 2_gb) {
+            new_regions.emplace_back(region);
+        }
+    }
+
+    // get allocation granularity
+    SYSTEM_INFO si{};
+    GetSystemInfo(&si);
+
+    const auto granularity = (uintptr_t)si.dwAllocationGranularity;
+
+    // Split the regions into 50mb chunks.
+    for (auto& region : new_regions) {
+        region.first = (region.first + granularity - 1) & ~(granularity - 1);
+        const auto size = region.second - region.first;
+
+        const auto chunk_size = (std::min<size_t>(10_mb, size) + granularity - 1) & ~(granularity - 1);
+        const auto num_chunks = size / chunk_size;
+
+        for (size_t i = 0; i < num_chunks; ++i) {
+            const auto wanted_alloc = region.first + i * chunk_size;
+            const auto alloced = VirtualAlloc((void*)wanted_alloc, chunk_size, MEM_RESERVE, PAGE_READWRITE);
+
+            printf("Allocated [0x%llX, 0x%llX] 0x%llX\n", (uintptr_t)alloced, (uintptr_t)alloced + chunk_size, chunk_size);
+
+            if (alloced == nullptr) {
+                const auto error = GetLastError();
+                printf("Failed to allocate to 0x%llX 0x%llX\n", wanted_alloc, chunk_size);
+                printf("Error: 0x%X\n", error);
+            }
+        }
+
+        // Handle remaining small fragment
+        auto remaining = size % chunk_size;
+        if (remaining > 0) {
+            printf("Allocating remaining 0x%llX\n", remaining);
+            auto last_alloc = region.first + num_chunks * chunk_size;
+            auto alloced = VirtualAlloc((void*)last_alloc, remaining, MEM_RESERVE, PAGE_READWRITE);
+
+            if (alloced == nullptr) {
+                const auto error = GetLastError();
+                printf("Failed to allocate to 0x%llX 0x%llX\n", last_alloc, remaining);
+                printf("Error: 0x%X\n", error);
+            }
+        }
+    }
+}
+
+int main() {
+    reserve_memory_2gb_around_target(reinterpret_cast<uintptr_t>(&say_hi));
+
+    printf("0x%llX\n", (uintptr_t)&say_hi);
+
+    uintptr_t real_say_hi = (uintptr_t)&say_hi;
+
+    if (*(uint8_t*)&say_hi == 0xE9) {
+        real_say_hi = (uintptr_t)&say_hi + *(int32_t*)((uintptr_t)&say_hi + 1) + 5;
+        printf("0x%llX\n", real_say_hi);
+    }
+
+    ZydisDecoder decoder{};
+
+#if defined(_M_X64)
+    ZydisDecoderInit(&decoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_STACK_WIDTH_64);
+#elif defined(_M_IX86)
+    ZydisDecoderInit(&decoder, ZYDIS_MACHINE_MODE_LEGACY_32, ZYDIS_STACK_WIDTH_32);
+#else
+#error "Unsupported architecture"
+#endif
+
+    ZydisFormatter formatter{};
+    ZydisFormatterInit(&formatter, ZYDIS_FORMATTER_STYLE_INTEL);
+    ZydisFormatterSetProperty(&formatter, ZYDIS_FORMATTER_PROP_FORCE_SEGMENT, ZYAN_TRUE);
+    ZydisFormatterSetProperty(&formatter, ZYDIS_FORMATTER_PROP_FORCE_SIZE, ZYAN_TRUE);
+
+    auto disassemble_say_hi = [&]() {
+        uintptr_t ip = real_say_hi;
+        for (auto i = 0; i < 10; ++i) {
+            ZydisDecodedInstruction ix{};
+            ZydisDecodedOperand operands[ZYDIS_MAX_OPERAND_COUNT];
+            ZydisDecoderDecodeFull(&decoder, reinterpret_cast<void*>(ip), 15, &ix, operands);
+
+            // Convert to text
+            char buffer[256]{};
+            ZydisFormatterInit(&formatter, ZYDIS_FORMATTER_STYLE_INTEL);
+            ZydisFormatterFormatInstruction(&formatter, &ix, operands, ZYDIS_MAX_OPERAND_COUNT, buffer, sizeof(buffer), ip, nullptr);
+
+            // Format to {bytes} {mnemonic} {operands}
+            std::stringstream bytehex{};
+
+            for (auto j = 0; j < ix.length; ++j) {
+                bytehex << std::hex << std::setfill('0') << std::setw(2) << (int)*(uint8_t*)(ip + j) << " ";
+            }
+
+            printf("0x%llX | %s | %s\n", ip, bytehex.str().c_str(), buffer);
+
+            ip += ix.length;
+        }
+    };
+
+    std::cout << "Before:" << std::endl;
+    disassemble_say_hi();
+
+    hook0 = safetyhook::create_inline(reinterpret_cast<void*>(real_say_hi), reinterpret_cast<void*>(hook0_fn));
+
+    if (!hook0) {
+        std::cout << "Failed to create hook\n";
+        return 1;
+    }
+
+    hook1 = safetyhook::create_inline(reinterpret_cast<void*>(real_say_hi), reinterpret_cast<void*>(hook1_fn));
+    hook2 = safetyhook::create_inline(reinterpret_cast<void*>(real_say_hi), reinterpret_cast<void*>(hook2_fn));
+    hook3 = safetyhook::create_inline(reinterpret_cast<void*>(real_say_hi), reinterpret_cast<void*>(hook3_fn));
+
+    std::cout << "After:" << std::endl;
+    disassemble_say_hi();
+
+    say_hi("world");
+    std::system("pause");
+
+    return 0;
+}