Initial support for tracking plugin heap usage and detect double frees or memory leaks

source/plugin/DefaultPluginHeapMemoryAllocator.cpp

@@ -0,0 +1,44 @@
+#include "DefaultPluginHeapMemoryAllocator.h"
+
+#include "utils/logger.h"
+
+#include <coreinit/dynload.h>
+
+DefaultPluginHeapMemoryAllocator DefaultPluginHeapMemoryAllocator::gDefaultPluginHeapMemoryAllocator;
+
+DefaultPluginHeapMemoryAllocator::DefaultPluginHeapMemoryAllocator() {
+    OSDynLoad_Module rplHandle;
+    if (OSDynLoad_Acquire("homebrew_memorymapping", &rplHandle) == OS_DYNLOAD_OK) {
+        OSDynLoad_FindExport(rplHandle, OS_DYNLOAD_EXPORT_DATA, "MEMAllocFromMappedMemory", reinterpret_cast<void **>(&mMemoryMappingModuleAllocPtr));
+        OSDynLoad_FindExport(rplHandle, OS_DYNLOAD_EXPORT_DATA, "MEMAllocFromMappedMemoryEx", reinterpret_cast<void **>(&mMemoryMappingModuleAllocExPtr));
+        OSDynLoad_FindExport(rplHandle, OS_DYNLOAD_EXPORT_DATA, "MEMFreeToMappedMemory", reinterpret_cast<void **>(&mMemoryMappingModuleFreePtr));
+    }
+}
+
+MEMAllocFromDefaultHeapFn *DefaultPluginHeapMemoryAllocator::GetAllocFunctionAddress() const {
+    if (mMemoryMappingModuleAllocPtr == nullptr) {
+        DEBUG_FUNCTION_LINE_WARN("### mMemoryMappingModuleAllocPtr was NULL!");
+        return &MEMAllocFromDefaultHeap;
+    }
+    return mMemoryMappingModuleAllocPtr;
+}
+
+MEMAllocFromDefaultHeapExFn *DefaultPluginHeapMemoryAllocator::GetAllocExFunctionAddress() const {
+    if (mMemoryMappingModuleAllocExPtr == nullptr) {
+        DEBUG_FUNCTION_LINE_WARN("### mMemoryMappingModuleAllocExPtr was NULL!");
+        return &MEMAllocFromDefaultHeapEx;
+    }
+    return mMemoryMappingModuleAllocExPtr;
+}
+
+MEMFreeToDefaultHeapFn *DefaultPluginHeapMemoryAllocator::GetFreeFunctionAddress() const {
+    if (mMemoryMappingModuleFreePtr == nullptr) {
+        DEBUG_FUNCTION_LINE_WARN("### mMemoryMappingModuleFreePtr was NULL!");
+        return &MEMFreeToDefaultHeap;
+    }
+    return mMemoryMappingModuleFreePtr;
+}
+
+std::optional<PluginMemorySnapshot> DefaultPluginHeapMemoryAllocator::GetHeapMemoryUsageSnapshot() const {
+    return {};
+}

source/plugin/DefaultPluginHeapMemoryAllocator.h

@@ -0,0 +1,20 @@
+#pragma once
+#include "IPluginHeapMemoryAllocator.h"
+
+class DefaultPluginHeapMemoryAllocator : public IPluginHeapMemoryAllocator {
+public:
+    ~DefaultPluginHeapMemoryAllocator() override = default;
+    MEMAllocFromDefaultHeapFn *GetAllocFunctionAddress() const override;
+    MEMAllocFromDefaultHeapExFn *GetAllocExFunctionAddress() const override;
+    MEMFreeToDefaultHeapFn *GetFreeFunctionAddress() const override;
+    std::optional<PluginMemorySnapshot> GetHeapMemoryUsageSnapshot() const override;
+
+    static DefaultPluginHeapMemoryAllocator gDefaultPluginHeapMemoryAllocator;
+
+private:
+    DefaultPluginHeapMemoryAllocator();
+
+    MEMAllocFromDefaultHeapFn *mMemoryMappingModuleAllocPtr     = nullptr;
+    MEMAllocFromDefaultHeapExFn *mMemoryMappingModuleAllocExPtr = nullptr;
+    MEMFreeToDefaultHeapFn *mMemoryMappingModuleFreePtr         = nullptr;
+};

