Use a custom heap for the relocator and modules

relocator/src/entry.cpp

@@ -7,6 +7,7 @@
 #include <nsysnet/socket.h>
 #include <map>
 #include <algorithm>
+#include <coreinit/memexpheap.h>
 #include "../../source/module/RelocationData.h"
 #include "../../source/module/ModuleData.h"
 #include "ModuleDataPersistence.h"
@@ -14,7 +15,9 @@
 #include "utils/dynamic.h"
 #include "globals.h"
 #include "hooks.h"
+#include "utils/memory.h"
 
+MEMHeapHandle gHeapHandle __attribute__((section(".data"))) = nullptr;
 uint8_t gFunctionsPatched __attribute__((section(".data"))) = 0;
 uint8_t gInitCalled __attribute__((section(".data"))) = 0;
 
@@ -25,6 +28,13 @@ extern "C" void doStart(int argc, char **argv);
 // The compiler tries to optimize this otherwise and calling the main function earlier
 extern "C" int _start(int argc, char **argv) {
     InitFunctionPointers();
+
+    static uint8_t ucSetupRequired = 1;
+    if (ucSetupRequired) {
+        gHeapHandle = MEMCreateExpHeapEx((void *) (MEMORY_REGION_USABLE_HEAP_START), MEMORY_REGION_USABLE_HEAP_END - MEMORY_REGION_USABLE_HEAP_START, 0);
+        ucSetupRequired = 0;
+    }
+
     socket_lib_init();
     log_init();
 
@@ -34,7 +44,7 @@ extern "C" int _start(int argc, char **argv) {
     return ((int (*)(int, char **)) (*(unsigned int *) 0x1005E040))(argc, argv);
 }
 
-bool doRelocation(std::vector<RelocationData> &relocData, relocation_trampolin_entry_t *tramp_data, uint32_t tramp_length) {
+bool doRelocation(std::vector<RelocationData> &relocData, relocation_trampolin_entry_t *tramp_data, uint32_t tramp_length, bool skipAllocReplacement) {
     std::map<std::string, OSDynLoad_Module> moduleCache;
     for (auto const &curReloc : relocData) {
         std::string functionName = curReloc.getName();
@@ -52,16 +62,26 @@ bool doRelocation(std::vector<RelocationData> &relocData, relocation_trampolin_e
             }
         }
 
+        if (!skipAllocReplacement) {
+            if (functionName == "MEMAllocFromDefaultHeap") {
+                functionAddress = reinterpret_cast<uint32_t>(&MEMAlloc);
+            } else if (functionName == "MEMAllocFromDefaultHeapEx") {
+                functionAddress = reinterpret_cast<uint32_t>(&MEMAllocEx);
+            } else if (functionName == "MEMFreeToDefaultHeap") {
+                functionAddress = reinterpret_cast<uint32_t>(&MEMFree);
+            }
+        }
+
         if (functionAddress == 0) {
             int32_t isData = curReloc.getImportRPLInformation().isData();
             OSDynLoad_Module rplHandle = nullptr;
             if (moduleCache.count(rplName) == 0) {
                 OSDynLoad_Error err = OSDynLoad_IsModuleLoaded(rplName.c_str(), &rplHandle);
-                if(err != OS_DYNLOAD_OK || rplHandle == nullptr) {
+                if (err != OS_DYNLOAD_OK || rplHandle == nullptr) {
                     DEBUG_FUNCTION_LINE("%s is not yet loaded\n", rplName.c_str());
                     // only acquire if not already loaded.
                     err = OSDynLoad_Acquire(rplName.c_str(), &rplHandle);
-                    if(err != OS_DYNLOAD_OK){
+                    if (err != OS_DYNLOAD_OK) {
                         DEBUG_FUNCTION_LINE("Failed to acquire %s\n", rplName.c_str());
                         //return false;
                     }
@@ -87,14 +107,21 @@ bool doRelocation(std::vector<RelocationData> &relocData, relocation_trampolin_e
     return true;
 }
 
-bool ResolveRelocations(std::vector<ModuleData> &loadedModules) {
+bool ResolveRelocations(std::vector<ModuleData> &loadedModules, bool skipMemoryMappingModule) {
     bool wasSuccessful = true;
 
     for (auto &curModule : loadedModules) {
         DEBUG_FUNCTION_LINE("Let's do the relocations for %s\n", curModule.getExportName().c_str());
         if (wasSuccessful) {
             std::vector<RelocationData> relocData = curModule.getRelocationDataList();
-            if (!doRelocation(relocData, gModuleData->trampolines, DYN_LINK_TRAMPOLIN_LIST_LENGTH)) {
+
+            // On first usage we can't redirect the alloc functions to our custom heap
+            // because threads can't run it on it. In order to patch the kernel
+            // to fully support our memory region, we have to run the FunctionPatcher and MemoryMapping
+            // once with the default heap. Afterwards we can just rely on the custom heap.
+            bool skipAllocFunction = skipMemoryMappingModule && (curModule.getExportName() == "homebrew_memorymapping" || curModule.getExportName() == "homebrew_functionpatcher");
+            DEBUG_FUNCTION_LINE("Skip alloc replace? %d\n", skipAllocFunction);
+            if (!doRelocation(relocData, gModuleData->trampolines, DYN_LINK_TRAMPOLIN_LIST_LENGTH, skipAllocFunction)) {
                 DEBUG_FUNCTION_LINE("FAIL\n");
                 wasSuccessful = false;
                 curModule.relocationsDone = false;
@@ -124,7 +151,6 @@ extern "C" void doStart(int argc, char **argv) {
     std::vector<ModuleData> loadedModulesUnordered = ModuleDataPersistence::loadModuleData(gModuleData);
     std::vector<ModuleData> loadedModules = OrderModulesByDependencies(loadedModulesUnordered);
 
-
     bool applicationEndHookLoaded = false;
     for (auto &curModule : loadedModules) {
         if (curModule.getExportName() == "homebrew_applicationendshook") {
@@ -150,7 +176,7 @@ extern "C" void doStart(int argc, char **argv) {
         gInitCalled = 1;
 
         DEBUG_FUNCTION_LINE("Resolve relocations without replacing alloc functions\n");
-        ResolveRelocations(loadedModules);
+        ResolveRelocations(loadedModules, true);
 
         for (auto &curModule : loadedModules) {
             if (curModule.isInitBeforeRelocationDoneHook()) {
@@ -181,7 +207,7 @@ extern "C" void doStart(int argc, char **argv) {
         }
     } else {
         DEBUG_FUNCTION_LINE("Resolve relocations and replace alloc functions\n");
-        ResolveRelocations(loadedModules);
+        ResolveRelocations(loadedModules, false);
         CallHook(loadedModules, WUMS_HOOK_RELOCATIONS_DONE);
     }
     CallHook(loadedModules, WUMS_HOOK_INIT_WUT);

relocator/src/utils/dynamic.c

@@ -15,11 +15,6 @@
 #define IMPORT_BEGIN(lib)  do{if(OSDynLoad_IsModuleLoaded(#lib ".rpl", &handle) != OS_DYNLOAD_OK) OSFatal(#lib ".rpl is not loaded");} while(0)
 #define IMPORT_END()
 
-#define EXPORT_VAR(type, var)           type var __attribute__((section(".data")));
-
-EXPORT_VAR(uint32_t *, pMEMAllocFromDefaultHeapEx);
-EXPORT_VAR(uint32_t *, pMEMAllocFromDefaultHeap);
-EXPORT_VAR(uint32_t *, pMEMFreeToDefaultHeap);
 
 void InitFunctionPointers(void) {
     OSDynLoad_Module handle;
@@ -27,12 +22,6 @@ void InitFunctionPointers(void) {
     addr_OSDynLoad_FindExport = (void *) 0x0102B828; // 0200f428 - 0xFE3C00
     addr_OSDynLoad_IsModuleLoaded = (void *) 0x0102A59C; // 0200e19c - 0xFE3C00
 
-    IMPORT_BEGIN(coreinit);
-    OSDynLoad_FindExport(handle, 1, "MEMAllocFromDefaultHeapEx", (void **) &pMEMAllocFromDefaultHeapEx);
-    OSDynLoad_FindExport(handle, 1, "MEMAllocFromDefaultHeap", (void **) &pMEMAllocFromDefaultHeap);
-    OSDynLoad_FindExport(handle, 1, "MEMFreeToDefaultHeap", (void **) &pMEMFreeToDefaultHeap);
-    IMPORT_END()
-
 #include "imports.h"
 
 }

relocator/src/utils/imports.h

@@ -11,6 +11,12 @@ IMPORT(MEMGetSizeForMBlockExpHeap);
 IMPORT(OSSleepTicks);
 IMPORT(OSEffectiveToPhysical);
 IMPORT(OSFatal);
+IMPORT(MEMFreeToExpHeap);
+IMPORT(MEMAllocFromExpHeapEx);
+IMPORT(MEMGetAllocatableSizeForExpHeapEx);
+IMPORT(MEMCreateExpHeapEx);
+IMPORT(OSUninterruptibleSpinLock_Acquire);
+IMPORT(OSUninterruptibleSpinLock_Release);
 
 IMPORT_END();
 

relocator/src/utils/memory.cpp

@@ -17,19 +17,69 @@
 #include <coreinit/memexpheap.h>
 #include <coreinit/memdefaultheap.h>
 #include <coreinit/memorymap.h>
+#include <coreinit/cache.h>
 #include <malloc.h>
 #include <string.h>
 #include <errno.h>
+#include "logger.h"
 
-extern uint32_t *pMEMAllocFromDefaultHeapEx;
+extern MEMHeapHandle gHeapHandle;
-extern uint32_t *pMEMAllocFromDefaultHeap;
+
-extern uint32_t *pMEMFreeToDefaultHeap;
+void *MEMAllocSafe(uint32_t size, uint32_t align) {
+    void *res = nullptr;
+    MEMHeapHandle heapHandle = gHeapHandle;
+    MEMExpHeap *heap = (MEMExpHeap *) heapHandle;
+    OSUninterruptibleSpinLock_Acquire(&heap->header.lock);
+    res = MEMAllocFromExpHeapEx(heapHandle, size, align);
+    auto cur = heap->usedList.head;
+    while (cur != nullptr) {
+        DCFlushRange(cur, sizeof(MEMExpHeapBlock));
+        cur = cur->next;
+    }
+    cur = heap->freeList.head;
+    while (cur != nullptr) {
+        DCFlushRange(cur, sizeof(MEMExpHeapBlock));
+        cur = cur->next;
+    }
+    OSUninterruptibleSpinLock_Release(&heap->header.lock);
+
+    return res;
+}
+
+
+void *MemoryAlloc(uint32_t size) {
+    void *res = MEMAllocSafe(size, 4);
+    if (res == nullptr) {
+        OSFatal_printf("Failed to MemoryAlloc %d", size);
+    }
+    return res;
+}
+
+void *MemoryAllocEx(uint32_t size, uint32_t align) {
+    void *res = MEMAllocSafe(size, align);
+    if (res == nullptr) {
+        OSFatal_printf("Failed to MemoryAllocEX %d %d", size, align);
+    }
+    return res;
+}
+
+void MemoryFree(void *ptr) {
+    if (ptr) {
+        MEMFreeToExpHeap(gHeapHandle, ptr);
+    } else {
+        OSFatal_printf("Failed to free");
+    }
+}
+
+uint32_t MEMAlloc __attribute__((__section__ (".data"))) = (uint32_t) MemoryAlloc;
+uint32_t MEMAllocEx __attribute__((__section__ (".data"))) = (uint32_t) MemoryAllocEx;
+uint32_t MEMFree __attribute__((__section__ (".data"))) = (uint32_t) MemoryFree;
 
 //!-------------------------------------------------------------------------------------------
 //! reent versions
 //!-------------------------------------------------------------------------------------------
 void *_malloc_r(struct _reent *r, size_t size) {
-    void *ptr = ((void *(*)(size_t)) (*pMEMAllocFromDefaultHeap))(size);
+    void *ptr = MemoryAllocEx(size, 4);
     if (!ptr) {
         r->_errno = ENOMEM;
     }
@@ -37,7 +87,7 @@ void *_malloc_r(struct _reent *r, size_t size) {
 }
 
 void *_calloc_r(struct _reent *r, size_t num, size_t size) {
-    void *ptr = ((void *(*)(size_t)) (*pMEMAllocFromDefaultHeap))(size);
+    void *ptr = MemoryAllocEx(num * size, 4);
     if (ptr) {
         memset(ptr, 0, num * size);
     } else {
@@ -48,17 +98,17 @@ void *_calloc_r(struct _reent *r, size_t num, size_t size) {
 }
 
 void *_memalign_r(struct _reent *r, size_t align, size_t size) {
-    return ((void *(*)(size_t, size_t)) (*pMEMAllocFromDefaultHeapEx))(size, align);
+    return MemoryAllocEx(size, align);
 }
 
 void _free_r(struct _reent *r, void *ptr) {
     if (ptr) {
-        ((void (*)(void *)) (*pMEMFreeToDefaultHeap))(ptr);
+        MemoryFree(ptr);
     }
 }
 
 void *_realloc_r(struct _reent *r, void *p, size_t size) {
-    void *new_ptr = ((void *(*)(size_t)) (*pMEMAllocFromDefaultHeap))(size);
+    void *new_ptr = MemoryAllocEx(size, 4);
     if (!new_ptr) {
         r->_errno = ENOMEM;
         return new_ptr;
@@ -67,7 +117,7 @@ void *_realloc_r(struct _reent *r, void *p, size_t size) {
     if (p) {
         size_t old_size = MEMGetSizeForMBlockExpHeap(p);
         memcpy(new_ptr, p, old_size <= size ? old_size : size);
-        ((void (*)(void *)) (*pMEMFreeToDefaultHeap))(p);
+        MemoryFree(p);
     }
     return new_ptr;
 }
@@ -93,12 +143,12 @@ _malloc_usable_size_r(struct _reent *r, void *ptr) {
 
 void *
 _valloc_r(struct _reent *r, size_t size) {
-    return ((void *(*)(size_t, size_t)) (*pMEMAllocFromDefaultHeapEx))(size, OS_PAGE_SIZE);
+    return MemoryAllocEx(size, OS_PAGE_SIZE);
 }
 
 void *
 _pvalloc_r(struct _reent *r, size_t size) {
-    return ((void *(*)(size_t, size_t)) (*pMEMAllocFromDefaultHeapEx))((size + (OS_PAGE_SIZE - 1)) & ~(OS_PAGE_SIZE - 1), OS_PAGE_SIZE);
+    return MemoryAllocEx((size + (OS_PAGE_SIZE - 1)) & ~(OS_PAGE_SIZE - 1), OS_PAGE_SIZE);
 }
 
 int

relocator/src/utils/memory.h

@@ -1,46 +1,5 @@
-/****************************************************************************
+#pragma once
- * Copyright (C) 2015 Dimok
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation, either version 3 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program.  If not, see <http://www.gnu.org/licenses/>.
- ****************************************************************************/
-#ifndef __MEMORY_H_
-#define __MEMORY_H_
 
-#ifdef __cplusplus
+extern uint32_t MEMAlloc;
-extern "C" {
+extern uint32_t MEMAllocEx;
-#endif
+extern uint32_t MEMFree;
-
-#include <malloc.h>
-
-void memoryInitialize(void);
-
-void memoryRelease(void);
-
-void *MEM2_alloc(u32 size, u32 align);
-
-void MEM2_free(void *ptr);
-
-void *MEM1_alloc(u32 size, u32 align);
-
-void MEM1_free(void *ptr);
-
-void *MEMBucket_alloc(u32 size, u32 align);
-
-void MEMBucket_free(void *ptr);
-
-#ifdef __cplusplus
-}
-#endif
-
-#endif // __MEMORY_H_

source/globals.h

@@ -5,7 +5,10 @@
 #define MEMORY_REGION_START	        0x00800000
 #define MEMORY_REGION_SIZE          0x00800000
 
-#define MEMORY_REGION_USABLE_START  MEMORY_REGION_START + 0x00080000
+#define MEMORY_REGION_USABLE_HEAP_START  (MEMORY_REGION_START + 0x00080000)
+#define MEMORY_REGION_USABLE_HEAP_END    (MEMORY_REGION_USABLE_HEAP_START + 0x00100000)
+
+#define MEMORY_REGION_USABLE_START  MEMORY_REGION_USABLE_HEAP_END
 #define MEMORY_REGION_USABLE_END    0x00FFF000
 
 #define gModuleData ((module_information_t *) (MEMORY_REGION_USABLE_START))