From d602d170344e0d58243b0eb1c6d767dd443fd9e9 Mon Sep 17 00:00:00 2001
From: Maschell <Maschell@gmx.de>
Date: Tue, 26 Jul 2022 23:11:47 +0200
Subject: [PATCH] Update the malloc/free pointers used for threads to fix apps
 which use a custom heap

---
 source/main.cpp          | 17 ++++++++++++++---
 source/utils/CThread.h   |  8 ++++----
 source/utils/globals.cpp |  4 ++--
 source/utils/globals.h   |  4 ++--
 4 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/source/main.cpp b/source/main.cpp
index 1eec86d..90c0216 100644
--- a/source/main.cpp
+++ b/source/main.cpp
@@ -4,13 +4,15 @@
 #include "utils/globals.h"
 #include "utils/logger.h"
 #include "utils/utils.h"
+#include <coreinit/memdefaultheap.h>
 #include <coreinit/memexpheap.h>
 #include <wums.h>
 
 WUMS_MODULE_EXPORT_NAME("homebrew_functionpatcher");
 WUMS_MODULE_INIT_BEFORE_RELOCATION_DONE_HOOK();
 
-WUMS_INITIALIZE() {
+
+void UpdateFunctionPointer() {
     // We need the real MEMAllocFromDefaultHeapEx/MEMFreeToDefaultHeap function pointer to force-allocate memory on the default heap.
     // Our custom heap doesn't work (yet) for threads and causes an app panic.
     OSDynLoad_Module coreinitModule;
@@ -30,9 +32,14 @@ WUMS_INITIALIZE() {
         OSFatal("OSDynLoad_FindExport for MEMFreeToDefaultHeap");
     }
 
-    gRealMEMAllocFromDefaultHeapEx = (void *(*) (uint32_t, int) ) * allocPtr;
-    gMEMFreeToDefaultHeap          = (void (*)(void *)) * freePtr;
+    gMEMAllocFromDefaultHeapExForThreads = (void *(*) (uint32_t, int) ) * allocPtr;
+    gMEMFreeToDefaultHeapForThreads      = (void (*)(void *)) * freePtr;
+
     OSDynLoad_Release(coreinitModule);
+}
+
+WUMS_INITIALIZE() {
+    UpdateFunctionPointer();
 
     memset(gJumpHeapData, 0, JUMP_HEAP_DATA_SIZE);
     gJumpHeapHandle = MEMCreateExpHeapEx((void *) (gJumpHeapData), JUMP_HEAP_DATA_SIZE, 1);
@@ -66,6 +73,10 @@ WUMS_APPLICATION_STARTS() {
         return;
     }
 
+    // Now we can update the pointer with the "real" functions
+    gMEMAllocFromDefaultHeapExForThreads = MEMAllocFromDefaultHeapEx;
+    gMEMFreeToDefaultHeapForThreads      = MEMFreeToDefaultHeap;
+
     initLogging();
 
     std::lock_guard<std::mutex> lock(gPatchedFunctionsMutex);
diff --git a/source/utils/CThread.h b/source/utils/CThread.h
index 6cf1dc6..910b1fd 100644
--- a/source/utils/CThread.h
+++ b/source/utils/CThread.h
@@ -33,9 +33,9 @@ public:
         //! save attribute assignment
         iAttributes = iAttr;
         //! allocate the thread on the default Cafe OS heap
-        pThread = (OSThread *) gRealMEMAllocFromDefaultHeapEx(sizeof(OSThread), 0x10);
+        pThread = (OSThread *) gMEMAllocFromDefaultHeapExForThreads(sizeof(OSThread), 0x10);
         //! allocate the stack on the default Cafe OS heap
-        pThreadStack = (uint8_t *) gRealMEMAllocFromDefaultHeapEx(iStackSize, 0x20);
+        pThreadStack = (uint8_t *) gMEMAllocFromDefaultHeapExForThreads(iStackSize, 0x20);
         //! create the thread
         if (pThread && pThreadStack) {
             // clang-format off
@@ -125,10 +125,10 @@ public:
         }
         //! free the thread stack buffer
         if (pThreadStack) {
-            gMEMFreeToDefaultHeap(pThreadStack);
+            gMEMFreeToDefaultHeapForThreads(pThreadStack);
         }
         if (pThread) {
-            gMEMFreeToDefaultHeap(pThread);
+            gMEMFreeToDefaultHeapForThreads(pThread);
         }
         pThread      = nullptr;
         pThreadStack = nullptr;
diff --git a/source/utils/globals.cpp b/source/utils/globals.cpp
index 409dfc7..d5a6619 100644
--- a/source/utils/globals.cpp
+++ b/source/utils/globals.cpp
@@ -7,5 +7,5 @@ std::shared_ptr<FunctionAddressProvider> gFunctionAddressProvider;
 std::mutex gPatchedFunctionsMutex;
 std::vector<std::shared_ptr<PatchedFunctionData>> gPatchedFunctions;
 
-void *(*gRealMEMAllocFromDefaultHeapEx)(uint32_t size, int align);
-void (*gMEMFreeToDefaultHeap)(void *ptr);
\ No newline at end of file
+void *(*gMEMAllocFromDefaultHeapExForThreads)(uint32_t size, int align);
+void (*gMEMFreeToDefaultHeapForThreads)(void *ptr);
\ No newline at end of file
diff --git a/source/utils/globals.h b/source/utils/globals.h
index 6642b18..df7d653 100644
--- a/source/utils/globals.h
+++ b/source/utils/globals.h
@@ -12,5 +12,5 @@ extern std::shared_ptr<FunctionAddressProvider> gFunctionAddressProvider;
 extern std::mutex gPatchedFunctionsMutex;
 extern std::vector<std::shared_ptr<PatchedFunctionData>> gPatchedFunctions;
 
-extern void *(*gRealMEMAllocFromDefaultHeapEx)(uint32_t size, int align);
-extern void (*gMEMFreeToDefaultHeap)(void *ptr);
\ No newline at end of file
+extern void *(*gMEMAllocFromDefaultHeapExForThreads)(uint32_t size, int align);
+extern void (*gMEMFreeToDefaultHeapForThreads)(void *ptr);
\ No newline at end of file