Check if the MEM2 heap is corrupted after calling a hook

wiiu-env · Apr 20, 2024 · 398be01 · 398be01
1 parent 1ac38cb
commit 398be01
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/wumsloader/src/utils/hooks.cpp b/wumsloader/src/utils/hooks.cpp
@@ -5,7 +5,6 @@
 #include <memory>
 #include <wums.h>
 
-#ifdef DEBUG
 static const char **hook_names = (const char *[]){
         "WUMS_HOOK_INIT_WUT_MALLOC",
         "WUMS_HOOK_FINI_WUT_MALLOC",
@@ -27,7 +26,6 @@ static const char **hook_names = (const char *[]){
         "WUMS_HOOK_RELOCATIONS_DONE",
         "WUMS_HOOK_APPLICATION_REQUESTS_EXIT",
         "WUMS_HOOK_DEINIT"};
-#endif
 
 void CallHook(const std::vector<std::shared_ptr<ModuleData>> &modules, wums_hook_type_t type, bool condition) {
     if (condition) {
@@ -48,7 +46,10 @@ void CallHook(const std::shared_ptr<ModuleData> &module, wums_hook_type_t type,
     }
 }
 
+extern "C" bool MEMCheckExpHeap(void *heap, bool logProblems);
+
 void CallHook(const std::shared_ptr<ModuleData> &module, wums_hook_type_t type) {
+    bool foundHook = false;
     for (auto &curHook : module->getHookDataList()) {
         auto func_ptr = (uint32_t) curHook->getTarget();
         if (func_ptr == 0) {
@@ -57,6 +58,7 @@ void CallHook(const std::shared_ptr<ModuleData> &module, wums_hook_type_t type)
         }
 
         if (type == curHook->getType()) {
+            foundHook = true;
             if ((type == WUMS_HOOK_APPLICATION_STARTS ||
                  type == WUMS_HOOK_APPLICATION_ENDS ||
                  type == WUMS_HOOK_INIT_WUT_MALLOC ||
@@ -90,4 +92,11 @@ void CallHook(const std::shared_ptr<ModuleData> &module, wums_hook_type_t type)
             break;
         }
     }
+
+    if (foundHook && !MEMCheckExpHeap(MEMGetBaseHeapHandle(MEM_BASE_HEAP_MEM2), true)) {
+        DEBUG_FUNCTION_LINE_ERR("MEM2 default heap is corrupted while calling hook %s for module %s", hook_names[type], module->getExportName().c_str());
+#ifdef DEBUG
+        OSFatal("WUMSLoader: MEM2 default heap is corrupted. \n Please restart the console.");
+#endif
+    }
 }
diff --git a/wumsloader/src/utils/imports.h b/wumsloader/src/utils/imports.h
@@ -40,6 +40,9 @@ IMPORT(OSGetCurrentThread);
 IMPORT(OSSetThreadCleanupCallback);
 IMPORT(OSIsDebuggerPresent);
 IMPORT(__os_snprintf);
+IMPORT(MEMCheckHeap);
+IMPORT(MEMGetBaseHeapHandle);
+IMPORT(MEMCheckExpHeap);
 
 IMPORT(FSTimeToCalendarTime);
 IMPORT(FSInit);