From 2882635f6f9a8578b80bd35428e5ac9270ae6fa1 Mon Sep 17 00:00:00 2001
From: Maschell <Maschell@gmx.de>
Date: Mon, 3 Oct 2022 21:51:12 +0200
Subject: [PATCH] Set a custom OSDynLoad Allocator when processing relocations

---
 source/PluginManagement.cpp | 12 +++++++++++-
 source/utils/utils.cpp      | 25 ++++++++++++++++++++++++-
 source/utils/utils.h        |  7 ++++++-
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/source/PluginManagement.cpp b/source/PluginManagement.cpp
index ad70c2e..cc1aead 100644
--- a/source/PluginManagement.cpp
+++ b/source/PluginManagement.cpp
@@ -52,6 +52,7 @@ bool PluginManagement::doRelocation(const std::vector<std::unique_ptr<Relocation
         } else {
             //DEBUG_FUNCTION_LINE("Found export for %s %s", rplName.c_str(), functionName.c_str());
         }
+
         if (!ElfUtils::elfLinkOne(cur->getType(), cur->getOffset(), cur->getAddend(), (uint32_t) cur->getDestination(), functionAddress, tramp_data, tramp_length, RELOC_TYPE_IMPORT, trampolineID)) {
             DEBUG_FUNCTION_LINE_ERR("elfLinkOne failed");
             return false;
@@ -73,19 +74,28 @@ bool PluginManagement::doRelocation(const std::vector<std::unique_ptr<Relocation
     return true;
 }
 
-
 bool PluginManagement::doRelocations(const std::vector<std::unique_ptr<PluginContainer>> &plugins, relocation_trampoline_entry_t *trampData, uint32_t tramp_size) {
     for (uint32_t i = 0; i < tramp_size; i++) {
         if (trampData[i].status == RELOC_TRAMP_IMPORT_DONE) {
             trampData[i].status = RELOC_TRAMP_FREE;
         }
     }
+
+    OSDynLoadAllocFn prevDynLoadAlloc = nullptr;
+    OSDynLoadFreeFn prevDynLoadFree   = nullptr;
+
+    OSDynLoad_GetAllocator(&prevDynLoadAlloc, &prevDynLoadFree);
+    OSDynLoad_SetAllocator(CustomDynLoadAlloc, CustomDynLoadFree);
+
     for (auto &pluginContainer : plugins) {
         DEBUG_FUNCTION_LINE_VERBOSE("Doing relocations for plugin: %s", pluginContainer->getMetaInformation()->getName().c_str());
         if (!PluginManagement::doRelocation(pluginContainer->getPluginInformation()->getRelocationDataList(), trampData, tramp_size, pluginContainer->getPluginInformation()->getTrampolineId())) {
             return false;
         }
     }
+
+    OSDynLoad_SetAllocator(prevDynLoadAlloc, prevDynLoadFree);
+
     return true;
 }
 
diff --git a/source/utils/utils.cpp b/source/utils/utils.cpp
index 33e87fd..9a8311b 100644
--- a/source/utils/utils.cpp
+++ b/source/utils/utils.cpp
@@ -1,4 +1,5 @@
-#include "utils/logger.h"
+#include "utils.h"
+#include "logger.h"
 #include <coreinit/ios.h>
 #include <cstring>
 #include <malloc.h>
@@ -53,3 +54,25 @@ void dumpHex(const void *data, size_t size) {
         }
     }
 }
+
+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;
+}
+
+void CustomDynLoadFree(void *addr) {
+    free(addr);
+}
\ No newline at end of file
diff --git a/source/utils/utils.h b/source/utils/utils.h
index 46e7f21..26a3bb2 100644
--- a/source/utils/utils.h
+++ b/source/utils/utils.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <coreinit/dynload.h>
 #include <cstdint>
 #include <forward_list>
 #include <malloc.h>
@@ -73,4 +74,8 @@ bool remove_locked_first_if(std::mutex &mutex, std::forward_list<T, Allocator> &
     return false;
 }
 
-std::string getPluginPath();
\ No newline at end of file
+std::string getPluginPath();
+
+OSDynLoad_Error CustomDynLoadAlloc(int32_t size, int32_t align, void **outAddr);
+
+void CustomDynLoadFree(void *addr);
\ No newline at end of file