source/plugin/IPluginHeapMemoryAllocator.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include <coreinit/memdefaultheap.h>
+#include <map>
+#include <mutex>
+#include <optional>
+#include <vector>
+
+#include <cstdint>
+
+struct AllocationInfo {
+    uint32_t size;
+    std::vector<uint32_t> stackTrace;
+};
+
+struct PluginMemorySnapshot {
+    uint32_t currentAllocated = 0;
+    uint32_t peakAllocated    = 0;
+    uint32_t allocCount       = 0;
+    uint32_t freeCount        = 0;
+    std::string pluginName    = {};
+
+    std::map<void *, AllocationInfo> allocationMap;
+};
+
+class IPluginHeapMemoryAllocator {
+public:
+    virtual ~IPluginHeapMemoryAllocator() = default;
+
+    /**
+     * The returned address needs to be valid for the whole time
+     */
+    virtual MEMAllocFromDefaultHeapFn *GetAllocFunctionAddress() const = 0;
+
+    /**
+     * The returned address needs to be valid for the whole time
+     */
+    virtual MEMAllocFromDefaultHeapExFn *GetAllocExFunctionAddress() const = 0;
+
+    /**
+     * The returned address needs to be valid for the whole time
+     */
+    virtual MEMFreeToDefaultHeapFn *GetFreeFunctionAddress() const = 0;
+
+    virtual std::optional<PluginMemorySnapshot> GetHeapMemoryUsageSnapshot() const = 0;
+};

source/plugin/PluginContainer.cpp

@@ -1,4 +1,5 @@
 #include "PluginContainer.h"
+
 #include "plugin/FunctionData.h"
 #include "plugin/HookData.h"
 #include "plugin/PluginConfigData.h"
@@ -6,7 +7,9 @@
 #include "plugin/PluginLinkInformation.h"
 #include "plugin/RelocationData.h"
 #include "plugin/SectionInfo.h"
+
 #include "utils/buttoncombo/ButtonComboUtils.h"
+#include "utils/logger.h"
 #include "utils/storage/StorageUtils.h"
 
 #include <optional>
