VmtHook: Maintain RTTI information when hooked

include/safetyhook/utility.hpp

@@ -1,8 +1,8 @@
 #pragma once
 
 #include <algorithm>
-#include <concepts>
 #include <cstdint>
+#include <type_traits>
 
 namespace safetyhook {
 template <typename T> constexpr void store(uint8_t* address, const T& value) {

include/safetyhook/vmt_hook.hpp

@@ -70,6 +70,7 @@ class VmtHook final {
     [[nodiscard]] std::expected<VmHook, Error> hook_method(size_t index, FnPtr auto new_function) {
         VmHook hook{};
 
+        ++index; // Skip RTTI pointer.
         hook.m_original_vm = m_new_vmt[index];
         hook.m_new_vm = reinterpret_cast<uint8_t*>(new_function);
         hook.m_vmt_entry = &m_new_vmt[index];

src/vmt_hook.cpp

@@ -39,11 +39,15 @@ std::expected<VmtHook, VmtHook::Error> VmtHook::create(void* object) {
     hook.m_object = object;
     hook.m_original_vmt = *reinterpret_cast<uint8_t***>(object);
 
+    // Copy pointer to RTTI.
+    hook.m_new_vmt.push_back(*(hook.m_original_vmt - 1));
+
+    // Copy virtual method pointers.
     for (auto vm = hook.m_original_vmt; *vm; ++vm) {
         hook.m_new_vmt.push_back(*vm);
     }
 
-    *reinterpret_cast<uint8_t***>(object) = hook.m_new_vmt.data();
+    *reinterpret_cast<uint8_t***>(object) = hook.m_new_vmt.data() + 1;
 
     return hook;
 }

unittest/vmt_hook.cpp

@@ -93,4 +93,42 @@ TEST_CASE("Resetting the VMT hook removes all VM hooks for that object", "[vmt_h
 
     REQUIRE(target->add_42(2) == 44);
     REQUIRE(target->add_43(2) == 45);
+}
+
+TEST_CASE("VMT hooking an object maintains correct RTTI", "[vmt_hook]") {
+    struct Interface {
+        virtual ~Interface() = default;
+        virtual int add_42(int a) = 0;
+    };
+
+    struct Target : Interface {
+        __declspec(noinline) int add_42(int a) override { return a + 42; }
+    };
+
+    std::unique_ptr<Interface> target = std::make_unique<Target>();
+
+    REQUIRE(target->add_42(0) == 42);
+    REQUIRE(dynamic_cast<Target*>(target.get()) != nullptr);
+
+    static SafetyHookVmt target_hook{};
+    static SafetyHookVm add_42_hook{};
+
+    struct Hook {
+        static int __thiscall add_42(Target* self, int a) { return add_42_hook.thiscall<int>(self, a) + 1337; }
+    };
+
+    auto vmt_result = SafetyHookVmt::create(target.get());
+
+    REQUIRE(vmt_result);
+
+    target_hook = std::move(*vmt_result);
+
+    auto vm_result = target_hook.hook_method(1, Hook::add_42);
+
+    REQUIRE(vm_result);
+
+    add_42_hook = std::move(*vm_result);
+
+    REQUIRE(target->add_42(1) == 1380);
+    REQUIRE(dynamic_cast<Target*>(target.get()) != nullptr);
 }