Avoid jump data heap corruption by flushing the whole heap

source/PatchedFunctionData.cpp

@@ -1,5 +1,6 @@
 #include "PatchedFunctionData.h"
 #include "utils/KernelFindExport.h"
+#include "utils/globals.h"
 #include "utils/utils.h"
 #include <coreinit/mcp.h>
 #include <coreinit/title.h>
@@ -102,16 +103,14 @@ bool PatchedFunctionData::allocateDataForJumps() {
     }
     if (this->replacementFunctionAddress > 0x01FFFFFC || this->targetProcess != FP_TARGET_PROCESS_ALL) {
         this->jumpDataSize = 15; // We could predict the actual size and save some memory, but at the moment we don't need it.
-        this->jumpData     = (uint32_t *) MEMAllocFromExpHeapEx(this->heapHandle, this->jumpDataSize * sizeof(uint32_t), 4);
+        this->jumpData     = (uint32_t *) MEMAllocFromExpHeapEx(this->heapHandle, this->jumpDataSize * sizeof(uint32_t), 0x20);
 
         if (!this->jumpData) {
             DEBUG_FUNCTION_LINE_ERR("Failed to alloc jump data");
             return false;
         }
     }
-
+    this->jumpToOriginal = (uint32_t *) MEMAllocFromExpHeapEx(this->heapHandle, 0x5 * sizeof(uint32_t), 0x20);
-    this->jumpToOriginal = (uint32_t *) MEMAllocFromExpHeapEx(this->heapHandle, 0x5 * sizeof(uint32_t), 4);
-
     if (!this->jumpToOriginal) {
         DEBUG_FUNCTION_LINE_ERR("Failed to alloc jump data");
         return false;
@@ -244,10 +243,11 @@ void PatchedFunctionData::generateJumpToOriginal() {
         this->jumpToOriginal[1] = 0x48000002 | (jumpToAddress & 0x01FFFFFC);
     }
 
-    DCFlushRange((void *) this->jumpToOriginal, sizeof(uint32_t) * 5);
+    DCFlushRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
-    ICInvalidateRange((void *) this->jumpToOriginal, sizeof(uint32_t) * 5);
+    ICInvalidateRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
 
     *(this->realCallFunctionAddressPtr) = (uint32_t) this->jumpToOriginal;
+
     OSMemoryBarrier();
 }
 
@@ -311,8 +311,8 @@ void PatchedFunctionData::generateReplacementJump() {
 
         this->replaceWithInstruction = 0x48000002 | ((uint32_t) this->jumpData & 0x01FFFFFC);
 
-        DCFlushRange((void *) this->jumpData, sizeof(uint32_t) * 15);
+        DCFlushRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
-        ICInvalidateRange((void *) this->jumpData, sizeof(uint32_t) * 15);
+        ICInvalidateRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
     }
 
     DCFlushRange((void *) &replaceWithInstruction, 4);
@@ -330,6 +330,9 @@ PatchedFunctionData::~PatchedFunctionData() {
         MEMFreeToExpHeap(this->heapHandle, this->jumpData);
         this->jumpData = nullptr;
     }
+
+    DCFlushRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
+    ICInvalidateRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
 }
 
 bool PatchedFunctionData::shouldBePatched() const {

source/function_patcher.cpp

@@ -101,6 +101,9 @@ bool RestoreFunction(std::shared_ptr<PatchedFunctionData> &patchedFunction) {
         return false;
     }
 
+    DCFlushRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
+    ICInvalidateRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
+
     auto targetAddrPhys = (uint32_t) patchedFunction->realPhysicalFunctionAddress;
 
     if (patchedFunction->library != LIBRARY_OTHER) {
@@ -137,6 +140,9 @@ bool RestoreFunction(std::shared_ptr<PatchedFunctionData> &patchedFunction) {
     ICInvalidateRange((void *) patchedFunction->realEffectiveFunctionAddress, 4);
     DCFlushRange((void *) patchedFunction->realEffectiveFunctionAddress, 4);
 
+    DCFlushRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
+    ICInvalidateRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE);
+
     patchedFunction->isPatched = false;
     return true;
 }

source/main.cpp

@@ -113,7 +113,7 @@ WUMS_INITIALIZE() {
     }
 
     memset(gJumpHeapData, 0, JUMP_HEAP_DATA_SIZE);
-    gJumpHeapHandle = MEMCreateExpHeapEx((void *) (gJumpHeapData), JUMP_HEAP_DATA_SIZE, 1);
+    gJumpHeapHandle = MEMCreateExpHeapEx((void *) (gJumpHeapData), JUMP_HEAP_DATA_SIZE, MEM_HEAP_FLAG_USE_LOCK);
     if (gJumpHeapHandle == nullptr) {
         DEBUG_FUNCTION_LINE_ERR("Failed to create heap for jump data");
         OSFatal("FunctionPatcherModule: Failed to create heap for jump data");
@@ -177,6 +177,8 @@ WUMS_APPLICATION_STARTS() {
         OSMemoryBarrier();
         OSDynLoad_AddNotifyCallback(notify_callback, nullptr);
     }
+
+    CheckMemExpHeapJumpData();
 }
 
 WUMS_APPLICATION_REQUESTS_EXIT() {

source/utils/utils.cpp

@@ -1,4 +1,11 @@
+#include "CThread.h"
+#include "globals.h"
+#include "logger.h"
+
+
 #include <coreinit/cache.h>
+#include <coreinit/core.h>
+#include <coreinit/memexpheap.h>
 #include <coreinit/memorymap.h>
 #include <kernel/kernel.h>
 
@@ -25,4 +32,92 @@ bool ReadFromPhysicalAddress(uint32_t srcPhys, uint32_t *out) {
     DCFlushRange((void *) &currentInstruction, 4);
     *out = currentInstruction;
     return true;
-}
+}
+
+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("[%p][%d][Exp Heap Check] Invalid block tag expected %04X, actual %04X", &cur->tag, OSGetCoreId(), 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;
+}
+
+
+static void CheckMemExpHeapJumpDataCallback(CThread *, void *) {
+    if (gJumpHeapHandle != nullptr) {
+        if (!CheckMemExpHeap(reinterpret_cast<MEMExpHeap *>(gJumpHeapHandle))) {
+            OSFatal("FunctionPatcherModule: Corrupted heap");
+        } else {
+            DEBUG_FUNCTION_LINE_VERBOSE("JumpData heap has no curruption. Checked on core %d", OSGetCoreId());
+        }
+    }
+}
+
+void CheckMemExpHeapJumpData() {
+    CThread::runOnAllCores(CheckMemExpHeapJumpDataCallback, nullptr, 0, 16, 0x1000);
+}

source/utils/utils.h

@@ -12,3 +12,5 @@ std::shared_ptr<T> make_shared_nothrow(Args &&...args) noexcept(noexcept(T(std::
 }
 
 bool ReadFromPhysicalAddress(uint32_t srcPhys, uint32_t *out);
+
+void CheckMemExpHeapJumpData();