@@ -17,10 +20,15 @@ PluginContainer::PluginContainer(PluginMetaInformation metaInformation, PluginLi
       mPluginData(std::move(pluginData)) {
     // Abuse this as a stable handle that references itself and survives std::move
     *mHandle = reinterpret_cast<uint32_t>(mHandle.get());
+
+    if (const bool res = useTrackingPluginHeapMemoryAllocator(mMetaInformation.getHeapTrackingOptions()); !res) {
+        DEBUG_FUNCTION_LINE_WARN("Failed to set heap tracking options for \"%s\"", mMetaInformation.getName().c_str());
+    }
 }
 
 PluginContainer::PluginContainer(PluginContainer &&src) noexcept : mMetaInformation(std::move(src.mMetaInformation)),
                                                                    mPluginLinkInformation(std::move(src.mPluginLinkInformation)),
+                                                                   mTrackingHeapAllocatorOpt(std::move(src.mTrackingHeapAllocatorOpt)),
                                                                    mPluginData(std::move(src.mPluginData)),
                                                                    mHandle(std::move(src.mHandle)),
                                                                    mPluginConfigData(std::move(src.mPluginConfigData)),
@@ -39,6 +47,7 @@ PluginContainer &PluginContainer::operator=(PluginContainer &&src) noexcept {
         this->mPluginData               = std::move(src.mPluginData);
         this->mPluginConfigData         = std::move(src.mPluginConfigData);
         this->mHandle                   = std::move(src.mHandle);
+        this->mTrackingHeapAllocatorOpt = std::move(src.mTrackingHeapAllocatorOpt);
         this->mStorageRootItem          = src.mStorageRootItem;
         this->mInitDone                 = src.mInitDone;
         this->mButtonComboManagerHandle = src.mButtonComboManagerHandle;
@@ -135,4 +144,31 @@ void PluginContainer::DeinitButtonComboData() {
 
 uint32_t PluginContainer::getButtonComboManagerHandle() const {
     return mButtonComboManagerHandle;
+}
+
+bool PluginContainer::useTrackingPluginHeapMemoryAllocator(PluginMetaInformation::HeapTrackingOptions options) {
+    if (options != PluginMetaInformation::TRACK_HEAP_OPTIONS_NONE) {
+        if (!this->mTrackingHeapAllocatorOpt) {
+            uint32_t stackTraceDepth = 0;
+            if (options == PluginMetaInformation::TRACK_HEAP_OPTIONS_TRACK_SIZE_AND_COLLECT_STACK_TRACES) {
+                stackTraceDepth = 8;
+            }
+            this->mTrackingHeapAllocatorOpt = TrackingPluginHeapMemoryAllocator::Create(this->mMetaInformation.getName(), stackTraceDepth);
+        }
+
+        return this->mTrackingHeapAllocatorOpt.has_value();
+    }
+    this->mTrackingHeapAllocatorOpt.reset();
+    return true;
+}
+
+bool PluginContainer::isUsingTrackingPluginHeapMemoryAllocator() const {
+    return mTrackingHeapAllocatorOpt.has_value();
+}
+
+const IPluginHeapMemoryAllocator &PluginContainer::getMemoryAllocator() const {
+    if (mTrackingHeapAllocatorOpt) {
+        return *mTrackingHeapAllocatorOpt;
+    }
+    return DefaultPluginHeapMemoryAllocator::gDefaultPluginHeapMemoryAllocator;
 }

source/plugin/PluginContainer.h

@@ -20,6 +20,7 @@
 #include "PluginConfigData.h"
 #include "PluginLinkInformation.h"
 #include "PluginMetaInformation.h"
+#include "TrackingPluginHeapMemoryAllocator.h"
 
 #include <wups/storage.h>
 
@@ -69,9 +70,19 @@ public:
 
     [[nodiscard]] uint32_t getButtonComboManagerHandle() const;
 
+    /**
+     * @return Returns true if setting the value was successful
+     */
+    bool useTrackingPluginHeapMemoryAllocator(PluginMetaInformation::HeapTrackingOptions options);
+
+    [[nodiscard]] bool isUsingTrackingPluginHeapMemoryAllocator() const;
+
+    [[nodiscard]] const IPluginHeapMemoryAllocator &getMemoryAllocator() const;
+
 private:
     PluginMetaInformation mMetaInformation;
     PluginLinkInformation mPluginLinkInformation;
+    std::optional<TrackingPluginHeapMemoryAllocator> mTrackingHeapAllocatorOpt;
     std::shared_ptr<PluginData> mPluginData;
 
     std::unique_ptr<uint32_t> mHandle = std::make_unique<uint32_t>();

source/plugin/PluginMetaInformation.cpp

@@ -38,6 +38,10 @@
     return mSize;
 }
 
+[[nodiscard]] PluginMetaInformation::HeapTrackingOptions PluginMetaInformation::getHeapTrackingOptions() const {
+    return mHeapTrackingOptions;
+}
+
 PluginMetaInformation::PluginMetaInformation() = default;
 
 void PluginMetaInformation::setName(std::string name) {
@@ -78,4 +82,8 @@ void PluginMetaInformation::setSize(const size_t size) {
 
 void PluginMetaInformation::setStorageId(std::string storageId) {
     mStorageId = std::move(storageId);
+}
+
+void PluginMetaInformation::setHeapTrackingOptions(HeapTrackingOptions value) {
+    mHeapTrackingOptions = value;
 }

source/plugin/PluginMetaInformation.h

@@ -26,7 +26,14 @@
 class WUPSVersion;
 
 class PluginMetaInformation {
+
 public:
+    enum HeapTrackingOptions {
+        TRACK_HEAP_OPTIONS_NONE                                = 0,
+        TRACK_HEAP_OPTIONS_TRACK_SIZE                          = 1,
+        TRACK_HEAP_OPTIONS_TRACK_SIZE_AND_COLLECT_STACK_TRACES = 2,
+    };
+
     [[nodiscard]] const std::string &getName() const;
 
     [[nodiscard]] const std::string &getAuthor() const;
@@ -45,6 +52,8 @@ public:
 
     [[nodiscard]] size_t getSize() const;
 
+    [[nodiscard]] HeapTrackingOptions getHeapTrackingOptions() const;
+
 private:
     PluginMetaInformation();
 
@@ -68,6 +77,8 @@ private:
 
     void setStorageId(std::string storageId);
 
+    void setHeapTrackingOptions(HeapTrackingOptions value);
+
     std::string mName;
     std::string mAuthor;
     std::string mVersion;
@@ -75,8 +86,9 @@ private:
     std::string mBuildTimestamp;
     std::string mDescription;
     std::string mStorageId;
-    size_t mSize             = {};
-    WUPSVersion mWUPSVersion = WUPSVersion(0, 0, 0);
+    size_t mSize                             = {};
+    WUPSVersion mWUPSVersion                 = WUPSVersion(0, 0, 0);
+    HeapTrackingOptions mHeapTrackingOptions = TRACK_HEAP_OPTIONS_NONE;
 
     friend class PluginMetaInformationFactory;
 

source/plugin/PluginMetaInformationFactory.cpp

@@ -108,6 +108,12 @@ std::optional<PluginMetaInformation> PluginMetaInformationFactory::loadPlugin(st
                         pluginInfo.setDescription(value);
                     } else if (key == "storage_id") {
                         pluginInfo.setStorageId(value);
+                    } else if (key == "debug") {
+                        if (value == "track_heap") {
+                            pluginInfo.setHeapTrackingOptions(PluginMetaInformation::TRACK_HEAP_OPTIONS_TRACK_SIZE);
+                        } else if (value == "track_heap_with_stack_trace") {
+                            pluginInfo.setHeapTrackingOptions(PluginMetaInformation::TRACK_HEAP_OPTIONS_TRACK_SIZE_AND_COLLECT_STACK_TRACES);
+                        }
                     } else if (key == "wups") {
                         if (value == "0.7.1") {
                             pluginInfo.setWUPSVersion(0, 7, 1);

source/plugin/TrackingPluginHeapMemoryAllocator.cpp

@@ -0,0 +1,177 @@
+#include "TrackingPluginHeapMemoryAllocator.h"
+
+#include "utils/logger.h"
+#include "utils/utils.h"
+
+#include <array>
+
+namespace {
+    struct PluginMemoryStatsInternal {
+        mutable std::mutex mutex;
+        bool inUse               = false;
+        uint32_t stackTraceDepth = 0;
+        PluginMemorySnapshot data;
+    };
+
+    // Generate 64 different thunk functions so we can track 64 plugins at a time which should be more than enough
+    constexpr int MAX_SLOTS = 64;
+    template<int N>
+    struct MemoryThunk {
+        inline static PluginMemoryStatsInternal stats;
+
+        static void *Alloc(uint32_t size) {
+            const auto ptr = (*DefaultPluginHeapMemoryAllocator::gDefaultPluginHeapMemoryAllocator.GetAllocFunctionAddress())(size);
+
+            if (ptr) {
+                std::lock_guard lock(stats.mutex);
+                if (stats.inUse) {
+                    stats.data.allocCount++;
+                    stats.data.currentAllocated += size;
+                    if (stats.data.currentAllocated > stats.data.peakAllocated) {
+                        stats.data.peakAllocated = stats.data.currentAllocated;
+                    }
+
+                    stats.data.allocationMap[ptr] = {size, stats.stackTraceDepth > 0 ? CaptureStackTrace(stats.stackTraceDepth) : std::vector<uint32_t>()};
+                }
+            }
+            return ptr;
+        }
+
+        static void *AllocEx(uint32_t size, int align) {
+            const auto ptr = (*DefaultPluginHeapMemoryAllocator::gDefaultPluginHeapMemoryAllocator.GetAllocExFunctionAddress())(size, align);
+            if (ptr) {
+                std::lock_guard lock(stats.mutex);
+                if (stats.inUse) {
+                    stats.data.allocCount++;
+                    stats.data.currentAllocated += size;
+                    if (stats.data.currentAllocated > stats.data.peakAllocated) {
+                        stats.data.peakAllocated = stats.data.currentAllocated;
+                    }
+                    stats.data.allocationMap[ptr] = {size, stats.stackTraceDepth > 0 ? CaptureStackTrace(stats.stackTraceDepth) : std::vector<uint32_t>()};
+                }
+            }
+            return ptr;
+        }
+
+        static void Free(void *ptr) {
+            if (!ptr) return;
+
+            std::lock_guard lock(stats.mutex);
+            if (stats.inUse) {
+                if (const auto it = stats.data.allocationMap.find(ptr); it != stats.data.allocationMap.end()) {
+                    uint32_t size = it->second.size;
+                    if (stats.data.currentAllocated >= size) {
+                        stats.data.currentAllocated -= size;
+                    } else {
+                        stats.data.currentAllocated = 0;
+                    }
+                    stats.data.allocationMap.erase(it);
+                } else {
+                    DEBUG_FUNCTION_LINE_ERR("free() of for unknown ptr detected (double free?). \"%s\": %p", stats.data.pluginName.c_str(), ptr);
+                    auto stackTrace = CaptureStackTrace(16);
+                    PrintCapturedStackTrace(stackTrace);
+                }
+                stats.data.freeCount++;
+            }
+
+            (*DefaultPluginHeapMemoryAllocator::gDefaultPluginHeapMemoryAllocator.GetFreeFunctionAddress())(ptr);
+        }
+    };
+
+    template<int... Is>
+    constexpr std::array<MEMAllocFromDefaultHeapFn, sizeof...(Is)> CreateAllocTable(std::integer_sequence<int, Is...>) { return {MemoryThunk<Is>::Alloc...}; }
+    template<int... Is>
+    constexpr std::array<MEMAllocFromDefaultHeapExFn, sizeof...(Is)> CreateAllocExTable(std::integer_sequence<int, Is...>) { return {MemoryThunk<Is>::AllocEx...}; }
+    template<int... Is>
+    constexpr std::array<MEMFreeToDefaultHeapFn, sizeof...(Is)> CreateFreeTable(std::integer_sequence<int, Is...>) { return {MemoryThunk<Is>::Free...}; }
+
+    template<int... Is>
+    constexpr std::array<PluginMemoryStatsInternal *, sizeof...(Is)> CreateStatsTable(std::integer_sequence<int, Is...>) {
+        return {&MemoryThunk<Is>::stats...};
+    }
+
+    auto sAllocThunks   = CreateAllocTable(std::make_integer_sequence<int, MAX_SLOTS>{});
+    auto sAllocExThunks = CreateAllocExTable(std::make_integer_sequence<int, MAX_SLOTS>{});
+    auto sFreeThunks    = CreateFreeTable(std::make_integer_sequence<int, MAX_SLOTS>{});
+    auto sStatsTable    = CreateStatsTable(std::make_integer_sequence<int, MAX_SLOTS>{});
+} // namespace
+
+TrackingPluginHeapMemoryAllocator::TrackingPluginHeapMemoryAllocator(const int32_t index) : sIndex(index) {
+}
+
+std::optional<TrackingPluginHeapMemoryAllocator> TrackingPluginHeapMemoryAllocator::Create(const std::string_view pluginName, uint32_t stackTraceDepth) {
+    int32_t index = 0;
+    for (auto *stat : sStatsTable) {
+        std::lock_guard lock(stat->mutex);
+        if (!stat->inUse) {
+            stat->inUse           = true;
+            stat->stackTraceDepth = stackTraceDepth;
+            stat->data            = {};
+            stat->data.pluginName = pluginName;
+            return TrackingPluginHeapMemoryAllocator{index};
+        }
+        ++index;
+    }
+    return std::nullopt;
+}
+
+TrackingPluginHeapMemoryAllocator::~TrackingPluginHeapMemoryAllocator() {
+    if (sIndex < 0 || sIndex >= MAX_SLOTS) {
+        return;
+    }
+    auto *stats = sStatsTable[sIndex];
+    std::lock_guard lock(stats->mutex);
+    stats->inUse = false;
+}
+
+TrackingPluginHeapMemoryAllocator::TrackingPluginHeapMemoryAllocator(TrackingPluginHeapMemoryAllocator &&src) noexcept {
+    this->sIndex = src.sIndex;
+    src.sIndex   = -1;
+}
+
+TrackingPluginHeapMemoryAllocator &TrackingPluginHeapMemoryAllocator::operator=(TrackingPluginHeapMemoryAllocator &&src) noexcept {
+    if (this != &src) {
+        // Release current slot if we have one
+        if (this->sIndex >= 0 && this->sIndex < MAX_SLOTS) {
+            auto *stats = sStatsTable[this->sIndex];
+            std::lock_guard lock(stats->mutex);
+            stats->inUse = false;
+            stats->data  = {};
+        }
+
+        this->sIndex = src.sIndex;
+        src.sIndex   = -1;
+    }
+    return *this;
+}
+
+MEMAllocFromDefaultHeapFn *TrackingPluginHeapMemoryAllocator::GetAllocFunctionAddress() const {
+    if (sIndex < 0 || sIndex >= MAX_SLOTS) {
+        return DefaultPluginHeapMemoryAllocator::gDefaultPluginHeapMemoryAllocator.GetAllocFunctionAddress();
+    }
+
+    return &sAllocThunks[sIndex];
+}
+
+MEMAllocFromDefaultHeapExFn *TrackingPluginHeapMemoryAllocator::GetAllocExFunctionAddress() const {
+    if (sIndex < 0 || sIndex >= MAX_SLOTS) {
+        return DefaultPluginHeapMemoryAllocator::gDefaultPluginHeapMemoryAllocator.GetAllocExFunctionAddress();
+    }
+    return &sAllocExThunks[sIndex];
+}
+
+MEMFreeToDefaultHeapFn *TrackingPluginHeapMemoryAllocator::GetFreeFunctionAddress() const {
+    if (sIndex < 0 || sIndex >= MAX_SLOTS) {
+        return DefaultPluginHeapMemoryAllocator::gDefaultPluginHeapMemoryAllocator.GetFreeFunctionAddress();
+    }
+    return &sFreeThunks[sIndex];
+}
+
+std::optional<PluginMemorySnapshot> TrackingPluginHeapMemoryAllocator::GetHeapMemoryUsageSnapshot() const {
+    if (sIndex < 0 || sIndex >= MAX_SLOTS) {
+        return std::nullopt;
+    }
+    const auto *stats = sStatsTable[sIndex];
+    std::lock_guard lock(stats->mutex);
+    return stats->data;
+}

source/plugin/TrackingPluginHeapMemoryAllocator.h

@@ -0,0 +1,29 @@
+#pragma once
+#include "DefaultPluginHeapMemoryAllocator.h"
+
+#include <optional>
+#include <string>
+
+#include <cstdint>
+
+class TrackingPluginHeapMemoryAllocator : public IPluginHeapMemoryAllocator {
+public:
+    static std::optional<TrackingPluginHeapMemoryAllocator> Create(std::string_view pluginName, uint32_t stackTraceDepth);
+
+    ~TrackingPluginHeapMemoryAllocator() override;
+
+    TrackingPluginHeapMemoryAllocator(const TrackingPluginHeapMemoryAllocator &) = delete;
+
+    TrackingPluginHeapMemoryAllocator(TrackingPluginHeapMemoryAllocator &&src) noexcept;
+
+    TrackingPluginHeapMemoryAllocator &operator=(TrackingPluginHeapMemoryAllocator &&src) noexcept;
+
+    MEMAllocFromDefaultHeapFn *GetAllocFunctionAddress() const override;
+    MEMAllocFromDefaultHeapExFn *GetAllocExFunctionAddress() const override;
+    MEMFreeToDefaultHeapFn *GetFreeFunctionAddress() const override;
+    std::optional<PluginMemorySnapshot> GetHeapMemoryUsageSnapshot() const override;
+
+private:
+    int sIndex = -1;
+    TrackingPluginHeapMemoryAllocator(int index);
+};