Use red4ext state hooks

maximegmd · Feb 29, 2024 · 7e08372 · 7e08372
1 parent 6d779f3
commit 7e08372
Show file tree

Hide file tree

Showing 5 changed files with 95 additions and 43 deletions.
diff --git a/src/CET.cpp b/src/CET.cpp
@@ -8,9 +8,9 @@ using namespace std::chrono_literals;
 static std::unique_ptr<CET> s_pInstance{nullptr};
 static bool s_isRunning{true};
 
-void CET::Initialize()
+void CET::Initialize(const RED4ext::Sdk* aSdk)
 {
-    s_pInstance.reset(new CET);
+    s_pInstance.reset(new CET(aSdk));
 
     s_pInstance->GetVM().Prepare();
 }
@@ -67,7 +67,7 @@ bool CET::IsRunning() noexcept
     return s_isRunning;
 }
 
-CET::CET()
+CET::CET(const RED4ext::Sdk* aSdk)
     : m_options(m_paths)
     , m_persistentState(m_paths, m_options)
     , m_bindings(m_paths, m_options)

diff --git a/src/CET.h b/src/CET.h
@@ -6,13 +6,14 @@
 #include "VKBindings.h"
 #include "d3d12/D3D12.h"
 #include "overlay/Overlay.h"
