diff --git a/wumsloader/src/entry.cpp b/wumsloader/src/entry.cpp index 800ad07..7a47b1d 100644 --- a/wumsloader/src/entry.cpp +++ b/wumsloader/src/entry.cpp @@ -52,7 +52,7 @@ extern "C" int _start(int argc, char **argv) { static uint8_t ucSetupRequired = 1; if (ucSetupRequired) { const auto heapSpan = getMemoryForHeap(); - gHeapHandle = MEMCreateExpHeapEx(heapSpan.data(), heapSpan.size_bytes(), 1); + gHeapHandle = MEMCreateExpHeapEx(heapSpan.data(), heapSpan.size_bytes(), MEM_HEAP_FLAG_USE_LOCK); if (!gHeapHandle) { OSFatal("Failed to alloc heap"); } diff --git a/wumsloader/src/module/RelocationUtils.cpp b/wumsloader/src/module/RelocationUtils.cpp index 1a32636..3a72199 100644 --- a/wumsloader/src/module/RelocationUtils.cpp +++ b/wumsloader/src/module/RelocationUtils.cpp @@ -47,12 +47,13 @@ OSDynLoad_Error CustomDynLoadAlloc(int32_t size, int32_t align, void **outAddr) } void CustomDynLoadFree(void *addr) { - free(addr); - // Remove from list if (const auto it = std::ranges::find(gAllocatedAddresses, addr); it != gAllocatedAddresses.end()) { gAllocatedAddresses.erase(it); } + + free(addr); + } bool doRelocation(const std::vector &relocData, diff --git a/wumsloader/src/utils/reent.cpp b/wumsloader/src/utils/reent.cpp index 6a699c0..30b91f6 100644 --- a/wumsloader/src/utils/reent.cpp +++ b/wumsloader/src/utils/reent.cpp @@ -8,7 +8,8 @@ #include #include -#include +#include +#include #include #include #include @@ -30,6 +31,7 @@ void wums_set_thread_specific_ex(const int id, void *value, OSThread *thread) { OSReport("wums_set_thread_specific: invalid thread\n"); OSFatal("wums_set_thread_specific: invalid thread"); } + OSMemoryBarrier(); } void *wums_get_thread_specific_ex(const int id, const OSThread *thread) { @@ -66,93 +68,73 @@ struct __wums_reent_node { OSThreadCleanupCallbackFn savedCleanup; }; + namespace { -std::unordered_set<__wums_reent_node *> sGlobalNodesSeen; + +class MutexWrapper { +public: + MutexWrapper() = default; + + void init(const char *name) { + OSInitMutexEx(&mutex, name); + } + + void lock() { + OSLockMutex(&mutex); + } + + void unlock() { + OSUnlockMutex(&mutex); + } + +private: + OSMutex mutex{}; +}; + +MutexWrapper gMutexWrapperReent; + +std::vector<__wums_reent_node *> sGlobalNodesCopy; std::vector<__wums_reent_node *> sGlobalNodes; -std::recursive_mutex sGlobalNodesMutex; void removeNodeFromListsSafe(__wums_reent_node *curr) { - std::lock_guard lock(sGlobalNodesMutex); + std::lock_guard lock(gMutexWrapperReent); if (const auto it = std::ranges::find(sGlobalNodes, curr); it != sGlobalNodes.end()) { *it = sGlobalNodes.back(); sGlobalNodes.pop_back(); } - sGlobalNodesSeen.erase(curr); } } // namespace +/* + * This function is expected to be called exactly once at the start of each new application cycle. + * It acts as a garbage collector for nodes left behind by the previous application. + * Leftover nodes typically occur when threads are forcefully killed before they can execute + * their cleanup callbacks, or if a thread's cleanup function was wrongly overridden. + * + * Mechanism: Since threads do not survive across application boundaries, any node we + * observe across multiple cycles is guaranteed to be a dangling pointer from a dead thread. + */ void ClearDanglingReentPtr() { - // This function is expected to be called exactly once at the start of each new application cycle. - // It acts as a garbage collector for nodes left behind by the previous application. - // Leftover nodes typically occur when threads are forcefully killed before they can execute - // their cleanup callbacks, or if a thread's cleanup function was wrongly overridden. - // - // Mechanism: Since threads do not survive across application boundaries, any node we - // observe across multiple cycles is guaranteed to be a dangling pointer from a dead thread. - std::vector<__wums_reent_node *> snapshot; + std::lock_guard lock(gMutexWrapperReent); - // Snapshot Phase - // - // We instantly move the global state to a local variable. - // This protects us from iterator invalidation and allocator hook re-entrancy. - // If a plugin registers a new context while we are iterating, it will safely - // push to the newly emptied global vector without interrupting us. { - std::lock_guard lock(sGlobalNodesMutex); - snapshot = std::move(sGlobalNodes); - sGlobalNodes.clear(); - } - - std::vector<__wums_reent_node *> nodesToFree; - std::vector<__wums_reent_node *> survivingNodes; - - // Iterate over our isolated snapshot - for (auto *ptr : snapshot) { - if (!ptr) continue; - - bool isDangling = false; - - { - std::lock_guard lock(sGlobalNodesMutex); - - if (auto it = sGlobalNodesSeen.find(ptr); it != sGlobalNodesSeen.end()) { - // We have already seen this address in a previous cycle. - // This means the node belongs to a dead thread from an older application. - isDangling = true; - sGlobalNodesSeen.erase(it); - } else { - // If it wasn't found, it might be a valid node created during the current - // application cycle (e.g., initialized via a hook in an RPL's init function). - // We add it to the seen list to check against in the NEXT cycle. - sGlobalNodesSeen.insert(ptr); + // Free everything that's currently in *sGlobalNodesCopy* + auto *oldHead = wums_backend_set_sentinel(); + for (const auto nodeToFree : sGlobalNodesCopy) { + if (nodeToFree->cleanupFn) { + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Call cleanupFn(%p) for node %p (dangling)", OSGetCurrentThread(), nodeToFree->reentPtr, nodeToFree); + nodeToFree->cleanupFn(nodeToFree->reentPtr); } + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Free node %p (dangling)", OSGetCurrentThread(), nodeToFree); + free(nodeToFree); } - - if (isDangling) { - nodesToFree.push_back(ptr); - } else { - survivingNodes.push_back(ptr); - } + sGlobalNodesCopy.clear(); + wums_backend_restore_head(oldHead); } - // Merge the surviving nodes back into the global pool - { - std::lock_guard lock(sGlobalNodesMutex); - // We append instead of overwriting, just in case a hook successfully - // registered a brand-new node into the empty global list while we were processing. - sGlobalNodes.insert(sGlobalNodes.end(), survivingNodes.begin(), survivingNodes.end()); - } - - // It is now safe to execute payload cleanups and free the memory. - // This adds nodes to the global list so we couldn't do is earlier - for (auto *nodeToFree : nodesToFree) { - if (nodeToFree->cleanupFn) { - nodeToFree->cleanupFn(nodeToFree->reentPtr); - } - free(nodeToFree); - } - - DEBUG_FUNCTION_LINE_INFO("Cleaned up %d dangling reent entries.", nodesToFree.size()); + // Then move the node current list into sGlobalNodesCopy. + sGlobalNodesCopy = std::move(sGlobalNodes); + sGlobalNodes.clear(); } static void __wums_thread_cleanup(OSThread *thread, void *stack) { @@ -170,14 +152,18 @@ static void __wums_thread_cleanup(OSThread *thread, void *stack) { auto *curr = head; while (curr) { __wums_reent_node *next = curr->next; + if (curr->magic == WUMS_REENT_NODE_MAGIC && curr->version >= 1) { + if (curr->cleanupFn) { + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Call cleanupFn(%p) for node %p", thread, curr->reentPtr, curr); + curr->cleanupFn(curr->reentPtr); + } - if (curr->cleanupFn) { - curr->cleanupFn(curr->reentPtr); + removeNodeFromListsSafe(curr); + + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Free node %p", thread, curr); + free(curr); } - removeNodeFromListsSafe(curr); - - free(curr); curr = next; } @@ -185,6 +171,7 @@ static void __wums_thread_cleanup(OSThread *thread, void *stack) { // Chain to previous OS callback if (savedCleanup) { + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Call saved cleanup function %p", thread, savedCleanup); savedCleanup(thread, stack); } } @@ -225,12 +212,14 @@ void *wums_backend_get_context(const void *moduleId, wums_loader_init_reent_erro } void *wums_backend_set_sentinel() { + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Set sentinel", OSGetCurrentThread()); auto *head = wums_get_thread_specific(__WUMS_CONTEXT_THREAD_SPECIFIC_ID); wums_set_thread_specific(__WUMS_CONTEXT_THREAD_SPECIFIC_ID, WUMS_REENT_ALLOC_SENTINEL); return head; } void wums_backend_restore_head(void *oldHead) { + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Set head to %p", OSGetCurrentThread(), oldHead); wums_set_thread_specific(__WUMS_CONTEXT_THREAD_SPECIFIC_ID, oldHead); } @@ -251,128 +240,22 @@ bool wums_backend_register_context(const void *moduleId, void *reentPtr, void (* newNode->savedCleanup = nullptr; if (oldHead == nullptr || oldHead == WUMS_REENT_ALLOC_SENTINEL || oldHead->magic != WUMS_REENT_NODE_MAGIC || oldHead->version < WUMS_REENT_NODE_VERSION) { + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Set OSSetThreadCleanupCallback for node %p", OSGetCurrentThread(), newNode); newNode->savedCleanup = OSSetThreadCleanupCallback(OSGetCurrentThread(), &__wums_thread_cleanup); } else { + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Add to existing cleanup chain for node %p", OSGetCurrentThread(), newNode); newNode->savedCleanup = oldHead->savedCleanup; oldHead->savedCleanup = nullptr; } + OSMemoryBarrier(); wums_set_thread_specific(__WUMS_CONTEXT_THREAD_SPECIFIC_ID, newNode); { - std::lock_guard lock(sGlobalNodesMutex); + std::lock_guard lock(gMutexWrapperReent); sGlobalNodes.push_back(newNode); + DEBUG_FUNCTION_LINE_VERBOSE("[%p] Registered reent ptr %p as node %p", OSGetCurrentThread(), reentPtr, newNode); } return true; -} - -/* -void ClearReentDataForModules(const std::vector &modules) { - auto *curThread = OSGetCurrentThread(); - - for (const auto &cur : modules) { - if (!cur.isLinkedAndLoaded()) { - continue; - } - - if (cur.getMetaInformation().getWUMSVersion() <= WUMSLoader::Modules::WUMSVersion(0, 3, 5)) { - continue; - } - - const auto startAddress = cur.getLinkInformation().getStartAddress(); - const auto endAddress = cur.getLinkInformation().getEndAddress(); - - // Zero-allocation deferred free list - __wums_reent_node *deferredFreeHead = nullptr; - - struct PendingRestore { - OSThread *thread; - OSThreadCleanupCallbackFn callback; - }; - std::vector pendingRestores; - constexpr int PENDING_RESTORES_SIZE = 128; - // Pre-allocate to prevent malloc() from firing while the scheduler is locked! - pendingRestores.reserve(PENDING_RESTORES_SIZE); - { - // Acquire GLOBAL scheduler lock - const int state = OSDisableInterrupts(); - __OSLockScheduler(curThread); - - OSThread *t = *reinterpret_cast(0x100567F8); - - while (t) { - auto *head = static_cast<__wums_reent_node *>(wums_get_thread_specific_ex(__WUMS_CONTEXT_THREAD_SPECIFIC_ID, t)); - - // Safety checks with Sentinel/Magic - if (!head || head == WUMS_REENT_ALLOC_SENTINEL || head->magic != WUMS_REENT_NODE_MAGIC) { - t = t->activeLink.next; - continue; - } - - __wums_reent_node *prev = nullptr; - auto *curr = head; - - while (curr) { - __wums_reent_node *next = curr->next; - auto moduleIdAddr = reinterpret_cast(curr->moduleId); - - // module id lives in the .data/.bss section of a module - if (moduleIdAddr >= startAddress && moduleIdAddr < endAddress) { - // remove from linked list - if (prev) { - prev->next = next; - } else { - head = next; - if (curr->savedCleanup) { - if (head) { - head->savedCleanup = curr->savedCleanup; - } else { - // No WUMS nodes left, mark for restoring. - if (pendingRestores.size() == PENDING_RESTORES_SIZE) { - OSFatal("WUMSBackend pendingRestores size limit hit"); - } - pendingRestores.push_back({t, curr->savedCleanup}); - } - } - } - - curr->next = deferredFreeHead; - deferredFreeHead = curr; - } else { - prev = curr; - } - curr = next; - } - - // Restore the updated head to the thread - wums_set_thread_specific_ex(__WUMS_CONTEXT_THREAD_SPECIFIC_ID, head, t); - t = t->activeLink.next; - } - - __OSUnlockScheduler(curThread); - OSRestoreInterrupts(state); - } - - for (const auto &restore : pendingRestores) { - DEBUG_FUNCTION_LINE_VERBOSE("Set cleanup function for thread %p to %p", restore.thread, restore.callback); - OSSetThreadCleanupCallback(restore.thread, restore.callback); - } - - // Free removed entries - auto *oldHead = wums_backend_set_sentinel(); - __wums_reent_node *nodeToFree = deferredFreeHead; - while (nodeToFree) { - __wums_reent_node *nextNode = nodeToFree->next; - if (nodeToFree->cleanupFn) { - nodeToFree->cleanupFn(nodeToFree->reentPtr); - } - removeNodeFromListsLocked(nodeToFree); - free(nodeToFree); - - nodeToFree = nextNode; - } - wums_backend_restore_head(oldHead); - } -} -*/ \ No newline at end of file +} \ No newline at end of file