Check heap on application start

wumsloader/src/entry.cpp

@@ -11,6 +11,7 @@
 #include "utils/SegmentedTimer.h"
 #include "utils/dynamic.h"
 #include "utils/logger.h"
+#include "utils/memory.h"
 #include "utils/reent.h"
 
 #include <coreinit/debug.h>
@@ -148,6 +149,8 @@ void doStart(int argc, char **argv) {
     Modules::HooksManagement::CallHook(gLoadedModules, WUMS_HOOK_APPLICATION_STARTS);
     Modules::HooksManagement::CallHook(gLoadedModules, WUMS_HOOK_ALL_APPLICATION_STARTS_DONE);
 
+    Memory_CheckHeaps();
+
     deinitLogging();
     fini_wut();
 }

wumsloader/src/module/RelocationUtils.cpp

@@ -53,7 +53,6 @@ void CustomDynLoadFree(void *addr) {
     }
 
     free(addr);
-
 }
 
 bool doRelocation(const std::vector<RelocationData> &relocData,

wumsloader/src/utils/memory.cpp

@@ -135,3 +135,79 @@ _pvalloc_r(struct _reent *r, size_t size) {
 int _malloc_trim_r(struct _reent *r, size_t pad) {
     return 0;
 }
+
+bool CheckMemExpHeapBlock(MEMExpHeap *heap, MEMExpHeapBlockList *block, uint32_t tag, const char *listName, uint32_t &totalSizeOut) {
+    MEMExpHeapBlock *prevBlock = nullptr;
+    for (auto *cur = block->head; cur != nullptr; cur = cur->next) {
+        if (cur->prev != prevBlock) {
+            DEBUG_FUNCTION_LINE_ERR("[Exp Heap Check] \"%s\" prev is invalid. expected %p actual %p", listName, prevBlock, cur->prev);
+            return false;
+        }
+        if (cur < heap->header.dataStart || cur > heap->header.dataEnd || ((uint32_t) cur + sizeof(MEMExpHeapBlock) + cur->blockSize) > (uint32_t) heap->header.dataEnd) {
+            DEBUG_FUNCTION_LINE_ERR("[Exp Heap Check] Block is not inside heap. block: %p size %d; heap start %p heap end %p", cur, sizeof(MEMExpHeapBlock) + cur->blockSize, heap->header.dataStart, heap->header.dataEnd);
+            return false;
+        }
+        if (cur->tag != tag) {
+            DEBUG_FUNCTION_LINE_ERR("[Exp Heap Check] Invalid block tag expected %04X, actual %04X", tag, cur->tag);
+            return false;
+        }
+
+        totalSizeOut = totalSizeOut + cur->blockSize + (cur->attribs >> 8 & 0x7fffff) + sizeof(MEMExpHeapBlock);
+        prevBlock    = cur;
+    }
+    if (prevBlock != block->tail) {
+        DEBUG_FUNCTION_LINE_ERR("[Exp Heap Check] \"%s\" tail is unexpected! expected %p, actual %p", listName, heap->usedList.tail, prevBlock);
+        return false;
+    }
+    return true;
+}
+
+bool CheckMemExpHeapCore(MEMExpHeap *heap) {
+    uint32_t totalSize = 0;
+#pragma GCC diagnostic ignored "-Waddress-of-packed-member"
+    if (!CheckMemExpHeapBlock(heap, &heap->usedList, 0x5544, "used", totalSize)) {
+        return false;
+    }
+
+#pragma GCC diagnostic ignored "-Waddress-of-packed-member"
+    if (!CheckMemExpHeapBlock(heap, &heap->freeList, 0x4652, "free", totalSize)) {
+        return false;
+    }
+
+    if (totalSize != (uint32_t) heap->header.dataEnd - (uint32_t) heap->header.dataStart) {
+        DEBUG_FUNCTION_LINE_ERR("[Exp Heap Check] heap size is unexpected! expected %08X, actual %08X", (uint32_t) heap->header.dataEnd - (uint32_t) heap->header.dataStart, totalSize);
+        return false;
+    }
+    return true;
+}
+
+bool CheckMemExpHeap(MEMExpHeap *heap) {
+    OSMemoryBarrier();
+    if (heap->header.tag != MEM_EXPANDED_HEAP_TAG) {
+        DEBUG_FUNCTION_LINE_ERR("[Exp Heap Check] Invalid heap handle. - %08X", heap->header.tag);
+        return false;
+    }
+
+    if (heap->header.flags & MEM_HEAP_FLAG_USE_LOCK) {
+#pragma GCC diagnostic ignored "-Waddress-of-packed-member"
+        OSUninterruptibleSpinLock_Acquire(&(heap->header).lock);
+    }
+
+    auto result = CheckMemExpHeapCore(heap);
+
+    if (heap->header.flags & MEM_HEAP_FLAG_USE_LOCK) {
+#pragma GCC diagnostic ignored "-Waddress-of-packed-member"
+        OSUninterruptibleSpinLock_Release(&(heap->header).lock);
+    }
+
+    return result;
+}
+
+void Memory_CheckHeaps() {
+    if (!CheckMemExpHeap(reinterpret_cast<MEMExpHeap *>(gHeapHandle))) {
+        DEBUG_FUNCTION_LINE_ERR("MemoryMapping heap %p is corrupted.", gHeapHandle);
+#ifdef DEBUG
+        OSFatal("MemoryMappingModule: Heap is corrupted");
+#endif
+    }
+}

wumsloader/src/utils/memory.h

@@ -2,4 +2,6 @@
 
 extern uint32_t MEMAlloc;
 extern uint32_t MEMAllocEx;
-extern uint32_t MEMFree;
+extern uint32_t MEMFree;
+
+void Memory_CheckHeaps();