Plugins/Lua: Add initial UFunction hooking

include/uevr/API.h

@@ -36,7 +36,7 @@ SOFTWARE.
 #define UEVR_OUT
 
 #define UEVR_PLUGIN_VERSION_MAJOR 2
-#define UEVR_PLUGIN_VERSION_MINOR 27
+#define UEVR_PLUGIN_VERSION_MINOR 28
 #define UEVR_PLUGIN_VERSION_PATCH 0
 
 #define UEVR_RENDERER_D3D11 0
@@ -191,6 +191,9 @@ typedef bool (*UEVR_Engine_TickFn)(UEVR_Engine_TickCb);
 typedef bool (*UEVR_Slate_DrawWindow_RenderThreadFn)(UEVR_Slate_DrawWindow_RenderThreadCb);
 typedef bool (*UEVR_Stereo_CalculateStereoViewOffsetFn)(UEVR_Stereo_CalculateStereoViewOffsetCb);
 typedef bool (*UEVR_ViewportClient_DrawFn)(UEVR_ViewportClient_DrawCb);
+typedef bool (*UEVR_UFunction_NativeFn)(UEVR_UObjectHandle, void*, void*); /* obj, frame, ret */
+typedef bool (*UEVR_UFunction_NativePreFn)(UEVR_UFunctionHandle, UEVR_UObjectHandle, void*, void*); /* obj, frame, ret */
+typedef bool (*UEVR_UFunction_NativePostFn)(UEVR_UFunctionHandle, UEVR_UObjectHandle, void*, void*); /* obj, frame, ret */
 
 typedef void (*UEVR_PluginRequiredVersionFn)(UEVR_PluginVersion*);
 