+#include "RED4ext/Api/Sdk.hpp"
 #include "scripting/LuaVM.h"
 
 struct CET
 {
     ~CET();
 
-    static void Initialize();
+    static void Initialize(const RED4ext::Sdk* aSdk);
     static void Shutdown();
     static CET& Get();
 
@@ -27,7 +28,7 @@ struct CET
     static bool IsRunning() noexcept;
 
 private:
-    CET();
+    CET(const RED4ext::Sdk* aSdk);
 
     Paths m_paths;
     Options m_options;

diff --git a/src/dllmain.cpp b/src/dllmain.cpp
@@ -21,13 +21,15 @@ static HANDLE s_modInstanceMutex = nullptr;
 
 using namespace std::chrono_literals;
 
-static bool Initialize()
+static bool Initialize(RED4ext::PluginHandle aHandle, const RED4ext::Sdk* aSdk)
 {
     try
     {
         MH_Initialize();
 
-        CET::Initialize();
+        GameMainThread::Create(aHandle, aSdk);
+
+        CET::Initialize(aSdk);
 
         const auto& options = CET::Get().GetOptions();
 
@@ -90,13 +92,12 @@ static void Shutdown()
 RED4EXT_C_EXPORT bool RED4EXT_CALL Main(RED4ext::PluginHandle aHandle, RED4ext::EMainReason aReason, const RED4ext::Sdk* aSdk)
 {
     RED4EXT_UNUSED_PARAMETER(aHandle);
-    RED4EXT_UNUSED_PARAMETER(aSdk);
 
     switch (aReason)
     {
     case RED4ext::EMainReason::Load:
     {
-        return Initialize();
+        return Initialize(aHandle, aSdk);
         break;
     }
     case RED4ext::EMainReason::Unload:

diff --git a/src/scripting/GameHooks.cpp b/src/scripting/GameHooks.cpp
@@ -23,42 +23,27 @@ void GameMainThread::RepeatedTaskQueue::Drain()
             ++taskIt;
     }
 }
-GameMainThread::StateTickOverride::StateTickOverride(uint32_t aHash, const char* acpRealFunctionName)
+GameMainThread::StateTickOverride::StateTickOverride(const char* acpRealFunctionName)
 {
-    const RED4ext::UniversalRelocPtr<uint8_t> func(aHash);
-    Location = func.GetAddr();
-
-    if (Location)
-    {
-        if (MH_CreateHook(Location, reinterpret_cast<void*>(&GameMainThread::HookStateTick), reinterpret_cast<void**>(&RealFunction)) != MH_OK || MH_EnableHook(Location) != MH_OK)
-            Log::Error("Could not hook main thread function {}! Main thread is not completely hooked!", acpRealFunctionName);
-        else
-            Log::Info("Main thread function {} hook complete!", acpRealFunctionName);
-    }
-    else
-        Log::Error("Could not locate {}! Main thread is not completely hooked!", acpRealFunctionName);
 }
 
 GameMainThread::StateTickOverride::~StateTickOverride()
 {
-    MH_DisableHook(Location);
-
-    Location = nullptr;
-    RealFunction = nullptr;
 }
 
 bool GameMainThread::StateTickOverride::OnTick(RED4ext::IGameState* apThisState, RED4ext::CGameApplication* apGameApplication)
 {
     Tasks.Drain();
 
-    return RealFunction(apThisState, apGameApplication);
+    return true;
 }
 
+static GameMainThread* s_gameMainThread = nullptr;
+
 GameMainThread& GameMainThread::Get()
 {
-    static GameMainThread s_gameMainThread;
-
-    return s_gameMainThread;
+    assert(s_gameMainThread);
+    return *s_gameMainThread;
 }
 
 void GameMainThread::AddBaseInitializationTask(const std::function<bool()>& aFunction)
@@ -90,14 +75,76 @@ void GameMainThread::AddGenericTask(const std::function<bool()>& aFunction)
     m_genericQueue.AddTask(aFunction);
 }
 
-bool GameMainThread::HookStateTick(RED4ext::IGameState* apThisState, RED4ext::CGameApplication* apGameApplication)
+void GameMainThread::Create(RED4ext::PluginHandle aHandle, const RED4ext::Sdk* aSdk)
+{
+    if (!s_gameMainThread)
+        s_gameMainThread = new GameMainThread(aHandle, aSdk);
+}
+
+GameMainThread::GameMainThread(RED4ext::PluginHandle aHandle, const RED4ext::Sdk* sdk)
+{
+    {
+        RED4ext::GameState state{nullptr, &HookBaseInitStateTick, nullptr};
+        sdk->gameStates->Add(aHandle, RED4ext::EGameStateType::BaseInitialization, &state);
+    }
+    {
+        RED4ext::GameState state{nullptr, &HookInitStateTick, nullptr};
+        sdk->gameStates->Add(aHandle, RED4ext::EGameStateType::Initialization, &state);
+    }
+    {
+        RED4ext::GameState state{nullptr, &HookRunningStateTick, nullptr};
+        sdk->gameStates->Add(aHandle, RED4ext::EGameStateType::Running, &state);
+    }
+    {
+        RED4ext::GameState state{nullptr, &HookShutdownStateTick, nullptr};
+        sdk->gameStates->Add(aHandle, RED4ext::EGameStateType::Shutdown, &state);
+    }
+}
+
+bool GameMainThread::HookBaseInitStateTick(RED4ext::CGameApplication* apGameApplication)
 {
     auto& gmt = Get();
 
     // drain generic tasks
     gmt.m_genericQueue.Drain();
 
-    // execute specific state tasks, including original function
-    const auto cStateIndex = static_cast<size_t>(apThisState->GetType());
-    return gmt.m_stateTickOverrides[cStateIndex].OnTick(apThisState, apGameApplication);
+    gmt.m_stateTickOverrides[0].OnTick(apGameApplication->currState, apGameApplication);
+
+    return true;
+}
+
+bool GameMainThread::HookInitStateTick(RED4ext::CGameApplication* apGameApplication)
+{
+    auto& gmt = Get();
+
+    // drain generic tasks
+    gmt.m_genericQueue.Drain();
+
+    gmt.m_stateTickOverrides[1].OnTick(apGameApplication->currState, apGameApplication);
+
+    return true;
+}
+
+bool GameMainThread::HookRunningStateTick(RED4ext::CGameApplication* apGameApplication)
+{
+    auto& gmt = Get();
+
+    // drain generic tasks
+    gmt.m_genericQueue.Drain();
+
+    gmt.m_stateTickOverrides[2].OnTick(apGameApplication->currState, apGameApplication);
+
+    return false;
+}
+
+bool GameMainThread::HookShutdownStateTick(RED4ext::CGameApplication* apGameApplication)
+{
+    auto& gmt = Get();
+
+    // drain generic tasks
+    gmt.m_genericQueue.Drain();
+
+    gmt.m_stateTickOverrides[3].OnTick(apGameApplication->currState, apGameApplication);
+
+    return true;
 }
diff --git a/src/scripting/GameHooks.h b/src/scripting/GameHooks.h
@@ -1,7 +1,9 @@
 #pragma once
+#include "RED4ext/Api/Sdk.hpp"
 
 struct GameMainThread
 {
+    static void Create(RED4ext::PluginHandle aHandle, const RED4ext::Sdk* aSdk);
     static GameMainThread& Get();
 
     void AddBaseInitializationTask(const std::function<bool()>& aFunction);
@@ -11,11 +13,14 @@ struct GameMainThread
     void AddGenericTask(const std::function<bool()>& aFunction);
 
 private:
-    GameMainThread() = default;
+    GameMainThread(RED4ext::PluginHandle aHandle, const RED4ext::Sdk* sdk);
 
     using TStateTick = bool(RED4ext::IGameState*, RED4ext::CGameApplication*);
 
-    static bool HookStateTick(RED4ext::IGameState* apThisState, RED4ext::CGameApplication* apGameApplication);
+    static bool HookBaseInitStateTick(RED4ext::CGameApplication* apGameApplication);
+    static bool HookInitStateTick(RED4ext::CGameApplication* apGameApplication);
+    static bool HookRunningStateTick(RED4ext::CGameApplication* apGameApplication);
+    static bool HookShutdownStateTick(RED4ext::CGameApplication* apGameApplication);
 
     // helper task queue which executes added tasks each drain until they are finished
     struct RepeatedTaskQueue
@@ -30,21 +35,19 @@ struct GameMainThread
 
     struct StateTickOverride
     {
-        StateTickOverride(uint32_t aHash, const char* acpRealFunctionName);
+        StateTickOverride(const char* acpRealFunctionName);
         ~StateTickOverride();
 
         bool OnTick(RED4ext::IGameState*, RED4ext::CGameApplication*);
 
-        uint8_t* Location = nullptr;
-        TStateTick* RealFunction = nullptr;
         RepeatedTaskQueue Tasks;
     };
 
     std::array<StateTickOverride, 4> m_stateTickOverrides{
-        StateTickOverride(CyberEngineTweaks::AddressHashes::CBaseInitializationState_OnTick, "CBaseInitializationState::OnTick"),
-        StateTickOverride(CyberEngineTweaks::AddressHashes::CInitializationState_OnTick, "CInitializationState::OnTick"),
-        StateTickOverride(CyberEngineTweaks::AddressHashes::CRunningState_OnTick, "CRunningState::OnTick"),
-        StateTickOverride(CyberEngineTweaks::AddressHashes::CShutdownState_OnTick, "CShutdownState::OnTick")};
+        StateTickOverride("CBaseInitializationState::OnTick"),
+        StateTickOverride("CInitializationState::OnTick"),
+        StateTickOverride( "CRunningState::OnTick"),
+        StateTickOverride("CShutdownState::OnTick")};
 
     RepeatedTaskQueue m_genericQueue;
 };