Update the malloc/free pointers used for threads to fix apps which use a custom heap

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);

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;

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);
+void *(*gMEMAllocFromDefaultHeapExForThreads)(uint32_t size, int align);
+void (*gMEMFreeToDefaultHeapForThreads)(void *ptr);

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);
+extern void *(*gMEMAllocFromDefaultHeapExForThreads)(uint32_t size, int align);
+extern void (*gMEMFreeToDefaultHeapForThreads)(void *ptr);