From 61db246c5ec86aa63215a13ad915f42da1f98426 Mon Sep 17 00:00:00 2001
From: squidbus <175574877+squidbus@users.noreply.github.com>
Date: Fri, 30 Aug 2024 20:47:07 -0700
Subject: [PATCH] core: Fix CPU patch stack issues

---
 CMakeLists.txt                                |  2 +-
 src/core/cpu_patches.cpp                      | 30 +++++++------------
 .../libraries/kernel/thread_management.cpp    |  3 +-
 src/core/linker.cpp                           |  4 +--
 4 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 8c7e03e3..a65e9c59 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -636,7 +636,7 @@ target_link_libraries(shadps4 PRIVATE Boost::headers GPUOpen::VulkanMemoryAlloca
 
 if (APPLE)
   # Reserve system-managed memory space.
-  target_link_options(shadps4 PRIVATE -Wl,-no_pie,-no_fixup_chains,-no_huge,-pagezero_size,0x400000,-segaddr,GUEST_SYSTEM,0x400000,-image_base,0x10000000000)
+  target_link_options(shadps4 PRIVATE -Wl,-no_pie,-no_fixup_chains,-no_huge,-pagezero_size,0x400000,-segaddr,GUEST_SYSTEM,0x400000,-image_base,0x20000000000)
 
   # Link MoltenVK for Vulkan support
   find_library(MOLTENVK MoltenVK REQUIRED)
diff --git a/src/core/cpu_patches.cpp b/src/core/cpu_patches.cpp
index 55bbf23b..e713155a 100644
--- a/src/core/cpu_patches.cpp
+++ b/src/core/cpu_patches.cpp
@@ -126,39 +126,35 @@ static Xbyak::Reg AllocateScratchRegister(
 static pthread_key_t stack_pointer_slot;
 static pthread_key_t patch_stack_slot;
 static std::once_flag patch_context_slots_init_flag;
+static constexpr u32 patch_stack_size = 0x1000;
 
 static_assert(sizeof(void*) == sizeof(u64),
               "Cannot fit a register inside a thread local storage slot.");
 
+static void FreePatchStack(void* patch_stack) {
+    // Subtract back to the bottom of the stack for free.
+    std::free(static_cast<u8*>(patch_stack) - patch_stack_size);
+}
+
 static void InitializePatchContextSlots() {
     ASSERT_MSG(pthread_key_create(&stack_pointer_slot, nullptr) == 0,
                "Unable to allocate thread-local register for stack pointer.");
-    ASSERT_MSG(pthread_key_create(&patch_stack_slot, nullptr) == 0,
+    ASSERT_MSG(pthread_key_create(&patch_stack_slot, FreePatchStack) == 0,
                "Unable to allocate thread-local register for patch stack.");
 }
 
 void InitializeThreadPatchStack() {
     std::call_once(patch_context_slots_init_flag, InitializePatchContextSlots);
 
-    const auto* patch_stack = std::malloc(0x1000);
-    pthread_setspecific(patch_stack_slot, patch_stack);
-}
-
-void CleanupThreadPatchStack() {
-    std::call_once(patch_context_slots_init_flag, InitializePatchContextSlots);
-
-    auto* patch_stack = pthread_getspecific(patch_stack_slot);
-    if (patch_stack != nullptr) {
-        std::free(patch_stack);
-        pthread_setspecific(patch_stack_slot, nullptr);
-    }
+    pthread_setspecific(patch_stack_slot,
+                        static_cast<u8*>(std::malloc(patch_stack_size)) + patch_stack_size);
 }
 
 /// Saves the stack pointer to thread local storage and loads the patch stack.
 static void SaveStack(Xbyak::CodeGenerator& c) {
     std::call_once(patch_context_slots_init_flag, InitializePatchContextSlots);
 
-    // Save stack pointer and load patch stack.
+    // Save original stack pointer and load patch stack.
     c.putSeg(gs);
     c.mov(qword[reinterpret_cast<void*>(stack_pointer_slot * sizeof(void*))], rsp);
     c.putSeg(gs);
@@ -184,10 +180,6 @@ void InitializeThreadPatchStack() {
     // No-op
 }
 
-void CleanupThreadPatchStack() {
-    // No-op
-}
-
 /// Saves the stack pointer to thread local storage and loads the patch stack.
 static void SaveStack(Xbyak::CodeGenerator& c) {
     UNIMPLEMENTED();
@@ -244,7 +236,7 @@ static void RestoreContext(Xbyak::CodeGenerator& c, const Xbyak::Operand& dst) {
         if (!dst.isREG() || dst.getIdx() != reg) {
             c.pop(Xbyak::Reg64(reg));
         } else {
-            c.add(rsp, 4);
+            c.add(rsp, 8);
         }
     }
     RestoreStack(c);
diff --git a/src/core/libraries/kernel/thread_management.cpp b/src/core/libraries/kernel/thread_management.cpp
index a2befd4c..5c84c35d 100644
--- a/src/core/libraries/kernel/thread_management.cpp
+++ b/src/core/libraries/kernel/thread_management.cpp
@@ -986,15 +986,14 @@ static void cleanup_thread(void* arg) {
             destructor(value);
         }
     }
-    Core::CleanupThreadPatchStack();
     thread->is_almost_done = true;
 }
 
 static void* run_thread(void* arg) {
     auto* thread = static_cast<ScePthread>(arg);
     Common::SetCurrentThreadName(thread->name.c_str());
-    auto* linker = Common::Singleton<Core::Linker>::Instance();
     Core::InitializeThreadPatchStack();
+    auto* linker = Common::Singleton<Core::Linker>::Instance();
     linker->InitTlsForThread(false);
     void* ret = nullptr;
     g_pthread_self = thread;
diff --git a/src/core/linker.cpp b/src/core/linker.cpp
index 4ef62c4a..0c914cef 100644
--- a/src/core/linker.cpp
+++ b/src/core/linker.cpp
@@ -85,8 +85,8 @@ void Linker::Execute() {
 
     // Init primary thread.
     Common::SetCurrentThreadName("GAME_MainThread");
-    Libraries::Kernel::pthreadInitSelfMainThread();
     InitializeThreadPatchStack();
+    Libraries::Kernel::pthreadInitSelfMainThread();
     InitTlsForThread(true);
 
     // Start shared library modules
@@ -106,8 +106,6 @@ void Linker::Execute() {
             RunMainEntry(m->GetEntryAddress(), &p, ProgramExitFunc);
         }
     }
-
-    CleanupThreadPatchStack();
 }
 
 s32 Linker::LoadModule(const std::filesystem::path& elf_name, bool is_dynamic) {