Use the default heap for threads, make sure to memset default heap allocations to 0 after using them

Dockerfile

@@ -1,4 +1,4 @@
-FROM ghcr.io/wiiu-env/devkitppc:20221228
+FROM ghcr.io/wiiu-env/devkitppc:20230326
 
 COPY --from=ghcr.io/wiiu-env/libkernel:20220904 /artifacts $DEVKITPRO
 COPY --from=ghcr.io/wiiu-env/libfunctionpatcher:20230106 /artifacts $DEVKITPRO

source/CThread.h

@@ -16,8 +16,10 @@
  ****************************************************************************/
 #pragma once
 
+#include "globals.h"
 #include <coreinit/thread.h>
 #include <cstdint>
+#include <cstring>
 #include <malloc.h>
 #include <unistd.h>
 
@@ -26,14 +28,15 @@ public:
     typedef void (*Callback)(CThread *thread, void *arg);
 
     //! constructor
-    explicit CThread(int32_t iAttr, int32_t iPriority = 16, int32_t iStackSize = 0x8000, CThread::Callback callback = nullptr, void *callbackArg = nullptr)
+    explicit CThread(int32_t iAttr, int32_t iPriority = 16, int32_t stacksize = 0x8000, CThread::Callback callback = nullptr, void *callbackArg = nullptr)
         : pThread(nullptr), pThreadStack(nullptr), pCallback(callback), pCallbackArg(callbackArg) {
         //! save attribute assignment
         iAttributes = iAttr;
-        //! allocate the thread
-        pThread = (OSThread *) memalign(8, 0x1000);
-        //! allocate the stack
-        pThreadStack = (uint8_t *) memalign(0x20, iStackSize);
+        iStackSize  = stacksize;
+        //! allocate the thread on the default Cafe OS heap
+        pThread = (OSThread *) gMEMAllocFromDefaultHeapExForThreads(sizeof(OSThread), 0x10);
+        //! allocate the stack on the default Cafe OS heap
+        pThreadStack = (uint8_t *) gMEMAllocFromDefaultHeapExForThreads(iStackSize, 0x20);
         //! create the thread
         if (pThread && pThreadStack) {
             // clang-format off
@@ -71,7 +74,9 @@ public:
     //! Resume thread
     virtual void resumeThread() {
         if (!isThreadSuspended()) return;
-        if (pThread) OSResumeThread(pThread);
+        if (pThread) {
+            OSResumeThread(pThread);
+        }
     }
 
     //! Set thread priority
@@ -111,12 +116,16 @@ public:
             }
             OSJoinThread(pThread, nullptr);
         }
+        // Some games (e.g. Minecraft) expect the default heap to be empty.
+        // Make sure to clean up the memory after using it
         //! free the thread stack buffer
         if (pThreadStack) {
-            free(pThreadStack);
+            memset(pThreadStack, 0, iStackSize);
+            gMEMFreeToDefaultHeapForThreads(pThreadStack);
         }
         if (pThread) {
-            free(pThread);
+            memset(pThread, 0, sizeof(OSThread));
+            gMEMFreeToDefaultHeapForThreads(pThread);
         }
         pThread      = nullptr;
         pThreadStack = nullptr;
@@ -139,6 +148,7 @@ private:
         return 0;
     }
 
+    uint32_t iStackSize;
     int32_t iAttributes;
     OSThread *pThread;
     uint8_t *pThreadStack;

source/globals.c

@@ -0,0 +1,4 @@
+#include "globals.h"
+
+void *(*gMEMAllocFromDefaultHeapExForThreads)(uint32_t size, int align);
+void (*gMEMFreeToDefaultHeapForThreads)(void *ptr);

source/globals.h

@@ -0,0 +1,5 @@
+#pragma once
+#include <stdint.h>
+
+extern void *(*gMEMAllocFromDefaultHeapExForThreads)(uint32_t size, int align);
+extern void (*gMEMFreeToDefaultHeapForThreads)(void *ptr);

source/main.cpp

@@ -1,4 +1,6 @@
 #include "function_replacements.h"
+#include "globals.h"
+#include "logger.h"
 #include "memory_mapping.h"
 #include "version.h"
 #include <coreinit/debug.h>
@@ -16,15 +18,45 @@ WUMS_MODULE_INIT_BEFORE_RELOCATION_DONE_HOOK();
 WUMS_DEPENDS_ON(homebrew_kernel);
 WUMS_DEPENDS_ON(homebrew_functionpatcher);
 
+#include <coreinit/dynload.h>
+#include <coreinit/memdefaultheap.h>
+
 // We can't use the functions from libfunctionpatcher. Defined in functionpatcher.def
 extern "C" FunctionPatcherStatus FPAddFunctionPatch(function_replacement_data_t *function_data, PatchedFunctionHandle *outHandle, bool *outHasBeenPatched);
 
+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;
+    if (OSDynLoad_Acquire("coreinit", &coreinitModule) != OS_DYNLOAD_OK) {
+        DEBUG_FUNCTION_LINE_ERR("Failed to acquire coreinit.rpl");
+        OSFatal("FunctionPatcherModule: Failed to acquire coreinit.rpl");
+    }
+    /* Memory allocation functions */
+    uint32_t *allocPtr, *freePtr;
+    if (OSDynLoad_FindExport(coreinitModule, OS_DYNLOAD_EXPORT_DATA, "MEMAllocFromDefaultHeapEx", reinterpret_cast<void **>(&allocPtr)) != OS_DYNLOAD_OK) {
+        DEBUG_FUNCTION_LINE_ERR("OSDynLoad_FindExport for MEMAllocFromDefaultHeapEx");
+        OSFatal("MemoryMappingModule: OSDynLoad_FindExport for MEMAllocFromDefaultHeapEx");
+    }
+    if (OSDynLoad_FindExport(coreinitModule, OS_DYNLOAD_EXPORT_DATA, "MEMFreeToDefaultHeap", reinterpret_cast<void **>(&freePtr)) != OS_DYNLOAD_OK) {
+        DEBUG_FUNCTION_LINE_ERR("OSDynLoad_FindExport for MEMFreeToDefaultHeap");
+        OSFatal("MemoryMappingModule: OSDynLoad_FindExport for MEMFreeToDefaultHeap");
+    }
+
+    gMEMAllocFromDefaultHeapExForThreads = (void *(*) (uint32_t, int) ) * allocPtr;
+    gMEMFreeToDefaultHeapForThreads      = (void (*)(void *)) * freePtr;
+
+    OSDynLoad_Release(coreinitModule);
+}
+
 WUMS_INITIALIZE(args) {
     static uint8_t ucSetupRequired = 1;
     if (!ucSetupRequired) {
         return;
     }
 
+    UpdateFunctionPointer();
+
     ucSetupRequired = 0;
     MemoryMapping_setupMemoryMapping();
     MemoryMapping_CreateHeaps();
@@ -45,6 +77,11 @@ WUMS_INITIALIZE(args) {
 
 WUMS_APPLICATION_STARTS() {
     OSReport("Running MemoryMappingModule " VERSION VERSION_EXTRA "\n");
+
+    // Now we can update the pointer with the "real" functions
+    gMEMAllocFromDefaultHeapExForThreads = MEMAllocFromDefaultHeapEx;
+    gMEMFreeToDefaultHeapForThreads      = MEMFreeToDefaultHeap;
+
 #ifdef DEBUG
     initLogging();
 #endif