vtable hook.

include/FunctionHook/Hook.h

@@ -5,6 +5,7 @@
 
 namespace FunctionHooking {
     class Detour;
+    class VTable;
 }
 
 class FunctionHooking::Detour {
@@ -46,4 +47,39 @@ public:
     /// @param destination_address The function we are redirecting execution to.
     Detour(void* source_address, void* destination_address);
     ~Detour();
+};
+
+// This is *mostly* cpu inspecific.
+// TODO test on Arm64 & Risc64.
+class FunctionHooking::VTable {
+private:
+    void* class_instance = nullptr;
+    void* destination = nullptr;
+    void* original = nullptr;
+    int vtable_offset = 0;
+private:
+    void CreateHook(void* class_instance, int source_function_vtable_offset, void* destination);
+    void RemoveHook();
+public:
+    template <typename T, typename... Args>
+    T CallOriginal(Args... args) { return reinterpret_cast<T(*)(Args...)>(original)(args...); }
+public:
+    /// @returns True if our hook is present.
+    /// @note This allows you to know if something has overwritten our hook.
+    [[nodiscard]] bool Valid();
+
+    /// Create a new VTable hook.
+    /// @param class_instance The pointer to an instance of the class.
+    /// @param vtable_offset The position of the virtual function in the vtable for the *base class* of our target class instance.
+    /// @param destination The function that we are redirecting execution to.
+    /// @note vtable offsets are in the same order they appear in the class definition.
+    /// @note *every* instance of the class will be hooked. not just the class instance.
+    /// @note if the class instance provided is destroyed before we remove our hook, We can't remove our hook anymore.
+    // TODO ^^ provide a way around this.
+
+    // TODO a constructor where we can pass in the function to be hooked and
+    // find it in the vtable so you don't have to provide the offset.
+    VTable(void* class_instance, int vtable_offset, void* destination);
+    ~VTable();
+public:
 };

main.cpp

@@ -3,35 +3,28 @@
 
 using namespace FunctionHooking;
 
-Detour* voidDetour;
-Detour* stringDetour;
+VTable* virtual_hook;
 
-void original_void() {
-    std::cout << "Original void function." << std::endl;
-}
+class Base {
+public:
+    virtual void test() {};
+    Base() = default;
+};
 
-void hook_void() {
-    std::cout << "Void hook function." << std::endl;
-    voidDetour->CallOriginal<void>();
-}
+class Inherited : public Base {
+public:
+    void test() override { std::cout << "test original" << std::endl; }
+    Inherited() = default;
+};
 
-std::string original_string(const std::string& string) {
-    std::cout << string + " original function." << std::endl;
-    return string;
-}
-
-std::string hook_string(const std::string& string) {
-    std::cout << string + " hook function." << std::endl;
-
-    return stringDetour->CallOriginal<std::string>(string);
-}
+void vtable_hook_function() { std::cout << "vtable hook function." << std::endl; virtual_hook->CallOriginal<void>(); }
 
 int main() {
-    //Set up hooks.
-    voidDetour = new Detour((void*) original_void, (void*) hook_void);
-    stringDetour = new Detour((void*) original_string, (void*) hook_string);
+    auto* some_class = new Inherited();
+    virtual_hook = new VTable(some_class, 0, (void*) vtable_hook_function );
+    some_class->test();
 
-    original_void();
-    std::cout << std::endl;
-    original_string("String");
+    std::cout << virtual_hook->Valid() << std::endl;
+    delete virtual_hook;
+    delete some_class;
 }

src/linux64/Hook.cpp

@@ -37,4 +37,41 @@ bool Detour::Hooked(void* function_address) {
     && current_data[11] == empty_hook[11] && current_data[12] == empty_hook[12])
         return true;
     return false;
-}
+}
+
+void VTable::CreateHook(void* class_inst, int offset, void* dest) {
+    class_instance = class_inst;
+    destination = dest;
+    vtable_offset = offset;
+    intptr_t vtable = *((intptr_t*) class_inst);
+    intptr_t entry = vtable + sizeof(intptr_t) * offset;
+    original = (intptr_t*) *((intptr_t*) entry);
+
+    mprotect((void*) ((uintptr_t) entry & ~(sysconf(_SC_PAGE_SIZE)-1)), sysconf(_SC_PAGE_SIZE), PROT_READ| PROT_WRITE| PROT_EXEC);
+    *((intptr_t *) entry) = (intptr_t) dest;
+    mprotect((void*) ((uintptr_t) entry & ~(sysconf(_SC_PAGE_SIZE)-1)), sysconf(_SC_PAGE_SIZE), PROT_READ | PROT_EXEC);
+}
+
+void VTable::RemoveHook() {
+    intptr_t vtable = *((intptr_t*) class_instance);
+    intptr_t entry = vtable + sizeof(intptr_t) * vtable_offset;
+
+    mprotect((void*) ((uintptr_t) entry & ~(sysconf(_SC_PAGE_SIZE)-1)), sysconf(_SC_PAGE_SIZE), PROT_READ| PROT_WRITE| PROT_EXEC);
+    *((intptr_t *) entry) = (intptr_t) original;
+    mprotect((void*) ((uintptr_t) entry & ~(sysconf(_SC_PAGE_SIZE)-1)), sysconf(_SC_PAGE_SIZE), PROT_READ | PROT_EXEC);
+}
+
+VTable::VTable(void* class_instance, int vtable_offset, void* destination) {
+    CreateHook(class_instance, vtable_offset, destination);
+}
+
+VTable::~VTable() {
+    if (Valid())
+        RemoveHook();
+}
+
+bool VTable::Valid() {
+    intptr_t vtable = *((intptr_t*) class_instance);
+    intptr_t entry = vtable + sizeof(intptr_t) * vtable_offset;
+    return *((intptr_t*) entry) == (intptr_t) destination;
+}