From 6555a8138f2bd5fb5593837ce97c30e3bf43391a Mon Sep 17 00:00:00 2001
From: Maschell <Maschell@gmx.de>
Date: Mon, 3 Oct 2022 22:04:38 +0200
Subject: [PATCH] Set a custom OSDynLoad Allocator when processing relocations

---
 wumsloader/src/utils/RelocationUtils.cpp | 31 ++++++++++++++++++++++++
 wumsloader/src/utils/imports.h           |  2 ++
 2 files changed, 33 insertions(+)

diff --git a/wumsloader/src/utils/RelocationUtils.cpp b/wumsloader/src/utils/RelocationUtils.cpp
index 62cf9f7..3518a75 100644
--- a/wumsloader/src/utils/RelocationUtils.cpp
+++ b/wumsloader/src/utils/RelocationUtils.cpp
@@ -1,13 +1,42 @@
 #include "RelocationUtils.h"
 #include "globals.h"
+#include "malloc.h"
 #include "utils/ElfUtils.h"
 #include "utils/memory.h"
 #include <coreinit/cache.h>
 #include <coreinit/dynload.h>
 
+static OSDynLoad_Error CustomDynLoadAlloc(int32_t size, int32_t align, void **outAddr) {
+    if (!outAddr) {
+        return OS_DYNLOAD_INVALID_ALLOCATOR_PTR;
+    }
+
+    if (align >= 0 && align < 4) {
+        align = 4;
+    } else if (align < 0 && align > -4) {
+        align = -4;
+    }
+
+    if (!(*outAddr = memalign(align, size))) {
+        return OS_DYNLOAD_OUT_OF_MEMORY;
+    }
+
+    return OS_DYNLOAD_OK;
+}
+
+static void CustomDynLoadFree(void *addr) {
+    free(addr);
+}
+
 bool ResolveRelocations(std::vector<std::shared_ptr<ModuleData>> &loadedModules, bool skipMemoryMappingModule) {
     bool wasSuccessful = true;
 
+    OSDynLoadAllocFn prevDynLoadAlloc = nullptr;
+    OSDynLoadFreeFn prevDynLoadFree   = nullptr;
+
+    OSDynLoad_GetAllocator(&prevDynLoadAlloc, &prevDynLoadFree);
+    OSDynLoad_SetAllocator(CustomDynLoadAlloc, CustomDynLoadFree);
+
     for (auto &curModule : loadedModules) {
         DEBUG_FUNCTION_LINE("Let's do the relocations for %s", curModule->getExportName().c_str());
 
@@ -25,6 +54,8 @@ bool ResolveRelocations(std::vector<std::shared_ptr<ModuleData>> &loadedModules,
             OSFatal("Failed to do Reloations");
         }
     }
+    OSDynLoad_SetAllocator(prevDynLoadAlloc, prevDynLoadFree);
+
     DCFlushRange((void *) MEMORY_REGION_START, MEMORY_REGION_SIZE);
     ICInvalidateRange((void *) MEMORY_REGION_START, MEMORY_REGION_SIZE);
     return wasSuccessful;
diff --git a/wumsloader/src/utils/imports.h b/wumsloader/src/utils/imports.h
index 5f14528..89b2c2d 100644
--- a/wumsloader/src/utils/imports.h
+++ b/wumsloader/src/utils/imports.h
@@ -34,6 +34,8 @@ IMPORT(OSSignalCond);
 IMPORT(MEMAllocFromDefaultHeapEx);
 IMPORT(MEMFreeToDefaultHeap);
 IMPORT(OSSwapAtomic);
+IMPORT(OSDynLoad_GetAllocator);
+IMPORT(OSDynLoad_SetAllocator);
 
 IMPORT(FSTimeToCalendarTime);
 IMPORT(FSInit);