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/export.cpp b/source/export.cpp index b6d15b6..d67ce26 100644 --- a/source/export.cpp +++ b/source/export.cpp @@ -12,50 +12,70 @@ WUT_CHECK_OFFSET(function_replacement_data_v2_t, 0x00, VERSION); WUT_CHECK_OFFSET(function_replacement_data_v3_t, 0x00, version); +FunctionPatcherStatus FPAddFunctionPatches(function_replacement_data_t **function_data_array, uint32_t count, PatchedFunctionHandle *outHandles, bool *outHasBeenPatchedArray); + FunctionPatcherStatus FPAddFunctionPatch(function_replacement_data_t *function_data, PatchedFunctionHandle *outHandle, bool *outHasBeenPatched) { - if (function_data == nullptr) { - DEBUG_FUNCTION_LINE_ERR("function_data was NULL"); + // Wrap the single patch into an array of size 1 and pass it to the batcher + function_replacement_data_t *arr[] = {function_data}; + return FPAddFunctionPatches(arr, 1, outHandle, outHasBeenPatched); +} + +FunctionPatcherStatus FPAddFunctionPatches(function_replacement_data_t **function_data_array, uint32_t count, PatchedFunctionHandle *outHandles, bool *outHasBeenPatchedArray) { + if (function_data_array == nullptr || count == 0) { + DEBUG_FUNCTION_LINE_ERR("function_data_array was NULL or count was 0"); return FUNCTION_PATCHER_RESULT_INVALID_ARGUMENT; } - if (function_data->version < 2 || function_data->version > 3) { - DEBUG_FUNCTION_LINE_ERR("Failed to patch function. struct version mismatch"); - return FUNCTION_PATCHER_RESULT_UNSUPPORTED_STRUCT_VERSION; + DEBUG_FUNCTION_LINE_ERR("Patching %d functions", count); + + std::vector> functionsToPatch; + functionsToPatch.reserve(count); + + for (uint32_t i = 0; i < count; ++i) { + auto *function_data = function_data_array[i]; + if (function_data == nullptr) { + DEBUG_FUNCTION_LINE_ERR("A function_data entry was NULL"); + return FUNCTION_PATCHER_RESULT_INVALID_ARGUMENT; + } + + if (function_data->version < 2 || function_data->version > 3) { + DEBUG_FUNCTION_LINE_ERR("Failed to patch function. struct version mismatch"); + return FUNCTION_PATCHER_RESULT_UNSUPPORTED_STRUCT_VERSION; + } + + std::optional> functionDataOpt; + if (function_data->version == 2) { + functionDataOpt = PatchedFunctionData::make_shared_v2(gFunctionAddressProvider, (function_replacement_data_v2_t *) function_data, gJumpHeapHandle); + } else if (function_data->version == 3) { + functionDataOpt = PatchedFunctionData::make_shared_v3(gFunctionAddressProvider, (function_replacement_data_v3_t *) function_data, gJumpHeapHandle); + } else { + DEBUG_FUNCTION_LINE_ERR("Unknown function_replacement_data_t struct version"); + OSFatal("Unknown function patching struct version. Update FunctionPatcherModule/Aroma."); + } + + if (!functionDataOpt) { + return FUNCTION_PATCHER_RESULT_UNKNOWN_ERROR; + } + + functionsToPatch.push_back(functionDataOpt.value()); } - std::optional> functionDataOpt; - if (function_data->version == 2) { - functionDataOpt = PatchedFunctionData::make_shared_v2(gFunctionAddressProvider, (function_replacement_data_v2_t *) function_data, gJumpHeapHandle); - } else if (function_data->version == 3) { - functionDataOpt = PatchedFunctionData::make_shared_v3(gFunctionAddressProvider, (function_replacement_data_v3_t *) function_data, gJumpHeapHandle); - } else { - // Should never happen. - DEBUG_FUNCTION_LINE_ERR("Unknown function_replacement_data_t struct version"); - OSFatal("Unknown function patching struct version. Update FunctionPatcherModule/Aroma."); - } - - if (!functionDataOpt) { - return FUNCTION_PATCHER_RESULT_UNKNOWN_ERROR; - } - - auto &functionData = functionDataOpt.value(); - - // PatchFunction calls OSFatal on fatal errors. - // If this function returns false the target function was not patched - // Usually this means the target RPL is not (yet) loaded. - auto patchResult = PatchFunction(functionData); - if (outHasBeenPatched) { - *outHasBeenPatched = patchResult; - } - - if (outHandle) { - *outHandle = functionData->getHandle(); - } + PatchFunctions(functionsToPatch); { std::lock_guard lock(gPatchedFunctionsMutex); - gPatchedFunctions.push_back(std::move(functionData)); + for (uint32_t i = 0; i < count; ++i) { + auto &funcData = functionsToPatch[i]; + if (outHasBeenPatchedArray) { + outHasBeenPatchedArray[i] = funcData->isPatched; + } + if (outHandles) { + outHandles[i] = funcData->getHandle(); + } + + gPatchedFunctions.push_back(std::move(funcData)); + } OSMemoryBarrier(); } @@ -127,7 +147,7 @@ FunctionPatcherStatus FPGetVersion(FunctionPatcherAPIVersion *outVersion) { if (outVersion == nullptr) { return FUNCTION_PATCHER_RESULT_INVALID_ARGUMENT; } - *outVersion = 2; + *outVersion = 3; return FUNCTION_PATCHER_RESULT_SUCCESS; } @@ -147,5 +167,6 @@ FunctionPatcherStatus FPIsFunctionPatched(PatchedFunctionHandle handle, bool *ou WUMS_EXPORT_FUNCTION(FPGetVersion); WUMS_EXPORT_FUNCTION(FPAddFunctionPatch); +WUMS_EXPORT_FUNCTION(FPAddFunctionPatches); WUMS_EXPORT_FUNCTION(FPRemoveFunctionPatch); WUMS_EXPORT_FUNCTION(FPIsFunctionPatched); \ No newline at end of file diff --git a/source/function_patcher.cpp b/source/function_patcher.cpp index 69892d6..c4a6f6f 100644 --- a/source/function_patcher.cpp +++ b/source/function_patcher.cpp @@ -5,93 +5,255 @@ #include "utils/logger.h" #include "utils/utils.h" +#include #include +#include #include #include +#include #include #include #include +#include -static void writeDataAndFlushIC(CThread *thread, void *arg) { +enum class PatchState { + PREPARE, + MAIN_CORE_PATCHING_DONE, + ALL_CORES_PATCHING_DONE +}; +static volatile PatchState gPatchState = PatchState::PREPARE; +static volatile int32_t gCoresReady = 0; +static volatile int32_t gCoresFlushed = 0; + +static std::recursive_mutex sPatch_RestoreMutex; +static OSSpinLock sGlobalSpinLock; +static bool sSpinLockInitialized = false; + +struct WorkerTask { + uint32_t targetPhys; + uint32_t sourcePhys; + uint32_t effectiveAddr; + void *jumpData; + uint32_t jumpDataSize; + void *jumpToOriginal; + void *realCallFunctionAddressPtr; +}; + +static void applyKernelPatchOnCore(CThread *thread, void *arg) { (void) thread; - auto *data = (PatchedFunctionData *) arg; + auto *tasks = (std::vector *) arg; - uint32_t replace_instruction = data->replaceWithInstruction; - uint32_t physical_address = data->realPhysicalFunctionAddress; - uint32_t effective_address = data->realEffectiveFunctionAddress; - DCFlushRange(&replace_instruction, 4); - DCFlushRange(&physical_address, 4); + OSSpinLock localLock; + OSInitSpinLock(&localLock); + OSUninterruptibleSpinLock_Acquire(&localLock); - auto replace_instruction_physical = (uint32_t) &replace_instruction; + OSAddAtomic(&gCoresReady, 1); - if (data->jumpData) { - DCFlushRange(data->jumpData, data->jumpDataSize * sizeof(uint32_t)); - ICInvalidateRange(data->jumpData, data->jumpDataSize * sizeof(uint32_t)); - } - if (data->jumpToOriginal) { - DCFlushRange(data->jumpToOriginal, 5 * sizeof(uint32_t)); - ICInvalidateRange(data->jumpToOriginal, 5 * sizeof(uint32_t)); - } - if (data->realCallFunctionAddressPtr) { - DCFlushRange(data->realCallFunctionAddressPtr, sizeof(uint32_t)); - ICInvalidateRange(data->realCallFunctionAddressPtr, sizeof(uint32_t)); + // Wait for the main core to finish preparing physical addresses + while (gPatchState == PatchState::PREPARE) { + asm volatile("nop"); } - if (replace_instruction_physical < 0x00800000 || replace_instruction_physical >= 0x01000000) { - replace_instruction_physical = OSEffectiveToPhysical(replace_instruction_physical); - } else { - replace_instruction_physical = replace_instruction_physical + 0x30800000 - 0x00800000; + for (const auto &task : *tasks) { + KernelCopyData(task.targetPhys, task.sourcePhys, 4); } - KernelCopyData(physical_address, replace_instruction_physical, 4); - ICInvalidateRange((void *) (effective_address), 4); + for (const auto &task : *tasks) { + if (task.jumpData) { + ICInvalidateRange(task.jumpData, task.jumpDataSize * sizeof(uint32_t)); + } + if (task.jumpToOriginal) { + ICInvalidateRange(task.jumpToOriginal, 5 * sizeof(uint32_t)); + } + if (task.realCallFunctionAddressPtr) { + ICInvalidateRange(task.realCallFunctionAddressPtr, sizeof(uint32_t)); + } + if (task.effectiveAddr) { + ICInvalidateRange((void *) task.effectiveAddr, 4); + } + } + + // Force pipeline flush + asm volatile("sync; isync"); + + // Atomically signal main core that our caches are clean + OSAddAtomic(&gCoresFlushed, 1); + + // Wait for the release signal + while (gPatchState == PatchState::ALL_CORES_PATCHING_DONE) {} + + // RESTORE INTERRUPTS + OSUninterruptibleSpinLock_Release(&localLock); } -bool PatchFunction(std::shared_ptr &patchedFunction) { - if (patchedFunction->isPatched) { +struct PatchDispatchCtx { + std::shared_ptr func; + bool result; +}; + +struct BatchPatchDispatchCtx { + std::vector> *list; + bool result; +}; + +static void PatchFunctionsBatchDispatcher(CThread *thread, void *arg) { + (void) thread; + auto *ctx = (BatchPatchDispatchCtx *) arg; + ctx->result = PatchFunctions(*(ctx->list)); +} + +static void RestoreFunctionDispatcher(CThread *thread, void *arg) { + (void) thread; + auto *ctx = (PatchDispatchCtx *) arg; + ctx->result = RestoreFunction(ctx->func); +} + +bool PatchFunctions(std::vector> &patchedFunctions) { + if (OSGetCoreId() != OSGetMainCoreId()) { + DEBUG_FUNCTION_LINE_INFO("PatchFunctions called from Core %d. Dispatching to Main Core %d...", OSGetCoreId(), OSGetMainCoreId()); + BatchPatchDispatchCtx ctx = {&patchedFunctions, false}; + { + CThread thread(CThread::eAttributeAffCore1, OSGetCurrentThread()->priority, 0x1000, PatchFunctionsBatchDispatcher, &ctx); + thread.resumeThread(); + } + return ctx.result; + } + std::lock_guard lock(sPatch_RestoreMutex); + + gPatchState = PatchState::PREPARE; + gCoresReady = 0; + gCoresFlushed = 0; + + std::vector> validToPatch; + validToPatch.reserve(patchedFunctions.size()); + + for (auto &patch : patchedFunctions) { + if (patch->isPatched) { continue; } + if (!patch->shouldBePatched()) { continue; } + if (!patch->updateFunctionAddresses()) { continue; } + + if (!ReadFromPhysicalAddress(patch->realPhysicalFunctionAddress, &patch->replacedInstruction)) { + DEBUG_FUNCTION_LINE_ERR("Failed to read instruction for %s", patch->functionName.value_or("").c_str()); + continue; + } + + patch->generateJumpToOriginal(); + patch->generateReplacementJump(); + validToPatch.push_back(patch); + } + + if (validToPatch.empty()) { return true; } - if (!patchedFunction->shouldBePatched()) { - return false; + + // This is very important. Otherwise the heap meta data might not up to date on all cores + DCFlushRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE); + ICInvalidateRange(gJumpHeapData, JUMP_HEAP_DATA_SIZE); + + std::vector tasks; + tasks.reserve(validToPatch.size()); + + for (auto &pf : validToPatch) { + WorkerTask task = {}; + task.targetPhys = pf->realPhysicalFunctionAddress; + task.effectiveAddr = pf->realEffectiveFunctionAddress; + task.jumpData = pf->jumpData; + task.jumpDataSize = pf->jumpDataSize; + task.jumpToOriginal = pf->jumpToOriginal; + task.realCallFunctionAddressPtr = pf->realCallFunctionAddressPtr; + + uint32_t replace_ptr = (uint32_t) &pf->replaceWithInstruction; + if (replace_ptr < 0x00800000 || replace_ptr >= 0x01000000) { + task.sourcePhys = OSEffectiveToPhysical(replace_ptr); + } else { + task.sourcePhys = replace_ptr + 0x30800000 - 0x00800000; + } + tasks.push_back(task); } - // The addresses of a function might change every time with run another application. - if (!patchedFunction->updateFunctionAddresses()) { - return false; + if (!sSpinLockInitialized) { + OSInitSpinLock(&sGlobalSpinLock); + sSpinLockInitialized = true; } - if (patchedFunction->functionName) { - DEBUG_FUNCTION_LINE("Patching function %s...", patchedFunction->functionName->c_str()); - } else { - DEBUG_FUNCTION_LINE("Patching function @ %08X", patchedFunction->realEffectiveFunctionAddress); + + CThread *threadA = CThread::create(applyKernelPatchOnCore, &tasks, CThread::eAttributeAffCore2, 0); + CThread *threadB = CThread::create(applyKernelPatchOnCore, &tasks, CThread::eAttributeAffCore0, 0); + threadA->resumeThread(); + threadB->resumeThread(); + + while (gCoresReady < 2) { + OSSleepTicks(1); } - if (!ReadFromPhysicalAddress(patchedFunction->realPhysicalFunctionAddress, &patchedFunction->replacedInstruction)) { - DEBUG_FUNCTION_LINE_ERR("Failed to read instruction."); - OSFatal("FunctionPatcherModule: Failed to read instruction."); - return false; + DEBUG_FUNCTION_LINE_VERBOSE("Applying %d patches on Core %d", validToPatch.size(), OSGetCoreId()); + + OSUninterruptibleSpinLock_Acquire(&sGlobalSpinLock); + + // D-Cache flush and Kernel Copy for all tasks + for (size_t i = 0; i < validToPatch.size(); ++i) { + auto &pf = validToPatch[i]; + auto &task = tasks[i]; + + uint32_t replace_ptr = (uint32_t) &pf->replaceWithInstruction; + DCFlushRange((void *) replace_ptr, 4); + + if (pf->jumpData) DCFlushRange(pf->jumpData, pf->jumpDataSize * sizeof(uint32_t)); + if (pf->jumpToOriginal) DCFlushRange(pf->jumpToOriginal, 5 * sizeof(uint32_t)); + if (pf->realCallFunctionAddressPtr) DCFlushRange(pf->realCallFunctionAddressPtr, sizeof(uint32_t)); + + KernelCopyData(task.targetPhys, task.sourcePhys, 4); } - // Generate a jump to the original function so the unpatched function can still be called - patchedFunction->generateJumpToOriginal(); + DCFlushRange(tasks.data(), sizeof(WorkerTask) * tasks.size()); + DCFlushRange(&tasks, sizeof(tasks)); + OSMemoryBarrier(); - // Generate a code that is run when somebody calls the patched function. - // If the correct process calls this, it'll jump the function replacement, otherwise the original function will be called. - patchedFunction->generateReplacementJump(); + // Release remote cores + gPatchState = PatchState::MAIN_CORE_PATCHING_DONE; - // Write this->replaceWithInstruction to the first instruction of the function we want to replace. - CThread::runOnAllCores(writeDataAndFlushIC, patchedFunction.get()); + // Invalidate Main Core I-Caches + for (auto &task : tasks) { + if (task.jumpData) { ICInvalidateRange(task.jumpData, task.jumpDataSize * sizeof(uint32_t)); } + if (task.jumpToOriginal) { ICInvalidateRange(task.jumpToOriginal, 5 * sizeof(uint32_t)); } + if (task.realCallFunctionAddressPtr) { ICInvalidateRange(task.realCallFunctionAddressPtr, sizeof(uint32_t)); } + ICInvalidateRange((void *) task.effectiveAddr, 4); + } + asm volatile("sync; isync"); - // Set patch status - patchedFunction->isPatched = true; + while (gCoresFlushed < 2) { asm volatile("nop"); } + + OSUninterruptibleSpinLock_Release(&sGlobalSpinLock); + gPatchState = PatchState::ALL_CORES_PATCHING_DONE; + + delete threadA; + delete threadB; + + for (auto &pf : validToPatch) { + pf->isPatched = true; + } return true; } +// Single Patch function simply wraps the Batch function +bool PatchFunction(std::shared_ptr &patchedFunction) { + std::vector list = {patchedFunction}; + return PatchFunctions(list); +} + bool RestoreFunction(std::shared_ptr &patchedFunction) { + if (OSGetCoreId() != OSGetMainCoreId()) { + DEBUG_FUNCTION_LINE_INFO("RestoreFunction called from Core %d. Dispatching to Main Core %d...", OSGetCoreId(), OSGetMainCoreId()); + PatchDispatchCtx ctx = {patchedFunction, false}; + CThread *dispatchThread = CThread::create(RestoreFunctionDispatcher, &ctx, CThread::eAttributeAffCore1, 0); + delete dispatchThread; + return ctx.result; + } + if (!patchedFunction->isPatched) { DEBUG_FUNCTION_LINE_VERBOSE("Skip restoring function because it's not patched"); return true; @@ -101,6 +263,8 @@ bool RestoreFunction(std::shared_ptr &patchedFunction) { return false; } + std::lock_guard lock(sPatch_RestoreMutex); + auto targetAddrPhys = (uint32_t) patchedFunction->realPhysicalFunctionAddress; if (patchedFunction->library != LIBRARY_OTHER) { @@ -133,9 +297,58 @@ bool RestoreFunction(std::shared_ptr &patchedFunction) { OSFatal("FunctionPatcherModule: Failed to get physical address"); } - KernelCopyData(targetAddrPhys, sourceAddrPhys, 4); - ICInvalidateRange((void *) patchedFunction->realEffectiveFunctionAddress, 4); - DCFlushRange((void *) patchedFunction->realEffectiveFunctionAddress, 4); + // Map Restore into the exact same Task structure for the generic worker + std::vector tasks; + WorkerTask task = {}; + task.targetPhys = targetAddrPhys; + task.sourcePhys = sourceAddrPhys; + task.effectiveAddr = patchedFunction->realEffectiveFunctionAddress; + tasks.push_back(task); + + DCFlushRange(tasks.data(), sizeof(task) * tasks.size()); + DCFlushRange(&tasks, sizeof(tasks)); + OSMemoryBarrier(); + + if (!sSpinLockInitialized) { + OSInitSpinLock(&sGlobalSpinLock); + sSpinLockInitialized = true; + } + + gPatchState = PatchState::PREPARE; + gCoresReady = 0; + gCoresFlushed = 0; + + CThread *threadA = CThread::create(applyKernelPatchOnCore, &tasks, CThread::eAttributeAffCore2, 0, 0x1000); + CThread *threadB = CThread::create(applyKernelPatchOnCore, &tasks, CThread::eAttributeAffCore0, 0, 0x1000); + threadA->resumeThread(); + threadB->resumeThread(); + + while (gCoresReady < 2) { + OSSleepTicks(1); + } + + DEBUG_FUNCTION_LINE_VERBOSE("Restore on thread for %08X on Core %d", patchedFunction->realPhysicalFunctionAddress, OSGetCoreId()); + + OSUninterruptibleSpinLock_Acquire(&sGlobalSpinLock); + + DCFlushRange(tasks.data(), tasks.size() * sizeof(WorkerTask)); + KernelCopyData(task.targetPhys, task.sourcePhys, 4); + OSMemoryBarrier(); + + DCFlushRange((void *) task.effectiveAddr, 4); + + gPatchState = PatchState::MAIN_CORE_PATCHING_DONE; + + ICInvalidateRange((void *) task.effectiveAddr, 4); + asm volatile("sync; isync"); + + while (gCoresFlushed < 2) { asm volatile("nop");} + + OSUninterruptibleSpinLock_Release(&sGlobalSpinLock); + gPatchState = PatchState::ALL_CORES_PATCHING_DONE; + + delete threadA; + delete threadB; patchedFunction->isPatched = false; return true; diff --git a/source/function_patcher.h b/source/function_patcher.h index 3cede57..1a01eec 100644 --- a/source/function_patcher.h +++ b/source/function_patcher.h @@ -1,8 +1,6 @@ #pragma once #include "PatchedFunctionData.h" -#include -#include #include #include @@ -13,6 +11,8 @@ extern "C" { bool PatchFunction(std::shared_ptr &patchedFunction); bool RestoreFunction(std::shared_ptr &patchedFunction); +bool PatchFunctions(std::vector> &patchedFunctions); + #ifdef __cplusplus } #endif \ No newline at end of file diff --git a/source/main.cpp b/source/main.cpp index 886140c..2615232 100644 --- a/source/main.cpp +++ b/source/main.cpp @@ -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..11c4df9 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("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