diff --git a/source/PatchedFunctionData.cpp b/source/PatchedFunctionData.cpp index 948eb46..6636b13 100644 --- a/source/PatchedFunctionData.cpp +++ b/source/PatchedFunctionData.cpp @@ -1,5 +1,6 @@ #include "PatchedFunctionData.h" #include "utils/KernelFindExport.h" +#include "utils/globals.h" #include "utils/utils.h" #include #include @@ -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), 4); - + this->jumpToOriginal = (uint32_t *) MEMAllocFromExpHeapEx(this->heapHandle, 0x5 * sizeof(uint32_t), 0x20); 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); - ICInvalidateRange((void *) this->jumpToOriginal, sizeof(uint32_t) * 5); + DCFlushRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE); + 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); - ICInvalidateRange((void *) this->jumpData, sizeof(uint32_t) * 15); + DCFlushRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE); + 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 { diff --git a/source/function_patcher.cpp b/source/function_patcher.cpp index 69892d6..715102e 100644 --- a/source/function_patcher.cpp +++ b/source/function_patcher.cpp @@ -101,6 +101,9 @@ bool RestoreFunction(std::shared_ptr &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 &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; } diff --git a/source/main.cpp b/source/main.cpp index 3349a0c..2615232 100644 --- a/source/main.cpp +++ b/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() { diff --git a/source/utils/utils.cpp b/source/utils/utils.cpp index 493836f..5448f37 100644 --- a/source/utils/utils.cpp +++ b/source/utils/utils.cpp @@ -1,4 +1,11 @@ +#include "CThread.h" +#include "globals.h" +#include "logger.h" + + #include +#include +#include #include #include @@ -25,4 +32,92 @@ bool ReadFromPhysicalAddress(uint32_t srcPhys, uint32_t *out) { DCFlushRange((void *) ¤tInstruction, 4); *out = currentInstruction; return true; -} \ No newline at end of file +} + +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(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); +} diff --git a/source/utils/utils.h b/source/utils/utils.h index e63a2e2..85364e3 100644 --- a/source/utils/utils.h +++ b/source/utils/utils.h @@ -12,3 +12,5 @@ std::shared_ptr make_shared_nothrow(Args &&...args) noexcept(noexcept(T(std:: } bool ReadFromPhysicalAddress(uint32_t srcPhys, uint32_t *out); + +void CheckMemExpHeapJumpData(); \ No newline at end of file