pm/romfs: first (broken?) pass at dynamic heap.

I cannot wait to figure out all the ways this is wrong.

stratosphere/ams_mitm/ams_mitm.json

@@ -65,6 +65,8 @@
                 "svcReplyAndReceive":	"0x43",
                 "svcReplyAndReceiveWithUserBuffer":	"0x44",
                 "svcCreateEvent":	"0x45",
+				"svcMapPhysicalMemoryUnsafe":	"0x48",
+				"svcUnmapPhysicalMemoryUnsafe":	"0x49",
                 "svcMapTransferMemory":	"0x51",
                 "svcUnmapTransferMemory":	"0x52",
                 "svcCreateInterruptEvent":	"0x53",

stratosphere/ams_mitm/source/fs_mitm/fsmitm_layered_romfs_storage.cpp

@@ -78,6 +78,7 @@ namespace ams::mitm::fs {
 
         constinit os::SdkRecursiveMutex g_storage_set_mutex;
         constinit LayeredRomfsStorageSet g_storage_set;
+        constinit os::SdkMutex g_initialization_mutex;
 
         void OpenReference(LayeredRomfsStorageImpl *impl) {
             std::scoped_lock lk(g_storage_set_mutex);
@@ -106,6 +107,8 @@ namespace ams::mitm::fs {
                 auto *impl = reinterpret_cast<LayeredRomfsStorageImpl *>(storage_uptr);
                 g_ack_mq.Send(storage_uptr);
 
+                std::scoped_lock lk(g_initialization_mutex);
+
                 impl->InitializeImpl();
 
                 /* Close the initial reference. */
@@ -255,6 +258,21 @@ namespace ams::mitm::fs {
         return std::make_shared<LayeredRomfsStorage>(impl);
     }
 
+    void FinalizeLayeredRomfsStorage(ncm::ProgramId program_id) {
+        std::scoped_lock lk(g_storage_set_mutex);
+        std::scoped_lock lk2(g_initialization_mutex);
+
+        /* Find an existing storage. */
+        if (auto it = g_storage_set.find_key(program_id.value); it != g_storage_set.end()) {
+            /* We need to delete the process romfs. Require invariant that it is unreferenced, by this point. */
+            AMS_ABORT_UNLESS(it->GetReferenceCount() == 0);
+
+            auto *holder = std::addressof(*it);
+            it = g_storage_set.erase(it);
+            delete holder;
+        }
+    }
+
     LayeredRomfsStorageImpl::LayeredRomfsStorageImpl(std::unique_ptr<IStorage> s_r, std::unique_ptr<IStorage> f_r, ncm::ProgramId pr_id) : m_storage_romfs(std::move(s_r)), m_file_romfs(std::move(f_r)), m_initialize_event(os::EventClearMode_ManualClear), m_program_id(std::move(pr_id)), m_is_initialized(false), m_started_initialize(false) {
         /* ... */
     }

stratosphere/ams_mitm/source/fs_mitm/fsmitm_layered_romfs_storage.hpp

@@ -51,4 +51,6 @@ namespace ams::mitm::fs {
 
     std::shared_ptr<ams::fs::IStorage> GetLayeredRomfsStorage(ncm::ProgramId program_id, ::FsStorage &data_storage, bool is_process_romfs);
 
+    void FinalizeLayeredRomfsStorage(ncm::ProgramId program_id);
+
 }

stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.cpp

@@ -16,6 +16,7 @@
 #include <stratosphere.hpp>
 #include "../amsmitm_fs_utils.hpp"
 #include "fsmitm_romfs.hpp"
+#include "fsmitm_layered_romfs_storage.hpp"
 
 namespace ams::mitm::fs {
 
@@ -25,18 +26,61 @@ namespace ams::mitm::fs {
 
         namespace {
 
-            /* TODO: Fancy Dynamic allocation globals. */
-            constinit os::SdkMutex g_romfs_build_lock;
-            //constinit size_t g_dynamic_heap_size = 0;
+            struct ApplicationWithDynamicHeapInfo {
+                ncm::ProgramId program_id;
+                size_t dynamic_heap_size;
+            };
 
-            void InitializeDynamicHeapForBuildRomfs(ncm::ProgramId program_id) {
-                /* TODO */
-                AMS_UNUSED(program_id);
+            constexpr const ApplicationWithDynamicHeapInfo ApplicationsWithDynamicHeap[] = {
+                { 0x0100A6301214E000,  40_MB }, /* Fire Emblem: Engage. */
+            };
+
+            constexpr size_t GetDynamicHeapSize(ncm::ProgramId program_id) {
+                for (const auto &info : ApplicationsWithDynamicHeap) {
+                    if (info.program_id == program_id) {
+                        return info.dynamic_heap_size;
+                    }
+                }
+
+                return 0;
             }
 
-            void FinalizeDynamicHeapForBuildRomfs(ncm::ProgramId program_id) {
-                /* TODO */
-                AMS_UNUSED(program_id);
+            /* TODO: Fancy Dynamic allocation globals. */
+            constinit os::SdkMutex g_romfs_build_lock;
+            constinit ncm::ProgramId g_dynamic_heap_program_id{};
+            constinit size_t g_dynamic_heap_size = 0;
+
+            constinit bool g_building_from_dynamic_heap = false;
+            constinit uintptr_t g_dynamic_heap_address = 0;
+            constinit size_t g_dynamic_heap_outstanding_allocations = 0;
+
+            constinit util::TypedStorage<mem::StandardAllocator> g_dynamic_heap{};
+
+            bool IsAllocatedFromDynamicHeap(void *p) {
+                const uintptr_t address = reinterpret_cast<uintptr_t>(p);
+
+                return g_dynamic_heap_address != 0 && (g_dynamic_heap_address <= address && address < g_dynamic_heap_address + g_dynamic_heap_size);
+            }
+
+            void InitializeDynamicHeapForBuildRomfs(ncm::ProgramId program_id) {
+                if (program_id == g_dynamic_heap_program_id && g_dynamic_heap_size > 0) {
+                    /* This romfs will build out of dynamic heap. */
+                    g_building_from_dynamic_heap = true;
+
+                    if (g_dynamic_heap_address == 0) {
+                        /* Map application memory as heap. */
+                        R_ABORT_UNLESS(os::AllocateUnsafeMemory(std::addressof(g_dynamic_heap_address), g_dynamic_heap_size));
+                        AMS_ABORT_UNLESS(g_dynamic_heap_address != 0);
+
+                        /* Create exp heap. */
+                        util::ConstructAt(g_dynamic_heap, reinterpret_cast<void *>(g_dynamic_heap_address), g_dynamic_heap_size);
+                    }
+                }
+            }
+
+            void FinalizeDynamicHeapForBuildRomfs() {
+                /* We are definitely no longer building out of dynamic heap. */
+                g_building_from_dynamic_heap = false;
             }
 
         }
@@ -44,16 +88,28 @@ namespace ams::mitm::fs {
         void *AllocateTracked(AllocationType type, size_t size) {
             AMS_UNUSED(type);
 
-            /* TODO: Fancy dynamic allocation with memory stealing from application pool. */
-            return std::malloc(size);
+            if (g_building_from_dynamic_heap) {
+                void * const ret = util::GetReference(g_dynamic_heap).Allocate(size);
+                AMS_ABORT_UNLESS(ret != nullptr);
+
+                ++g_dynamic_heap_outstanding_allocations;
+
+                return ret;
+            } else {
+                return std::malloc(size);
+            }
         }
 
         void FreeTracked(AllocationType type, void *p, size_t size) {
             AMS_UNUSED(type);
             AMS_UNUSED(size);
 
-            /* TODO: Fancy dynamic allocation with memory stealing from application pool. */
-            return std::free(p);
+            if (IsAllocatedFromDynamicHeap(p)) {
+                --g_dynamic_heap_outstanding_allocations;
+                util::GetReference(g_dynamic_heap).Free(p);
+            } else {
+                std::free(p);
+            }
         }
 
         namespace {
@@ -342,7 +398,7 @@ namespace ams::mitm::fs {
 
         Builder::~Builder() {
             /* If we have nothing remaining in dynamic heap, release it. */
-            FinalizeDynamicHeapForBuildRomfs(m_program_id);
+            FinalizeDynamicHeapForBuildRomfs();
 
             /* Release the romfs build lock. */
             g_romfs_build_lock.Unlock();
@@ -801,6 +857,49 @@ namespace ams::mitm::fs {
             }
         }
 
+        Result ConfigureDynamicHeap(u64 *out_size, ncm::ProgramId program_id, const cfg::OverrideStatus &status, bool is_application) {
+            /* Baseline: use no dynamic heap. */
+            *out_size = 0;
+
+            /* If the process is not an application, we do not care about dynamic heap. */
+            R_SUCCEED_IF(!is_application);
+
+            /* First, we need to ensure that, if the game used dynamic heap, we clear it. */
+            if (g_dynamic_heap_size > 0) {
+                mitm::fs::FinalizeLayeredRomfsStorage(g_dynamic_heap_program_id);
+
+                /* This should have freed any remaining allocations. */
+                AMS_ABORT_UNLESS(g_dynamic_heap_outstanding_allocations == 0);
+
+                /* Free the heap. */
+                if (g_dynamic_heap_address != 0) {
+                    util::DestroyAt(g_dynamic_heap);
+                    g_dynamic_heap = {};
+
+                    R_ABORT_UNLESS(os::FreeUnsafeMemory(g_dynamic_heap_address, g_dynamic_heap_size));
+                }
+
+                /* Clear the heap globals. */
+                g_dynamic_heap_address = 0;
+                g_dynamic_heap_size    = 0;
+            }
+
+            /* Next, if we aren't going to end up building a romfs, we can ignore dynamic heap. */
+            R_SUCCEED_IF(!status.IsProgramSpecific());
+
+            /* Only mitm if there is actually an override romfs. */
+            R_SUCCEED_IF(!mitm::fs::HasSdRomfsContent(program_id));
+
+            /* Next, set the new program id for dynamic heap. */
+            g_dynamic_heap_program_id = program_id;
+            g_dynamic_heap_size       = GetDynamicHeapSize(g_dynamic_heap_program_id);
+
+            /* Set output. */
+            *out_size = g_dynamic_heap_size;
+
+            R_SUCCEED();
+        }
+
     }
 
 }

stratosphere/ams_mitm/source/fs_mitm/fsmitm_romfs.hpp

@@ -374,4 +374,6 @@ namespace ams::mitm::fs::romfs {
             void Build(SourceInfoVector *out_infos);
     };
 
+    Result ConfigureDynamicHeap(u64 *out_size, ncm::ProgramId program_id, const cfg::OverrideStatus &status, bool is_application);
+
 }

stratosphere/ams_mitm/source/mitm_pm/mitm_pm_service.cpp

@@ -16,13 +16,19 @@
 #include <stratosphere.hpp>
 #include "../amsmitm_initialization.hpp"
 #include "mitm_pm_service.hpp"
+#include "mitm_pm_service.hpp"
+#include "../fs_mitm/fsmitm_romfs.hpp"
 
 namespace ams::mitm::pm {
 
     Result PmService::PrepareLaunchProgram(sf::Out<u64> out, ncm::ProgramId program_id, const cfg::OverrideStatus &status, bool is_application) {
-        /* TODO */
+        /* Default to zero heap. */
         *out = 0;
-        AMS_UNUSED(program_id, status, is_application);
+
+        /* Actually configure the required boost size for romfs. */
+        R_TRY(mitm::fs::romfs::ConfigureDynamicHeap(out.GetPointer(), program_id, status, is_application));
+
+        /* TODO: Is there anything else we should do, while we have the opportunity? */
         R_SUCCEED();
     }
 

stratosphere/fatal/fatal.json

@@ -79,7 +79,7 @@
 				"svcReplyAndReceiveWithUserBuffer":	"0x44",
 				"svcCreateEvent":	"0x45",
 				"svcMapPhysicalMemoryUnsafe":	"0x48",
-				"svcUnmapPhysicalMemoryUnsafe":	"0x48",
+				"svcUnmapPhysicalMemoryUnsafe":	"0x49",
 				"svcSetUnsafeLimit":	"0x4A",
 				"svcReadWriteRegister":	"0x4E",
                 "svcDebugActiveProcess": "0x60",

stratosphere/fatal/source/fatal_task_screen.cpp

@@ -69,7 +69,7 @@ namespace ams::fatal::srv {
             /* Neither heap nor insecure is available, so we're going to have to try to raid the unsafe pool. */
             {
                 /* First, increase the limit to an extremely high value. */
-                size_t large_size = std::max(64_MB, FrameBufferRequiredSizeHeapAligned);
+                size_t large_size = std::max(128_MB, FrameBufferRequiredSizeHeapAligned);
                 while (svc::ResultLimitReached::Includes(svc::SetUnsafeLimit(large_size))) {
                     large_size *= 2;
                 }

stratosphere/pm/source/impl/pm_process_manager.cpp

@@ -257,8 +257,10 @@ namespace ams::pm::impl {
                 u64 mitm_boost_size = 0;
                 R_TRY(mitm::pm::PrepareLaunchProgram(std::addressof(mitm_boost_size), program_info.program_id, override_status, is_application));
 
-                R_ABORT_UNLESS(BoostSystemMemoryResourceLimitForMitm(mitm_boost_size));
-                ON_RESULT_FAILURE_2 { R_ABORT_UNLESS(BoostSystemMemoryResourceLimitForMitm(0)); };
+                if (mitm_boost_size > 0 || is_application) {
+                    R_ABORT_UNLESS(BoostSystemMemoryResourceLimitForMitm(mitm_boost_size));
+                    ON_RESULT_FAILURE_2 { R_ABORT_UNLESS(BoostSystemMemoryResourceLimitForMitm(0)); };
+                }
 
                 /* Ensure resources are available. */
                 resource::WaitResourceAvailable(std::addressof(program_info));