Add UObjectHook mod

CMakeLists.txt

@@ -884,6 +884,7 @@ list(APPEND ue4poc_SOURCES
 	"src/hooks/XInputHook.cpp"
 	"src/mods/FrameworkConfig.cpp"
 	"src/mods/PluginLoader.cpp"
+	"src/mods/UObjectHook.cpp"
 	"src/mods/VR.cpp"
 	"src/mods/vr/Bindings.cpp"
 	"src/mods/vr/CVarManager.cpp"
@@ -913,6 +914,7 @@ list(APPEND ue4poc_SOURCES
 	"src/hooks/XInputHook.hpp"
 	"src/mods/FrameworkConfig.hpp"
 	"src/mods/PluginLoader.hpp"
+	"src/mods/UObjectHook.hpp"
 	"src/mods/VR.hpp"
 	"src/mods/vr/CVarManager.hpp"
 	"src/mods/vr/D3D11Component.hpp"
@@ -1038,6 +1040,7 @@ list(APPEND ue4poc-nolog_SOURCES
 	"src/hooks/XInputHook.cpp"
 	"src/mods/FrameworkConfig.cpp"
 	"src/mods/PluginLoader.cpp"
+	"src/mods/UObjectHook.cpp"
 	"src/mods/VR.cpp"
 	"src/mods/vr/Bindings.cpp"
 	"src/mods/vr/CVarManager.cpp"
@@ -1067,6 +1070,7 @@ list(APPEND ue4poc-nolog_SOURCES
 	"src/hooks/XInputHook.hpp"
 	"src/mods/FrameworkConfig.hpp"
 	"src/mods/PluginLoader.hpp"
+	"src/mods/UObjectHook.hpp"
 	"src/mods/VR.hpp"
 	"src/mods/vr/CVarManager.hpp"
 	"src/mods/vr/D3D11Component.hpp"

src/Mods.cpp

@@ -5,11 +5,13 @@
 #include "mods/FrameworkConfig.hpp"
 #include "mods/VR.hpp"
 #include "mods/PluginLoader.hpp"
+#include "mods/UObjectHook.hpp"
 #include "Mods.hpp"
 
 Mods::Mods() {
     m_mods.emplace_back(FrameworkConfig::get());
     m_mods.emplace_back(VR::get());
+    m_mods.emplace_back(UObjectHook::get());
 
     m_mods.emplace_back(PluginLoader::get());
 }

src/mods/UObjectHook.cpp

@@ -0,0 +1,210 @@
+#include <utility/Logging.hpp>
+#include <utility/String.hpp>
+
+#include <sdk/UObjectBase.hpp>
+#include <sdk/UObjectArray.hpp>
+#include <sdk/UClass.hpp>
+#include <sdk/threading/GameThreadWorker.hpp>
+
+#include "UObjectHook.hpp"
+
+//#define VERBOSE_UOBJECTHOOK
+
+std::shared_ptr<UObjectHook>& UObjectHook::get() {
+    static std::shared_ptr<UObjectHook> instance = std::make_shared<UObjectHook>();
+    return instance;
+}
+
+void UObjectHook::activate() {
+    if (m_hooked) {
+        return;
+    }
+
+    if (GameThreadWorker::get().is_same_thread()) {
+        hook();
+        return;
+    }
+
+    m_wants_activate = true;
+}
+
+void UObjectHook::hook() {
+    if (m_hooked) {
+        return;
+    }
+
+    SPDLOG_INFO("[UObjectHook} Hooking UObjectBase");
+
+    m_hooked = true;
+    m_wants_activate = false;
+
+    auto destructor_fn = sdk::UObjectBase::get_destructor();
+
+    if (!destructor_fn) {
+        SPDLOG_ERROR("[UObjectHook] Failed to find UObjectBase::destructor, cannot hook UObjectBase");
+        return;
+    }
+
+    auto add_object_fn = sdk::UObjectBase::get_add_object();
+
+    if (!add_object_fn) {
+        SPDLOG_ERROR("[UObjectHook] Failed to find UObjectBase::AddObject, cannot hook UObjectBase");
+        return;
+    }
+
+    m_destructor_hook = safetyhook::create_inline((void**)destructor_fn.value(), &destructor);
+
+    if (!m_destructor_hook) {
+        SPDLOG_ERROR("[UObjectHook] Failed to hook UObjectBase::destructor, cannot hook UObjectBase");
+        return;
+    }
+
+    m_add_object_hook = safetyhook::create_inline((void**)add_object_fn.value(), &add_object);
+
+    if (!m_add_object_hook) {
+        SPDLOG_ERROR("[UObjectHook] Failed to hook UObjectBase::AddObject, cannot hook UObjectBase");
+        return;
+    }
+
+    SPDLOG_INFO("[UObjectHook] Hooked UObjectBase");
+
+    // Add all the objects that already exist
+    auto uobjectarray = sdk::FUObjectArray::get();
+
+    for (auto i = 0; i < uobjectarray->get_object_count(); ++i) {
+        auto object = uobjectarray->get_object(i);
+
+        if (object == nullptr || object->object == nullptr) {
+            continue;
+        }
+
+        add_new_object(object->object);
+    }
+
+    SPDLOG_INFO("[UObjectHook] Added {} existing objects", m_objects.size());
+
+    m_fully_hooked = true;
+}
+
+void UObjectHook::add_new_object(sdk::UObjectBase* object) {
+    std::unique_lock _{m_mutex};
+    std::unique_ptr<MetaObject> meta_object{};
+
+    if (!m_reusable_meta_objects.empty()) {
+        meta_object = std::move(m_reusable_meta_objects.back());
+        m_reusable_meta_objects.pop_back();
+    } else {
+        meta_object = std::make_unique<MetaObject>();
+    }
+
+    m_objects.insert(object);
+    meta_object->full_name = object->get_full_name();
+    meta_object->uclass = object->get_class();
+
+    m_meta_objects[object] = std::move(meta_object);
+    m_objects_by_class[object->get_class()].insert(object);
+
+#ifdef VERBOSE_UOBJECTHOOK
+    SPDLOG_INFO("Adding object {:x} {:s}", (uintptr_t)object, utility::narrow(m_meta_objects[object]->full_name));
+#endif
+}
+
+void UObjectHook::on_pre_engine_tick(sdk::UGameEngine* engine, float delta) {
+    if (m_wants_activate) {
+        hook();
+    }
+}
+
+void UObjectHook::on_draw_ui() {
+    if (ImGui::CollapsingHeader("UObjectHook")) {
+        activate();
+
+        if (!m_fully_hooked) {
+            ImGui::Text("Waiting for UObjectBase to be hooked...");
+            return;
+        }
+
+        std::shared_lock _{m_mutex};
+
+        ImGui::Text("Objects: %zu (%zu actual)", m_objects.size(), sdk::FUObjectArray::get()->get_object_count());
+
+        if (ImGui::TreeNode("Objects")) {
+            for (auto& [object, meta_object] : m_meta_objects) {
+                ImGui::Text("%s", utility::narrow(meta_object->full_name).data());
+            }
+
+            ImGui::TreePop();
+        }
+
+        if (ImGui::TreeNode("Objects by class")) {
+            for (auto& [uclass, objects] : m_objects_by_class) {
+                const auto uclass_name = utility::narrow(uclass->get_full_name());
+
+                if (ImGui::TreeNode(uclass_name.data())) {
+                    for (auto& object : objects) {
+                        ImGui::Text("%s", utility::narrow(m_meta_objects[object]->full_name).data());
+                    }
+
+                    ImGui::TreePop();
+                }
+            }
+
+            ImGui::TreePop();
+        }
+    }
+}
+
+void* UObjectHook::add_object(void* rcx, void* rdx, void* r8, void* r9) {
+    auto& hook = UObjectHook::get();
+    auto result = hook->m_add_object_hook.unsafe_call<void*>(rcx, rdx, r8, r9);
+
+    {
+        static bool is_rcx = [&]() {
+            if (!IsBadReadPtr(rcx, sizeof(void*)) && 
+                !IsBadReadPtr(*(void**)rcx, sizeof(void*)) &&
+                !IsBadReadPtr(**(void***)rcx, sizeof(void*))) 
+            {
+                SPDLOG_INFO("[UObjectHook] RCX is UObjectBase*");
+                return true;
+            } else {
+                SPDLOG_INFO("[UObjectHook] RDX is UObjectBase*");
+                return false;
+            }
+        }();
+
+        sdk::UObjectBase* obj = nullptr;
+
+        if (is_rcx) {
+            obj = (sdk::UObjectBase*)rcx;
+        } else {
+            obj = (sdk::UObjectBase*)rdx;
+        }
+
+        hook->add_new_object(obj);
+    }
+
+    return result;
+}
+
+void* UObjectHook::destructor(sdk::UObjectBase* object, void* rdx, void* r8, void* r9) {
+    auto& hook = UObjectHook::get();
+
+    {
+        std::unique_lock _{hook->m_mutex};
+
+        if (auto it = hook->m_meta_objects.find(object); it != hook->m_meta_objects.end()) {
+#ifdef VERBOSE_UOBJECTHOOK
+            SPDLOG_INFO("Removing object {:x} {:s}", (uintptr_t)object, utility::narrow(it->second->full_name));
+#endif
+            hook->m_objects.erase(object);
+            hook->m_objects_by_class[it->second->uclass].erase(object);
+
+            hook->m_reusable_meta_objects.push_back(std::move(it->second));
+            hook->m_meta_objects.erase(object);
+        }
+    }
+
+    auto result = hook->m_destructor_hook.unsafe_call<void*>(object, rdx, r8, r9);
+
+    return result;
+}

src/mods/UObjectHook.hpp

@@ -0,0 +1,62 @@
+#pragma once
+
+#include <shared_mutex>
+#include <unordered_set>
+#include <memory>
+#include <deque>
+
+#include <safetyhook.hpp>
+
+#include "Mod.hpp"
+
+namespace sdk {
+class UObjectBase;
+class UObject;
+class UClass;
+}
+
+class UObjectHook : public Mod {
+public:
+    static std::shared_ptr<UObjectHook>& get();
+    std::unordered_set<sdk::UObjectBase*> get_objects_by_class(sdk::UClass* uclass) const {
+        std::shared_lock _{m_mutex};
+        if (auto it = m_objects_by_class.find(uclass); it != m_objects_by_class.end()) {
+            return it->second;
+        }
+
+        return {};
+    }
+
+    void activate();
+
+protected:
+    void on_pre_engine_tick(sdk::UGameEngine* engine, float delta) override;
+    void on_draw_ui() override;
+
+private:
+    void hook();
+    void add_new_object(sdk::UObjectBase* object);
+
+    static void* add_object(void* rcx, void* rdx, void* r8, void* r9);
+    static void* destructor(sdk::UObjectBase* object, void* rdx, void* r8, void* r9);
+
+    bool m_hooked{false};
+    bool m_fully_hooked{false};
+    bool m_wants_activate{false};
+
+    mutable std::shared_mutex m_mutex{};
+
+    struct MetaObject {
+        std::wstring full_name{};
+        sdk::UClass* uclass{nullptr};
+    };
+
+    std::unordered_set<sdk::UObjectBase*> m_objects{};
+    std::unordered_map<sdk::UObjectBase*, std::unique_ptr<MetaObject>> m_meta_objects{};
+    std::unordered_map<sdk::UClass*, std::unordered_set<sdk::UObjectBase*>> m_objects_by_class{};
+
+    std::deque<std::unique_ptr<MetaObject>> m_reusable_meta_objects{};
+
+    SafetyHookInline m_add_object_hook{};
+    SafetyHookInline m_destructor_hook{};
+};