Merge branch 'praydog:master' into master

joeyhodge · web-flow · commit 0109df9647dc · 2024-06-06T17:51:23.000-04:00
diff --git a/shared/sdk/REContext.cpp b/shared/sdk/REContext.cpp
@@ -24,12 +24,21 @@ namespace sdk {
 
     sdk::VM* VM::get() {
         update_pointers();
+
+        if (s_global_context == nullptr) {
+            return nullptr;
+        }
+
         return *s_global_context;
     }
 
     REThreadContext* VM::get_thread_context(int32_t unk /*= -1*/) {
         update_pointers();
 
+        if (s_get_thread_context == nullptr) {
+            return nullptr;
+        }
+
         return s_get_thread_context(this, unk);
     }
 
@@ -86,7 +95,7 @@ namespace sdk {
         std::unique_lock lock{ s_mutex };
 
         utility::ScopeGuard sg{ [&]() {
-                s_fully_updated_pointers = true;
+                s_fully_updated_pointers = s_global_context != nullptr && s_get_thread_context != nullptr;
             } 
         };
 
@@ -148,11 +157,16 @@ namespace sdk {
             return;
         }
 
-        s_global_context = (decltype(s_global_context))utility::calculate_absolute(*ref + context_pattern->ctx_offset);
-        s_get_thread_context = (decltype(s_get_thread_context))utility::calculate_absolute(*ref + context_pattern->get_thread_context_offset);
+        const auto potential_context = (decltype(s_global_context))utility::calculate_absolute(*ref + context_pattern->ctx_offset);
+        bool found_tdb = false;
+
+        if (*potential_context == nullptr) {
+            spdlog::info("[VM::update_pointers] Context is null.");
+            return;
+        }
 
         for (auto i = 0; i < 0x20000; i += sizeof(void*)) {
-            auto ptr = *(sdk::RETypeDB**)((uintptr_t)*s_global_context + i);
+            auto ptr = *(sdk::RETypeDB**)((uintptr_t)*potential_context + i);
 
             if (ptr == nullptr || IsBadReadPtr(ptr, sizeof(void*)) || ((uintptr_t)ptr & (sizeof(void*) - 1)) != 0) {
                 continue;
@@ -164,6 +178,7 @@ namespace sdk {
                 s_tdb_version = version;
                 s_type_db_offset = i;
                 s_static_tbl_offset = s_type_db_offset - 0x30; // hope this holds true for the older gameS!!!!!!!!!!!!!!!!!!!
+                found_tdb = true;
                 spdlog::info("[VM::update_pointers] s_type_db_offset: {:x}", s_type_db_offset);
                 spdlog::info("[VM::update_pointers] s_static_tbl_offset: {:x}", s_static_tbl_offset);
                 spdlog::info("[VM::update_pointers] TDB Version: {}", version);
@@ -172,6 +187,14 @@ namespace sdk {
             }
         }
 
+        if (!found_tdb) {
+            spdlog::error("[VM::update_pointers] Unable to find TDB inside VM");
+            return;
+        }
+
+        s_global_context = potential_context;
+        s_get_thread_context = (decltype(s_get_thread_context))utility::calculate_absolute(*ref + context_pattern->get_thread_context_offset);
+
         spdlog::info("[VM::update_pointers] s_global_context: {:x}", (uintptr_t)s_global_context);
         spdlog::info("[VM::update_pointers] s_get_thread_context: {:x}", (uintptr_t)s_get_thread_context);
 
@@ -181,13 +204,14 @@ namespace sdk {
 #if TDB_VER >= 71
         if (s_global_context != nullptr && *s_global_context != nullptr) {
             auto static_tbl = (REStaticTbl**)((uintptr_t)*s_global_context + s_static_tbl_offset);
+            bool found_static_tbl_offset = false;
             if (IsBadReadPtr(*static_tbl, sizeof(void*)) || ((uintptr_t)*static_tbl & (sizeof(void*) - 1)) != 0) {
                 spdlog::info("[VM::update_pointers] Static table offset is bad, correcting...");
 
                 // We are looking for the two arrays, the static field table, and the static field "initialized table"
                 // The initialized table tells whether a specific entry in the static field table has been initialized or not
                 // so they both should have the same size, easy to find
-                for (auto i = sizeof(void*); i < 0x100; i+= sizeof(void*)) {
+                for (auto i = sizeof(void*); i < 0x100; i+= sizeof(void*)) try {
                     const auto& ptr = *(REStaticTbl**)((uintptr_t)*s_global_context + (s_type_db_offset - i));
 
                     if (IsBadReadPtr(ptr, sizeof(void*)) || ((uintptr_t)ptr & (sizeof(void*) - 1)) != 0) {
@@ -210,9 +234,23 @@ namespace sdk {
                     if (previous_count == potential_count) {
                         spdlog::info("[VM::update_pointers] Found static table at {:x} (offset {:x})", (uintptr_t)ptr, previous_offset);
                         s_static_tbl_offset = previous_offset;
+                        found_static_tbl_offset = true;
                         break;
                     }
+                } catch (...) {
+                    continue;
                 }
+            } else {
+                found_static_tbl_offset = true;
+            }
+
+            // Just make it return null if we can't find it
+            // We do this so the consumer can do while(sdk::VM::get() != nullptr) { ... } to wait for everything to be valid
+            if (!found_static_tbl_offset) {
+                spdlog::error("[VM::update_pointers] Unable to find static table offset.");
+                s_global_context = nullptr;
+                s_get_thread_context = nullptr;
+                return;
             }
         }
 #endif
diff --git a/shared/sdk/Renderer.cpp b/shared/sdk/Renderer.cpp
@@ -710,6 +710,17 @@ void RenderContext::copy_texture(Texture* dest, Texture* src, Fence& fence) {
     func(this, dest, src, fence);
 }
 
+std::optional<uint32_t> Renderer::get_render_frame() const {
+    static auto tdef = sdk::find_type_definition("via.render.Renderer");
+    static auto m = tdef != nullptr ? tdef->get_method("get_RenderFrame") : nullptr;
+
+    if (m == nullptr) {
+        return std::nullopt;
+    }
+
+    return m->call<uint32_t>(sdk::get_thread_context(), this);
+}
+
 ConstantBuffer* Renderer::get_constant_buffer(std::string_view name) const {
     static auto tdef = sdk::find_type_definition("via.render.Renderer");
     static auto t = tdef->get_type();
diff --git a/shared/sdk/Renderer.hpp b/shared/sdk/Renderer.hpp
@@ -413,6 +413,12 @@ class RenderContext {
 
 class Renderer {
 public:
+    void* get_device() const {
+        return *(void**)((uintptr_t)this + sizeof(void*)); // simple!
+    }
+
+    std::optional<uint32_t> get_render_frame() const;
+
     ConstantBuffer* get_constant_buffer(std::string_view name) const;
 
     ConstantBuffer* get_scene_info() const {
diff --git a/src/D3D12Hook.cpp b/src/D3D12Hook.cpp
@@ -33,6 +33,10 @@ void* D3D12Hook::Streamline::link_swapchain_to_cmd_queue(void* rcx, void* rdx, v
         return hook->get_original<decltype(link_swapchain_to_cmd_queue)>()(rcx, rdx, r8, r9);
     }
 
+    while (g_framework == nullptr) {
+        std::this_thread::yield();
+    }
+
     std::scoped_lock _{g_framework->get_hook_monitor_mutex()};
 
     spdlog::info("[Streamline] linkSwapchainToCmdQueue: {:x}", (uintptr_t)_ReturnAddress());
@@ -57,6 +61,39 @@ void* D3D12Hook::Streamline::link_swapchain_to_cmd_queue(void* rcx, void* rdx, v
     return result;
 }
 
+HRESULT WINAPI D3D12Hook::create_swapchain(IDXGIFactory4* factory, IUnknown* device, HWND hwnd, const DXGI_SWAP_CHAIN_DESC* desc, const DXGI_SWAP_CHAIN_FULLSCREEN_DESC* p_fullscreen_desc, IDXGIOutput* p_restrict_to_output, IDXGISwapChain** swap_chain) {
+    auto create_swap_chain_fn = s_create_swapchain_hook->get_original<decltype(D3D12Hook::create_swapchain)*>();
+
+    if (g_inside_d3d12_hook) {
+        spdlog::info("create_swapchain (inside D3D12 hook)");
+        return create_swap_chain_fn(factory, device, hwnd, desc, p_fullscreen_desc, p_restrict_to_output, swap_chain);
+    }
+
+    spdlog::info("create_swapchain called");
+
+    while (g_framework == nullptr) {
+        std::this_thread::yield();
+    }
+
+    std::scoped_lock _{g_framework->get_hook_monitor_mutex()};
+
+    bool hook_was_nullptr = g_d3d12_hook == nullptr;
+
+    if (g_d3d12_hook != nullptr && g_framework->get_d3d12_hook() != nullptr) {
+        g_framework->on_reset(); // Needed to prevent a crash due to resources hanging around
+        g_d3d12_hook->unhook(); // Removes all vtable hooks
+    }
+
+    const auto result = create_swap_chain_fn(factory, device, hwnd, desc, p_fullscreen_desc, p_restrict_to_output, swap_chain);
+
+    // rather than waiting on the hook monitor to notice the hook isn't working
+    if (!hook_was_nullptr) {
+        g_framework->hook_d3d12();
+    }
+
+    return result;
+}
+
 void D3D12Hook::hook_streamline(HMODULE dlssg_module) try {
     if (D3D12Hook::s_streamline.setup) {
         return;
@@ -325,10 +362,13 @@ bool D3D12Hook::hook() {
         return false;
     }
 
-    const auto ti = utility::rtti::get_type_info(swap_chain1);
-
     try {
-        const auto swapchain_classname = ti != nullptr ? std::string_view{ti->name()} : "unknown";
+        const auto ti = utility::rtti::get_type_info(swap_chain1);
+        const auto swapchain_classname = ti != nullptr && ti->name() != nullptr ? std::string_view{ti->name()} : "unknown";
+        const auto raw_name = ti != nullptr && ti->raw_name() != nullptr ? std::string_view{ti->raw_name()} : "unknown";
+
+        spdlog::info("Swapchain type info: {}", swapchain_classname);
+        spdlog::info("Swapchain raw type info: {}", raw_name);
         
         if (swapchain_classname.contains("interposer::DXGISwapChain")) { // DLSS3
             spdlog::info("Found Streamline (DLSSFG) swapchain during dummy initialization: {:x}", (uintptr_t)swap_chain1);
@@ -346,6 +386,8 @@ bool D3D12Hook::hook() {
 
     spdlog::info("Finding command queue offset");
 
+    m_command_queue_offset = 0;
+
     // Find the command queue offset in the swapchain
     for (auto i = 0; i < 512 * sizeof(void*); i += sizeof(void*)) {
         const auto base = (uintptr_t)swap_chain1 + i;
@@ -442,6 +484,12 @@ bool D3D12Hook::hook() {
 
         auto& present_fn = (*(void***)target_swapchain)[8]; // Present
         m_present_hook = std::make_unique<PointerHook>(&present_fn, &D3D12Hook::present);
+
+        if (s_create_swapchain_hook == nullptr) {
+            auto& create_swapchain_fn = (*(void***)factory)[15]; // CreateSwapChainForHwnd
+            s_create_swapchain_hook = std::make_unique<PointerHook>(&create_swapchain_fn, &D3D12Hook::create_swapchain);
+        }
+
         m_hooked = true;
     } catch (const std::exception& e) {
         spdlog::error("Failed to initialize hooks: {}", e.what());
@@ -468,6 +516,10 @@ bool D3D12Hook::hook() {
 }
 
 bool D3D12Hook::unhook() {
+    while (g_framework == nullptr) {
+        std::this_thread::yield();
+    }
+
     std::scoped_lock _{g_framework->get_hook_monitor_mutex()};
 
     if (!m_hooked) {
@@ -488,6 +540,10 @@ bool D3D12Hook::unhook() {
 thread_local int32_t g_present_depth = 0;
 
 HRESULT WINAPI D3D12Hook::present(IDXGISwapChain3* swap_chain, uint64_t sync_interval, uint64_t flags, void* r9) {
+    while (g_framework == nullptr) {
+        std::this_thread::yield();
+    }
+
     std::scoped_lock _{g_framework->get_hook_monitor_mutex()};
 
     auto d3d12 = g_d3d12_hook;
@@ -620,6 +676,10 @@ HRESULT WINAPI D3D12Hook::present(IDXGISwapChain3* swap_chain, uint64_t sync_int
 thread_local int32_t g_resize_buffers_depth = 0;
 
 HRESULT WINAPI D3D12Hook::resize_buffers(IDXGISwapChain3* swap_chain, UINT buffer_count, UINT width, UINT height, DXGI_FORMAT new_format, UINT swap_chain_flags) {
+    while (g_framework == nullptr) {
+        std::this_thread::yield();
+    }
+
     std::scoped_lock _{g_framework->get_hook_monitor_mutex()};
 
     spdlog::info("D3D12 resize buffers called");
@@ -705,6 +765,10 @@ HRESULT WINAPI D3D12Hook::resize_buffers(IDXGISwapChain3* swap_chain, UINT buffe
 thread_local int32_t g_resize_target_depth = 0;
 
 HRESULT WINAPI D3D12Hook::resize_target(IDXGISwapChain3* swap_chain, const DXGI_MODE_DESC* new_target_parameters) {
+    while (g_framework == nullptr) {
+        std::this_thread::yield();
+    }
+
     std::scoped_lock _{g_framework->get_hook_monitor_mutex()};
 
     spdlog::info("D3D12 resize target called");
@@ -786,21 +850,3 @@ HRESULT WINAPI D3D12Hook::resize_target(IDXGISwapChain3* swap_chain, const DXGI_
 
     return result;
 }
-
-/*HRESULT WINAPI D3D12Hook::create_swap_chain(IDXGIFactory4* factory, IUnknown* device, HWND hwnd, const DXGI_SWAP_CHAIN_DESC* desc, const DXGI_SWAP_CHAIN_FULLSCREEN_DESC* p_fullscreen_desc, IDXGIOutput* p_restrict_to_output, IDXGISwapChain** swap_chain)
-{
-    spdlog::info("D3D12 create swapchain called");
-
-    auto d3d12 = g_d3d12_hook;
-
-    d3d12->m_command_queue = (ID3D12CommandQueue*)device;
-    
-    if (d3d12->m_on_create_swap_chain) {
-        d3d12->m_on_create_swap_chain(*d3d12);
-    }
-
-    auto create_swap_chain_fn = d3d12->m_create_swap_chain_hook->get_original<decltype(D3D12Hook::create_swap_chain)>();
-
-    return create_swap_chain_fn(factory, device, hwnd, desc, p_fullscreen_desc, p_restrict_to_output, swap_chain);
-}*/
-
diff --git a/src/D3D12Hook.hpp b/src/D3D12Hook.hpp
@@ -122,9 +122,7 @@ class D3D12Hook
     bool m_ignore_next_present{false};
 
     std::unique_ptr<PointerHook> m_present_hook{};
-    //std::unique_ptr<PointerHook> m_release_hook{};
     std::unique_ptr<VtableHook> m_swapchain_hook{};
-    //std::unique_ptr<FunctionHook> m_create_swap_chain_hook{};
 
     struct Streamline {
         static void* link_swapchain_to_cmd_queue(void* rcx, void* rdx, void* r8, void* r9);
@@ -135,6 +133,9 @@ class D3D12Hook
     };
 
     static inline Streamline s_streamline{};
+
+    // This is static because unhooking it seems to cause a crash sometimes
+    static inline std::unique_ptr<PointerHook> s_create_swapchain_hook{};
     
     OnPresentFn m_on_present{ nullptr };
     OnPresentFn m_on_post_present{ nullptr };
@@ -145,6 +146,6 @@ class D3D12Hook
     static HRESULT WINAPI present(IDXGISwapChain3* swap_chain, uint64_t sync_interval, uint64_t flags, void* r9);
     static HRESULT WINAPI resize_buffers(IDXGISwapChain3* swap_chain, UINT buffer_count, UINT width, UINT height, DXGI_FORMAT new_format, UINT swap_chain_flags);
     static HRESULT WINAPI resize_target(IDXGISwapChain3* swap_chain, const DXGI_MODE_DESC* new_target_parameters);
-    //static HRESULT WINAPI create_swap_chain(IDXGIFactory4* factory, IUnknown* device, HWND hwnd, const DXGI_SWAP_CHAIN_DESC* desc, const DXGI_SWAP_CHAIN_FULLSCREEN_DESC* p_fullscreen_desc, IDXGIOutput* p_restrict_to_output, IDXGISwapChain** swap_chain);
+    static HRESULT WINAPI create_swapchain(IDXGIFactory4* factory, IUnknown* device, HWND hwnd, const DXGI_SWAP_CHAIN_DESC* desc, const DXGI_SWAP_CHAIN_FULLSCREEN_DESC* p_fullscreen_desc, IDXGIOutput* p_restrict_to_output, IDXGISwapChain** swap_chain);
 };
 
diff --git a/src/REFramework.cpp b/src/REFramework.cpp
