Keep track of acquired RPLs and release them on exit

source/PluginManagement.cpp

@@ -13,8 +13,11 @@
 #include <memory>
 #include <ranges>
 
-bool PluginManagement::doRelocation(const std::vector<std::unique_ptr<RelocationData>> &relocData, relocation_trampoline_entry_t *tramp_data, uint32_t tramp_length, uint32_t trampolineID) {
+bool PluginManagement::doRelocation(const std::vector<std::unique_ptr<RelocationData>> &relocData,
-    std::map<std::string, OSDynLoad_Module> moduleHandleCache;
+                                    relocation_trampoline_entry_t *tramp_data,
+                                    uint32_t tramp_length,
+                                    uint32_t trampolineID,
+                                    std::map<std::string, OSDynLoad_Module> &usedRPls) {
     for (auto const &cur : relocData) {
         uint32_t functionAddress        = 0;
         const std::string &functionName = cur->getName();
@@ -37,12 +40,22 @@ bool PluginManagement::doRelocation(const std::vector<std::unique_ptr<Relocation
             auto rplName               = cur->getImportRPLInformation()->getRPLName();
             int32_t isData             = cur->getImportRPLInformation()->isData();
             OSDynLoad_Module rplHandle = nullptr;
-            if (moduleHandleCache.count(rplName) > 0) {
+
-                rplHandle = moduleHandleCache[rplName];
+            if (!usedRPls.contains(rplName)) {
+                DEBUG_FUNCTION_LINE_VERBOSE("Acquire %s", rplName.c_str());
+                // Always acquire to increase refcount and make sure it won't get unloaded while we're using it.
+                OSDynLoad_Error err = OSDynLoad_Acquire(rplName.c_str(), &rplHandle);
+                if (err != OS_DYNLOAD_OK) {
+                    DEBUG_FUNCTION_LINE_ERR("Failed to acquire %s", rplName.c_str());
+                    return false;
+                }
+                // Keep track RPLs we are using.
+                // They will be released on exit
+                usedRPls[rplName] = rplHandle;
             } else {
-                OSDynLoad_Acquire(rplName.c_str(), &rplHandle);
+                rplHandle = usedRPls[rplName];
-                moduleHandleCache[rplName] = rplHandle;
             }
+
             OSDynLoad_FindExport(rplHandle, isData, functionName.c_str(), (void **) &functionAddress);
         }
 
@@ -74,7 +87,9 @@ 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) {
+bool PluginManagement::doRelocations(const std::vector<std::unique_ptr<PluginContainer>> &plugins,
+                                     relocation_trampoline_entry_t *trampData, uint32_t tramp_size,
+                                     std::map<std::string, OSDynLoad_Module> &usedRPls) {
     for (uint32_t i = 0; i < tramp_size; i++) {
         if (trampData[i].status == RELOC_TRAMP_IMPORT_DONE) {
             trampData[i].status = RELOC_TRAMP_FREE;
@@ -89,7 +104,11 @@ bool PluginManagement::doRelocations(const std::vector<std::unique_ptr<PluginCon
 
     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())) {
+        if (!PluginManagement::doRelocation(pluginContainer->getPluginInformation()->getRelocationDataList(),
+                                            trampData,
+                                            tramp_size,
+                                            pluginContainer->getPluginInformation()->getTrampolineId(),
+                                            usedRPls)) {
             return false;
         }
     }

source/PluginManagement.h

@@ -1,9 +1,10 @@
 #pragma once
 
 #include "plugin/PluginContainer.h"
+#include <coreinit/dynload.h>
 #include <forward_list>
+#include <map>
 #include <memory>
-#include <vector>
 #include <wums/defines/relocation_defines.h>
 
 class PluginManagement {
@@ -12,9 +13,16 @@ public:
 
     static void callInitHooks(const std::vector<std::unique_ptr<PluginContainer>> &plugins);
 
-    static bool doRelocations(const std::vector<std::unique_ptr<PluginContainer>> &plugins, relocation_trampoline_entry_t *trampData, uint32_t tramp_size);
+    static bool doRelocations(const std::vector<std::unique_ptr<PluginContainer>> &plugins,
+                              relocation_trampoline_entry_t *trampData,
+                              uint32_t tramp_size,
+                              std::map<std::string, OSDynLoad_Module> &usedRPls);
 
-    static bool doRelocation(const std::vector<std::unique_ptr<RelocationData>> &relocData, relocation_trampoline_entry_t *tramp_data, uint32_t tramp_length, uint32_t trampolineID);
+    static bool doRelocation(const std::vector<std::unique_ptr<RelocationData>> &relocData,
+                             relocation_trampoline_entry_t *tramp_data,
+                             uint32_t tramp_length,
+                             uint32_t trampolineID,
+                             std::map<std::string, OSDynLoad_Module> &usedRPls);
 
     static bool DoFunctionPatches(const std::vector<std::unique_ptr<PluginContainer>> &plugins);
 

source/globals.cpp

@@ -9,3 +9,5 @@ relocation_trampoline_entry_t *gTrampData __attribute__((section(".data"))) = nu
 std::forward_list<std::shared_ptr<PluginData>> gLoadedData __attribute__((section(".data")));
 std::forward_list<std::shared_ptr<PluginData>> gLoadOnNextLaunch __attribute__((section(".data")));
 std::mutex gLoadedDataMutex __attribute__((section(".data")));
+std::map<std::string, OSDynLoad_Module> gUsedRPLs __attribute__((section(".data")));
+std::vector<void *> gAllocatedAddresses __attribute__((section(".data")));

source/globals.h

@@ -2,6 +2,7 @@
 #include "plugin/PluginContainer.h"
 #include "utils/ConfigUtils.h"
 #include "version.h"
+#include <coreinit/dynload.h>
 #include <forward_list>
 #include <memory>
 #include <mutex>
@@ -21,3 +22,5 @@ extern std::vector<std::unique_ptr<PluginContainer>> gLoadedPlugins;
 extern std::forward_list<std::shared_ptr<PluginData>> gLoadedData;
 extern std::forward_list<std::shared_ptr<PluginData>> gLoadOnNextLaunch;
 extern std::mutex gLoadedDataMutex;
+extern std::map<std::string, OSDynLoad_Module> gUsedRPLs;
+extern std::vector<void *> gAllocatedAddresses;

source/main.cpp

@@ -39,6 +39,11 @@ WUMS_APPLICATION_ENDS() {
     CallHook(gLoadedPlugins, WUPS_LOADER_HOOK_FINI_WUT_SOCKETS);
     CallHook(gLoadedPlugins, WUPS_LOADER_HOOK_FINI_WUT_DEVOPTAB);
 
+    for (auto &pair : gUsedRPLs) {
+        OSDynLoad_Release(pair.second);
+    }
+    gUsedRPLs.clear();
+
     deinitLogging();
 }
 
@@ -50,6 +55,17 @@ WUMS_APPLICATION_STARTS() {
     }
 
     OSReport("Running WiiUPluginLoaderBackend " VERSION_FULL "\n");
+
+    gUsedRPLs.clear();
+
+    // If an allocated rpl was not released properly (e.g. if something else calls OSDynload_Acquire without releasing it) memory get 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();
+
     initLogging();
     bool initNeeded = false;
 
@@ -93,7 +109,7 @@ WUMS_APPLICATION_STARTS() {
     gLoadedData.clear();
 
     if (!gLoadedPlugins.empty()) {
-        if (!PluginManagement::doRelocations(gLoadedPlugins, gTrampData, TRAMP_DATA_SIZE)) {
+        if (!PluginManagement::doRelocations(gLoadedPlugins, gTrampData, TRAMP_DATA_SIZE, gUsedRPLs)) {
             DEBUG_FUNCTION_LINE_ERR("Relocations failed");
             OSFatal("Relocations failed");
         }

source/utils/utils.cpp

@@ -1,8 +1,8 @@
 #include "utils.h"
+#include "globals.h"
 #include "logger.h"
+#include <algorithm>
 #include <coreinit/ios.h>
-#include <cstring>
-#include <malloc.h>
 #include <string>
 
 std::string getPluginPath() {
@@ -69,10 +69,18 @@ OSDynLoad_Error CustomDynLoadAlloc(int32_t size, int32_t align, void **outAddr)
     if (!(*outAddr = memalign(align, size))) {
         return OS_DYNLOAD_OUT_OF_MEMORY;
     }
+    // keep track of allocated memory to clean it up if RPLs won't get unloaded properly
+    gAllocatedAddresses.push_back(*outAddr);
 
     return OS_DYNLOAD_OK;
 }
 
 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);
+    }
 }