Merge pull request #629 from psiberx/master

Overrides fixes, TweakDB hook

src/scripting/FunctionOverride.h

@@ -38,8 +38,8 @@ struct FunctionOverride
     static bool HookRunPureScriptFunction(RED4ext::CClassFunction* apFunction, RED4ext::CScriptStack* apContext, RED4ext::CStackFrame* a3);
     static bool ExecuteChain(const CallChain& aChain, std::shared_lock<std::shared_mutex>& aLock,
                              RED4ext::IScriptable* apContext, TiltedPhoques::Vector<sol::object>* apArgs,
-                             RED4ext::CStackType* apResult, RED4ext::CScriptStack* apStack,
-                             RED4ext::CStackFrame* apFrame, char* apCode, uint8_t aParam);
+                             RED4ext::CStackType* apResult, TiltedPhoques::Vector<RED4ext::CStackType>* apOutArgs, 
+                             RED4ext::CScriptStack* apStack, RED4ext::CStackFrame* apFrame, char* apCode, uint8_t aParam);
     static sol::function WrapNextOverride(const CallChain& aChain, int aStep, sol::state& aLuaState,
                                           sol::object& aLuaContext, TiltedPhoques::Vector<sol::object>& aLuaArgs,
                                           RED4ext::CBaseFunction* apRealFunction, RED4ext::IScriptable* apRealContext,
@@ -57,7 +57,8 @@ struct FunctionOverride
     void* m_pBufferStart;
     void* m_pBuffer;
     size_t m_size{ kExecutableSize };
-    TiltedPhoques::Map<RED4ext::CClassFunction*, CallChain> m_functions;
+    TiltedPhoques::Map<RED4ext::CBaseFunction*, CallChain> m_functions;
+    TiltedPhoques::Map<RED4ext::CBaseFunction*, RED4ext::CBaseFunction*> m_trampolines;
     Scripting* m_pScripting;
     std::shared_mutex m_lock;
 };

src/scripting/LuaVM.cpp

@@ -25,7 +25,7 @@ void LuaVM::Update(float aDeltaTime)
     if (!m_initialized)
     {
         if (m_logCount.load(std::memory_order_relaxed) > 0)
-            PostInitialize();
+            PostInitializeStage2();
 
         return;
     }
@@ -46,6 +46,7 @@ void LuaVM::ReloadAllMods()
     if (m_initialized)
     {
         m_scripting.ReloadAllMods();
+        m_scripting.TriggerOnTweak();
         m_scripting.TriggerOnInit();
 
         if (CET::Get().GetOverlay().IsEnabled())
@@ -180,12 +181,17 @@ void LuaVM::RegisterTDBIDString(uint64_t aValue, uint64_t aBase, const std::stri
         m_tdbidDerivedLookup[aBase].insert(aValue);
 }
 
-void LuaVM::PostInitialize()
+void LuaVM::PostInitializeStage1()
 {
-    assert(!m_initialized);
+    m_scripting.PostInitializeStage1();
+    m_scripting.TriggerOnTweak();
+}
 
-    m_scripting.PostInitialize();
+void LuaVM::PostInitializeStage2()
+{
+    assert(!m_initialized);
 
+    m_scripting.PostInitializeStage2();
     m_scripting.TriggerOnInit();
     if (CET::Get().GetOverlay().IsEnabled())
         m_scripting.TriggerOnOverlayOpen();

src/scripting/LuaVM.h

@@ -11,6 +11,7 @@ using TSetMousePosition = BOOL(void*, HWND, long, long);
 using TTDBIDCtorDerive = TDBID*(const TDBID*, TDBID*, const char*);
 using TRunningStateRun = bool(uintptr_t, uintptr_t);
 using TSetLoadingState = uintptr_t(uintptr_t, int);
+using TTweakDBLoad = uint64_t(uintptr_t, uintptr_t);
 
 struct TDBIDLookupEntry
 {
@@ -51,7 +52,8 @@ struct LuaVM
 
     void RegisterTDBIDString(uint64_t aValue, uint64_t aBase, const std::string& acString);
 
-    void PostInitialize();
+    void PostInitializeStage1();
+    void PostInitializeStage2();
 
 protected:
 
@@ -63,6 +65,7 @@ struct LuaVM
     static TDBID* HookTDBIDCtorDerive(TDBID* apBase, TDBID* apThis, const char* acpName);
     static bool HookRunningStateRun(uintptr_t aThis, uintptr_t aApp);
     static uintptr_t HookSetLoadingState(uintptr_t aThis, int aState);
+    static uint64_t HookTweakDBLoad(uintptr_t aThis, uintptr_t aParam);
 
 private:
 
@@ -77,6 +80,7 @@ struct LuaVM
     TTDBIDCtorDerive* m_realTDBIDCtorDerive{ nullptr };
     TRunningStateRun* m_realRunningStateRun{ nullptr };
     TSetLoadingState* m_realSetLoadingState{ nullptr };
+    TTweakDBLoad* m_realTweakDBLoad{ nullptr };
 
     std::chrono::time_point<std::chrono::high_resolution_clock> m_lastframe;
 

src/scripting/LuaVM_Hooks.cpp

@@ -141,13 +141,24 @@ uintptr_t LuaVM::HookSetLoadingState(uintptr_t aThis, int aState)
     {
         std::call_once(s_initBarrier, [&luavm]()
         {
-            luavm.PostInitialize();
+            luavm.PostInitializeStage2();
         });
     }
 
     return luavm.m_realSetLoadingState(aThis, aState);
 }
 
+uint64_t LuaVM::HookTweakDBLoad(uintptr_t aThis, uintptr_t aParam)
+{
+    auto& luavm = CET::Get().GetVM();
+
+    const auto ret = luavm.m_realTweakDBLoad(aThis, aParam);
+
+    luavm.PostInitializeStage1();
+
+    return ret;
+}
+
 void LuaVM::Hook(Options& aOptions)
 {
     auto& gameImage = aOptions.GameImage;
@@ -241,6 +252,23 @@ void LuaVM::Hook(Options& aOptions)
             }
         }
     }