@@ -313,6 +316,7 @@ typedef struct {
 
 typedef struct {
     void* (*get_native_function)(UEVR_UFunctionHandle function);
+    bool (*hook_ptr)(UEVR_UFunctionHandle function, UEVR_UFunction_NativePreFn pre_hook, UEVR_UFunction_NativePostFn post_hook);
 } UEVR_UFunctionFunctions;
 
 typedef struct {

include/uevr/API.hpp

@@ -517,6 +517,14 @@ class API {
             return fn(to_handle());
         }
 
+        using UEVR_UFunction_CPPPreNative = bool(*)(API::UFunction*, API::UObject*, void*, void*);
+        using UEVR_UFunction_CPPPostNative = void(*)(API::UFunction*, API::UObject*, void*, void*);
+
+        bool hook_ptr(UEVR_UFunction_CPPPreNative pre, UEVR_UFunction_CPPPostNative post) {
+            static const auto fn = initialize()->hook_ptr;
+            return fn(to_handle(), (UEVR_UFunction_NativePreFn)pre, (UEVR_UFunction_NativePostFn)post);
+        }
+
     private:
         static inline const UEVR_UFunctionFunctions* s_functions{nullptr};
         inline static const UEVR_UFunctionFunctions* initialize() {

lua-api/lib/include/ScriptContext.hpp

@@ -2,6 +2,7 @@
 
 #include <iostream>
 #include <memory>
+#include <shared_mutex>
 
 #include <sol/sol.hpp>
 #include <uevr/API.hpp>
@@ -118,6 +119,16 @@ class ScriptContext : public std::enable_shared_from_this<ScriptContext> {
     std::vector<sol::protected_function> m_on_draw_ui_callbacks{};
     std::vector<sol::protected_function> m_on_script_reset_callbacks{};
 
+    struct UFunctionHookState {
+        std::vector<sol::protected_function> pre_hooks{};
+        std::vector<sol::protected_function> post_hooks{};
+    };
+
+    static inline std::shared_mutex m_ufunction_hooks_mtx{};
+    static inline std::unordered_map<uevr::API::UFunction*, std::unique_ptr<UFunctionHookState>> m_ufunction_hooks{};
+    static bool global_ufunction_pre_handler(uevr::API::UFunction* fn, uevr::API::UObject* obj, void* params, void* result);
+    static void global_ufunction_post_handler(uevr::API::UFunction* fn, uevr::API::UObject* obj, void* params, void* result);
+
     static void on_pre_engine_tick(UEVR_UGameEngineHandle engine, float delta_seconds);
     static void on_post_engine_tick(UEVR_UGameEngineHandle engine, float delta_seconds);
     static void on_pre_slate_draw_window_render_thread(UEVR_FSlateRHIRendererHandle renderer, UEVR_FViewportInfoHandle viewport_info);

lua-api/lib/src/ScriptContext.cpp

@@ -468,7 +468,39 @@ int ScriptContext::setup_bindings() {
         sol::base_classes, sol::bases<uevr::API::UStruct, uevr::API::UObject>(),
         "static_class", &uevr::API::UFunction::static_class,
         "call", &uevr::API::UFunction::call,
-        "get_native_function", &uevr::API::UFunction::get_native_function
+        "get_native_function", &uevr::API::UFunction::get_native_function,
+        "hook_ptr", [this](sol::this_state s, uevr::API::UFunction* fn, sol::function pre, sol::function post) {
+            if (fn == nullptr) {
+                return;
+            }
+
+            std::unique_lock _{ m_ufunction_hooks_mtx };
+
+            if (auto it = m_ufunction_hooks.find(fn); it != m_ufunction_hooks.end()) {
+                if (pre != sol::nil) {
+                    it->second->pre_hooks.push_back(pre);
+                }
+
+                if (post != sol::nil) {
+                    it->second->post_hooks.push_back(post);
+                }
+
+                return;
+            }
+
+            auto& hook = m_ufunction_hooks[fn];
+            hook = std::make_unique<UFunctionHookState>();
+
+            fn->hook_ptr(global_ufunction_pre_handler, global_ufunction_post_handler);
+
+            if (pre != sol::nil) {
+                hook->pre_hooks.push_back(pre);
+            }
+
+            if (post != sol::nil) {
+                hook->post_hooks.push_back(post);
+            }
+        }
     );
 
     m_lua.new_usertype<uevr::API::FField>("UEVR_FField",
@@ -630,6 +662,44 @@ int ScriptContext::setup_bindings() {
     return out.push(m_lua.lua_state());
 }
 
+bool ScriptContext::global_ufunction_pre_handler(uevr::API::UFunction* fn, uevr::API::UObject* obj, void* params, void* out_result) {
+    bool any_false = false;
+
+    g_contexts.for_each([=, &any_false](auto ctx) {
+        std::scoped_lock _{ ctx->m_mtx };
+        std::scoped_lock __{ ctx->m_ufunction_hooks_mtx };
+
+        auto it = ctx->m_ufunction_hooks.find(fn);
+
+        if (it != ctx->m_ufunction_hooks.end()) {
+            for (auto& cb : it->second->pre_hooks) {
+                bool result = cb(fn, obj, params, out_result);
+
+                if (!result) {
+                    any_false = true;
+                }
+            }
+        }
+    });
+
+    return !any_false;
+}
+
+void ScriptContext::global_ufunction_post_handler(uevr::API::UFunction* fn, uevr::API::UObject* obj, void* params, void* result) {
+    g_contexts.for_each([=](auto ctx) {
+        std::scoped_lock _{ ctx->m_mtx };
+        std::scoped_lock __{ ctx->m_ufunction_hooks_mtx };
+
+        auto it = ctx->m_ufunction_hooks.find(fn);
+
+        if (it != ctx->m_ufunction_hooks.end()) {
+            for (auto& cb : it->second->post_hooks) {
+                cb(fn, obj, params, result);
+            }
+        }
+    });
+}
+
 void ScriptContext::on_pre_engine_tick(UEVR_UGameEngineHandle engine, float delta_seconds) {
     g_contexts.for_each([=](auto ctx) {
         std::scoped_lock _{ ctx->m_mtx };

src/mods/PluginLoader.cpp

@@ -484,11 +484,125 @@ UEVR_UClassFunctions g_uclass_functions {
 
 #define UFUNCTION(x) ((sdk::UFunction*)x)
 
+void PluginLoader::ufunction_hook_intermediary(UEVR_UObjectHandle obj, void* params, void* out_result, sdk::UFunction* func) {
+    auto& plugin_loader = PluginLoader::get();
+    std::shared_lock _{plugin_loader->m_ufunction_hooks_mtx};
+
+    auto it = plugin_loader->m_ufunction_hooks.find(func);
+
+    // uh...
+    if (it == plugin_loader->m_ufunction_hooks.end()) {
+        return;
+    }
+
+    auto& hook = it->second;
+    std::scoped_lock __{hook->mux};
+
+    bool any_false = false;
+
+    for (auto&& cb : hook->pre_callbacks) {
+        bool result = cb((UEVR_UFunctionHandle)func, obj, params, out_result);
+
+        if (!result) {
+            any_false = true;
+        }
+    }
+
+    // Call the original
+    if (!any_false) {
+        auto orig = hook->hook->get_original<UEVR_UFunction_NativeFn>();
+
+        if (orig != nullptr) {
+            orig(obj, params, out_result);
+        }
+    }
+
+    for (auto&& cb : hook->post_callbacks) {
+        cb((UEVR_UFunctionHandle)func, obj, params, out_result);
+    }
+}
+
+bool PluginLoader::hook_ufunction_ptr(UEVR_UFunctionHandle func, UEVR_UFunction_NativePreFn pre, UEVR_UFunction_NativePostFn post) {
+    std::unique_lock _{m_ufunction_hooks_mtx};
+
+    if (func == nullptr || (pre == nullptr && post == nullptr)) {
+        spdlog::error("hook_ufunction_ptr: Invalid arguments");
+        return false;
+    }
+
+    const auto offset = sdk::UFunction::get_native_function_offset();
+
+    if (offset == 0) {
+        spdlog::error("UFunction::get_native_function_offset() returned 0");
+        return false;
+    }
+
+    auto ufunc = UFUNCTION(func);
+
+    void** native = (void**)((uintptr_t)ufunc + offset);
+
+    if (*native == nullptr) {
+        return false;
+    }
+
+    auto& existing_hook = m_ufunction_hooks[ufunc];
+
+    if (existing_hook == nullptr) {
+        existing_hook = std::make_unique<UFunctionHookState>();
+        std::scoped_lock __{existing_hook->mux};
+
+        //existing_hook->hook = std::make_unique<PointerHook>(native, dst);
+
+        using namespace asmjit;
+        using namespace asmjit::x86;
+        {
+            CodeHolder code{};
+            code.init(m_jit_runtime.environment());
+            Assembler a{&code};
+
+            a.mov(r9, ufunc);
+            a.movabs(r10, &PluginLoader::ufunction_hook_intermediary);
+            a.jmp(r10);
+
+            m_jit_runtime.add((uintptr_t*)&existing_hook->jitted_pre, &code);
+        }
+
+        existing_hook->hook = std::make_unique<PointerHook>(native, (void*)existing_hook->jitted_pre);
+
+        if (pre != nullptr) {
+            existing_hook->pre_callbacks.push_back(pre);
+        }
+
+        if (post != nullptr) {
+            existing_hook->post_callbacks.push_back(post);
+        }
+    } else {
+        std::scoped_lock __{existing_hook->mux};
+
+        if (pre != nullptr) {
+            // We dont want to call the same function multiple times, could cause issues
+            if (std::find(existing_hook->pre_callbacks.begin(), existing_hook->pre_callbacks.end(), pre) == existing_hook->pre_callbacks.end()) {
+                existing_hook->pre_callbacks.push_back(pre);
+            }
+        }
+
+        if (post != nullptr) {
+            if (std::find(existing_hook->post_callbacks.begin(), existing_hook->post_callbacks.end(), post) == existing_hook->post_callbacks.end()) {
+                existing_hook->post_callbacks.push_back(post);
+            }
+        }
+    }
+
+    return true;
+}
+
 UEVR_UFunctionFunctions g_ufunction_functions {
-    // get_native_function
-    [](UEVR_UFunctionHandle func) {
+    .get_native_function = [](UEVR_UFunctionHandle func) {
         return (void*)UFUNCTION(func)->get_native_function();
     },
+    .hook_ptr = [](UEVR_UFunctionHandle func, UEVR_UFunction_NativePreFn pre, UEVR_UFunction_NativePostFn post) -> bool {
+        return PluginLoader::get()->hook_ufunction_ptr(func, pre, post);
+    }
 };
 
 namespace uevr {
@@ -1617,6 +1731,14 @@ void PluginLoader::attempt_unload_plugins() {
         callbacks->clear();
     }
 
+    {
+        std::unique_lock _{m_ufunction_hooks_mtx};
+
+        for (auto& [ufunction, hook] : m_ufunction_hooks) {
+            hook->remove_callbacks();
+        }
+    }
+
     for (auto& pair : m_plugins) {
         FreeLibrary(pair.second);
     }

src/mods/PluginLoader.hpp

@@ -7,6 +7,7 @@
 #include <Windows.h>
 
 #include <safetyhook.hpp>
+#include <asmjit/asmjit.h>
 
 #include "Mod.hpp"
 #include "uevr/API.h"
@@ -88,13 +89,29 @@ class PluginLoader : public Mod {
     bool add_on_post_viewport_client_draw(UEVR_ViewportClient_DrawCb cb);
 
     bool remove_callback(void* cb) {
-        std::unique_lock lock{m_api_cb_mtx};
+        {
+            std::unique_lock lock{m_api_cb_mtx};
+
+            for (auto& pcb_list : m_plugin_callback_lists) {
+                auto& cb_list = *pcb_list;
+                std::erase_if(cb_list, [cb](auto& cb_func) {
+                    return cb_func == cb;
+                });
+            }
+        }
 
-        for (auto& pcb_list : m_plugin_callback_lists) {
-            auto& cb_list = *pcb_list;
-            std::erase_if(cb_list, [cb](auto& cb_func) {
-                return cb_func == cb;
-            });
+        {
+            std::unique_lock lock{ m_ufunction_hooks_mtx };
+
+            for (auto& [_, hook] : m_ufunction_hooks) {
+                std::scoped_lock _{hook->mux};
+                std::erase_if(hook->pre_callbacks, [cb](auto& cb_func) {
+                    return (void*)cb_func == cb;
+                });
+                std::erase_if(hook->post_callbacks, [cb](auto& cb_func) {
+                    return (void*)cb_func == cb;
+                });
+            }
         }
 
         return true;
@@ -124,6 +141,8 @@ class PluginLoader : public Mod {
         m_inline_hooks.erase(idx);
     }
 
+    bool hook_ufunction_ptr(UEVR_UFunctionHandle func, UEVR_UFunction_NativePreFn pre, UEVR_UFunction_NativePostFn post);
+
 private:
     std::shared_mutex m_api_cb_mtx;
     std::vector<UEVR_OnPresentCb> m_on_present_cbs{};
@@ -189,4 +208,26 @@ class PluginLoader : public Mod {
     //std::vector<InlineHookState> m_inline_hooks{};
     std::unordered_map<size_t, std::unique_ptr<InlineHookState>> m_inline_hooks{};
     size_t m_inline_hook_idx{0};
+
+    asmjit::JitRuntime m_jit_runtime{};
+
+    struct UFunctionHookState {
+        using Intermediary = void(*)(UEVR_UObjectHandle, void*, void*, sdk::UFunction*);
+        Intermediary jitted_pre{};
+
+        std::unique_ptr<PointerHook> hook{};
+        std::vector<UEVR_UFunction_NativePreFn> pre_callbacks{};
+        std::vector<UEVR_UFunction_NativePostFn> post_callbacks{};
+        std::recursive_mutex mux{};
+
+        void remove_callbacks() {
+            std::scoped_lock _{mux};
+
+            pre_callbacks.clear();
+            post_callbacks.clear();
+        }
+    };
+    std::shared_mutex m_ufunction_hooks_mtx{};
+    std::unordered_map<sdk::UFunction*, std::unique_ptr<UFunctionHookState>> m_ufunction_hooks{};
+    static void ufunction_hook_intermediary(UEVR_UObjectHandle obj, void* frame, void* result, sdk::UFunction* func);
 };