From bf995414496f7cee0d6688b3735415bb42caf7a0 Mon Sep 17 00:00:00 2001
From: Maschell <Maschell@gmx.de>
Date: Fri, 26 Apr 2024 23:32:39 +0200
Subject: [PATCH] Use custom function to check exp heap

---
 source/main.cpp           |  7 ++--
 source/memory_mapping.cpp | 69 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/source/main.cpp b/source/main.cpp
index 648f4cb..f0679d1 100644
--- a/source/main.cpp
+++ b/source/main.cpp
@@ -82,6 +82,9 @@ WUMS_INITIALIZE(args) {
 }
 
 WUMS_APPLICATION_STARTS() {
+#ifdef DEBUG
+    initLogging();
+#endif
     OSReport("Running MemoryMappingModule " VERSION VERSION_EXTRA "\n");
 
     MemoryMapping_checkHeaps();
@@ -89,10 +92,6 @@ WUMS_APPLICATION_STARTS() {
     // Now we can update the pointer with the "real" functions
     gMEMAllocFromDefaultHeapExForThreads = MEMAllocFromDefaultHeapEx;
     gMEMFreeToDefaultHeapForThreads      = MEMFreeToDefaultHeap;
-
-#ifdef DEBUG
-    initLogging();
-#endif
 }
 
 
diff --git a/source/memory_mapping.cpp b/source/memory_mapping.cpp
index 7d257da..ff6f57e 100644
--- a/source/memory_mapping.cpp
+++ b/source/memory_mapping.cpp
@@ -591,6 +591,73 @@ void *MemoryMapping_allocEx(uint32_t size, int32_t align, bool videoOnly) {
     return res;
 }
 
+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 %08X actual %08X", 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: %08X size %d; heap start %08X heap end %08X", 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 %08X, actual %08X", 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 MemoryMapping_checkHeaps() {
     OSLockMutex(&allocMutex);
     for (int32_t i = 0; /* waiting for a break */; i++) {
@@ -598,7 +665,7 @@ void MemoryMapping_checkHeaps() {
             break;
         }
         auto heapHandle = (MEMHeapHandle) mem_mapping[i].effective_start_address;
-        if (!MEMCheckExpHeap(heapHandle, MEM_EXP_HEAP_CHECK_FLAGS_LOG_ERRORS)) {
+        if (!CheckMemExpHeap(reinterpret_cast<MEMExpHeap *>(heapHandle))) {
             DEBUG_FUNCTION_LINE_ERR("MemoryMapping heap %08X (index %d) is corrupted.", heapHandle, i);
 #ifdef DEBUG
             OSFatal("MemoryMappingModule: Heap is corrupted");