Keep track of acquired RPLs and memory, release/free them if needed

Dockerfile

@@ -1,5 +1,5 @@
-FROM wiiuenv/devkitppc:20220806
+FROM wiiuenv/devkitppc:20220917
 
-COPY --from=wiiuenv/wiiumodulesystem:20220904 /artifacts $DEVKITPRO
+COPY --from=wiiuenv/wiiumodulesystem:20221005 /artifacts $DEVKITPRO
 
 WORKDIR project

wumsloader/src/entry.cpp

@@ -40,10 +40,37 @@ extern "C" int _start(int argc, char **argv) {
     return ((int (*)(int, char **))(*(unsigned int *) 0x1005E040))(argc, argv);
 }
 
+void SaveLoadedRPLsInGlobalInformation(module_information_t *globalInformation,
+                                       std::vector<OSDynLoad_Module> loadedRPLs) {
+    // free previous allocations.
+    if (globalInformation->acquired_rpls) {
+        free(globalInformation->acquired_rpls);
+    }
+    globalInformation->number_acquired_rpls = loadedRPLs.size();
+    globalInformation->acquired_rpls        = (uint32_t *) malloc(loadedRPLs.size() * sizeof(uint32_t));
+    if (!globalInformation->acquired_rpls) {
+        OSFatal("Failed to allocate memory");
+    }
+    uint32_t i = 0;
+    for (auto &rpl : loadedRPLs) {
+        globalInformation->acquired_rpls[i] = (uint32_t) rpl;
+        ++i;
+    }
+}
+
 void doStart(int argc, char **argv) {
     init_wut();
     initLogging();
 
+    gLoadedRPLs.clear();
+    // If an allocated rpl was not released properly (e.g. if something else calls OSDynload_Acquire without releasing it)
+    // memory gets leaked. Let's clean this up!
+    for (auto &addr : gAllocatedAddresses) {
+        DEBUG_FUNCTION_LINE_WARN("Memory allocated by OSDynload was not freed properly, let's clean it up! (%08X)", addr);
+        free((void *) addr);
+    }
+    gAllocatedAddresses.clear();
+
     if (!gInitCalled) {
         gInitCalled = 1;
 
@@ -108,7 +135,7 @@ void doStart(int argc, char **argv) {
         }
 
         DEBUG_FUNCTION_LINE_VERBOSE("Resolve relocations without replacing alloc functions");
-        ResolveRelocations(gLoadedModules, true);
+        ResolveRelocations(gLoadedModules, true, gLoadedRPLs);
 
         for (auto &curModule : gLoadedModules) {
             if (curModule->isInitBeforeRelocationDoneHook()) {
@@ -126,10 +153,12 @@ void doStart(int argc, char **argv) {
         }
     } else {
         DEBUG_FUNCTION_LINE("Resolve relocations and replace alloc functions");
-        ResolveRelocations(gLoadedModules, false);
+        ResolveRelocations(gLoadedModules, false, gLoadedRPLs);
         CallHook(gLoadedModules, WUMS_HOOK_RELOCATIONS_DONE);
     }
 
+    SaveLoadedRPLsInGlobalInformation(&gModuleInformation, gLoadedRPLs);
+
     CallHook(gLoadedModules, WUMS_HOOK_INIT_WUT_DEVOPTAB);
     CallHook(gLoadedModules, WUMS_HOOK_INIT_WUT_SOCKETS);
     CallHook(gLoadedModules, WUMS_HOOK_APPLICATION_STARTS);

wumsloader/src/globals.cpp

@@ -5,3 +5,6 @@ uint8_t gInitCalled __attribute__((section(".data")))       = 0;
 module_information_t gModuleInformation __attribute__((section(".data")));
 std::vector<std::shared_ptr<ModuleData>> gLoadedModules __attribute__((section(".data")));
 std::unique_ptr<module_information_single_t[]> gModuleDataInfo __attribute__((section(".data")));
+
+std::vector<OSDynLoad_Module> gLoadedRPLs __attribute__((section(".data")));
+std::vector<void *> gAllocatedAddresses __attribute__((section(".data")));

wumsloader/src/globals.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "module/ModuleData.h"
+#include <coreinit/dynload.h>
 #include <coreinit/memheap.h>
 #include <cstdint>
 #include <memory>
@@ -10,6 +11,8 @@ extern MEMHeapHandle gHeapHandle;
 extern module_information_t gModuleInformation;
 extern std::vector<std::shared_ptr<ModuleData>> gLoadedModules;
 extern std::unique_ptr<module_information_single_t[]> gModuleDataInfo;
+extern std::vector<OSDynLoad_Module> gLoadedRPLs;
+extern std::vector<void *> gAllocatedAddresses;
 
 #define MEMORY_REGION_START                   0x00800000
 #define MEMORY_REGION_SIZE                    0x00800000

wumsloader/src/utils/RelocationUtils.cpp

@@ -1,10 +1,13 @@
 #include "RelocationUtils.h"
+#include "ElfUtils.h"
 #include "globals.h"
-#include "malloc.h"
-#include "utils/ElfUtils.h"
-#include "utils/memory.h"
+#include "memory.h"
+#include <algorithm>
 #include <coreinit/cache.h>
-#include <coreinit/dynload.h>
+#include <iostream>
+#include <iterator>
+#include <malloc.h>
+#include <vector>
 
 static OSDynLoad_Error CustomDynLoadAlloc(int32_t size, int32_t align, void **outAddr) {
     if (!outAddr) {
@@ -21,14 +24,23 @@ static OSDynLoad_Error CustomDynLoadAlloc(int32_t size, int32_t align, void **ou
         return OS_DYNLOAD_OUT_OF_MEMORY;
     }
 
+    // keep track of allocated memory to clean it up in case the RPLs won't get unloaded properly
+    gAllocatedAddresses.push_back(*outAddr);
+
     return OS_DYNLOAD_OK;
 }
 
 static void CustomDynLoadFree(void *addr) {
     free(addr);
+
+    // Remove from list
+    auto it = std::find(gAllocatedAddresses.begin(), gAllocatedAddresses.end(), addr);
+    if (it != gAllocatedAddresses.end()) {
+        gAllocatedAddresses.erase(it);
+    }
 }
 
-bool ResolveRelocations(std::vector<std::shared_ptr<ModuleData>> &loadedModules, bool skipMemoryMappingModule) {
+bool ResolveRelocations(std::vector<std::shared_ptr<ModuleData>> &loadedModules, bool skipMemoryMappingModule, std::vector<OSDynLoad_Module> &loadedRPLs) {
     bool wasSuccessful = true;
 
     OSDynLoadAllocFn prevDynLoadAlloc = nullptr;
@@ -48,10 +60,10 @@ bool ResolveRelocations(std::vector<std::shared_ptr<ModuleData>> &loadedModules,
         // Afterwards we can just rely on the custom heap.
         bool skipAllocFunction = skipMemoryMappingModule && (std::string_view(curModule->getExportName()) == "homebrew_memorymapping");
         DEBUG_FUNCTION_LINE_VERBOSE("Skip alloc replace? %d", skipAllocFunction);
-        if (!doRelocation(gLoadedModules, relocData, nullptr, 0, skipAllocFunction)) {
+        if (!doRelocation(gLoadedModules, relocData, nullptr, 0, skipAllocFunction, loadedRPLs)) {
             wasSuccessful = false;
             DEBUG_FUNCTION_LINE_ERR("Failed to do Relocations for %s", curModule->getExportName().c_str());
-            OSFatal("Failed to do Reloations");
+            OSFatal("Failed to do Relocations");
         }
     }
     OSDynLoad_SetAllocator(prevDynLoadAlloc, prevDynLoadFree);
@@ -66,7 +78,8 @@ bool doRelocation(const std::vector<std::shared_ptr<ModuleData>> &moduleList,
                   const std::vector<std::unique_ptr<RelocationData>> &relocData,
                   relocation_trampoline_entry_t *tramp_data,
                   uint32_t tramp_length,
-                  bool skipAllocReplacement) {
+                  bool skipAllocReplacement,
+                  std::vector<OSDynLoad_Module> &loadedRPLs) {
     std::map<std::string, OSDynLoad_Module> moduleCache;
     for (auto const &curReloc : relocData) {
         auto &functionName       = curReloc->getName();
@@ -106,6 +119,8 @@ bool doRelocation(const std::vector<std::shared_ptr<ModuleData>> &moduleList,
                         DEBUG_FUNCTION_LINE_ERR("Failed to acquire %s", rplName.c_str());
                         return false;
                     }
+                    // Keep track RPLs we have acquired, they will be released on exit (see: AromaBaseModule)
+                    loadedRPLs.push_back(rplHandle);
                 }
                 moduleCache[rplName] = rplHandle;
             }

wumsloader/src/utils/RelocationUtils.h

@@ -1,12 +1,15 @@
 #pragma once
 #include "module/ModuleData.h"
+#include <coreinit/dynload.h>
 #include <vector>
 
 bool ResolveRelocations(std::vector<std::shared_ptr<ModuleData>> &loadedModules,
-                        bool skipMemoryMappingModule);
+                        bool skipMemoryMappingModule,
+                        std::vector<OSDynLoad_Module> &loadedRPLs);
 
 bool doRelocation(const std::vector<std::shared_ptr<ModuleData>> &moduleList,
                   const std::vector<std::unique_ptr<RelocationData>> &relocData,
                   relocation_trampoline_entry_t *tramp_data,
                   uint32_t tramp_length,
-                  bool skipAllocReplacement);
+                  bool skipAllocReplacement,
+                  std::vector<OSDynLoad_Module> &loadedRPLs);