Make the reent implementation even more robust

This commit is contained in:
Maschell 2026-04-13 16:10:01 +02:00
parent 47e3080b30
commit 6d1645c1f9
3 changed files with 74 additions and 190 deletions

View File

@ -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");
}

View File

@ -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<RelocationData> &relocData,

View File

@ -8,7 +8,8 @@
#include <coreinit/debug.h>
#include <coreinit/thread.h>
#include <coreinit/memexpheap.h>
#include <coreinit/cache.h>
#include <coreinit/mutex.h>
#include <cstdlib>
#include <mutex>
#include <vector>
@ -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<WUMSLoader::Modules::ModuleContainer> &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<PendingRestore> 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<OSThread **>(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<uint32_t>(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);
}
}
*/
}