+
+    {
+        const mem::pattern cPattern("48 89 5C 24 10 48 89 7C 24 18 4C 89 74 24 20 55 48 8B EC 48 83 EC 70 48");
+        const mem::default_scanner cScanner(cPattern);
+        uint8_t* pLocation = cScanner(gameImage.TextRegion).as<uint8_t*>();
+
+        if (pLocation)
+        {
+            if (MH_CreateHook(pLocation, &HookTweakDBLoad, reinterpret_cast<void**>(&m_realTweakDBLoad)) != MH_OK ||
+                MH_EnableHook(pLocation) != MH_OK)
+                spdlog::error("Could not hook TweakDB::Load function!");
+            else
+            {
+                spdlog::info("TweakDB::Load function hook complete!");
+            }
+        }
+    }
 
     // Disable SetLoadingState hook temporarily and get back to log count workaround
     // as it introduces major breaking change for onInit handler.

src/scripting/ScriptContext.cpp

@@ -43,6 +43,8 @@ ScriptContext::ScriptContext(LuaSandbox& aLuaSandbox, const std::filesystem::pat
     {
         if(acName == "onInit")
             m_onInit = aCallback;
+        else if(acName == "onTweak")
+            m_onTweak = aCallback;
         else if(acName == "onShutdown")
             m_onShutdown = aCallback;
         else if(acName == "onUpdate")
@@ -155,6 +157,13 @@ const TiltedPhoques::Vector<VKBindInfo>& ScriptContext::GetBinds() const
     return m_vkBindInfos;
 }
 
+void ScriptContext::TriggerOnTweak() const
+{
+    auto state = m_sandbox.GetState();
+
+    TryLuaFunction(m_logger, m_onTweak);
+}
+
 void ScriptContext::TriggerOnInit() const
 {
     auto state = m_sandbox.GetState();

src/scripting/ScriptContext.h

@@ -12,6 +12,7 @@ struct ScriptContext
 
     const TiltedPhoques::Vector<VKBindInfo>& GetBinds() const;
 
+    void TriggerOnTweak() const;
     void TriggerOnInit() const;
     void TriggerOnUpdate(float aDeltaTime) const;
     void TriggerOnDraw() const;
@@ -30,6 +31,7 @@ struct ScriptContext
     LuaSandbox& m_sandbox;
     size_t m_sandboxID;
     sol::object m_object{ };
+    sol::function m_onTweak{ };
     sol::function m_onInit{ };
     sol::function m_onShutdown{ };
     sol::function m_onUpdate{ };

src/scripting/ScriptStore.cpp

@@ -25,10 +25,16 @@ void ScriptStore::LoadAll()
 
         auto fPath = file.path();
 
-        if (is_symlink(fPath))
-            fPath = read_symlink(fPath);
-        else if (is_symlink(fPath / "init.lua"))
-            fPath = read_symlink(fPath / "init.lua").parent_path();
+        try
+        {
+            if (is_symlink(fPath))
+                fPath = read_symlink(fPath);
+            else if (is_symlink(fPath / "init.lua"))
+                fPath = read_symlink(fPath / "init.lua").parent_path();
+        }
+        catch (std::exception& e)
+        {
+        }
 
         fPath = absolute(fPath);
         auto fPathStr = fPath.string();
@@ -68,6 +74,12 @@ const TiltedPhoques::Vector<VKBindInfo>& ScriptStore::GetBinds() const
     return m_vkBindInfos;
 }
 
+void ScriptStore::TriggerOnTweak() const
+{
+    for (const auto& kvp : m_contexts)
+        kvp.second.TriggerOnTweak();
+}
+
 void ScriptStore::TriggerOnInit() const
 {
     for (const auto& kvp : m_contexts)

src/scripting/ScriptStore.h

@@ -11,6 +11,7 @@ struct ScriptStore
 
     const TiltedPhoques::Vector<VKBindInfo>& GetBinds() const;
 
+    void TriggerOnTweak() const;
     void TriggerOnInit() const;
     void TriggerOnUpdate(float aDeltaTime) const;
     void TriggerOnDraw() const;