romfs: fix issues in close-during-build

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

libraries/libvapours/include/vapours/results/fs_results.hpp

@@ -41,8 +41,9 @@ namespace ams::fs {
         R_DEFINE_ERROR_RESULT(MountNameAlreadyExists, 60);
 
     R_DEFINE_ERROR_RANGE(HandledBySystemProcess, 1000, 2999);
-        R_DEFINE_ERROR_RESULT(PartitionNotFound,    1001);
+        R_DEFINE_ERROR_RESULT(PartitionNotFound,      1001);
-        R_DEFINE_ERROR_RESULT(TargetNotFound,       1002);
+        R_DEFINE_ERROR_RESULT(TargetNotFound,         1002);
+        R_DEFINE_ERROR_RESULT(NcaExternalKeyNotFound, 1004);
 
         R_DEFINE_ERROR_RANGE(SdCardAccessFailed, 2000, 2499);
             R_DEFINE_ERROR_RESULT(SdCardNotPresent, 2001);

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_initialization_mutex);
+        std::scoped_lock lk2(g_storage_set_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,82 @@ namespace ams::mitm::fs {
 
         namespace {
 
-            /* TODO: Fancy Dynamic allocation globals. */
+            struct ApplicationWithDynamicHeapInfo {
-            constinit os::SdkMutex g_romfs_build_lock;
+                ncm::ProgramId program_id;
-            //constinit size_t g_dynamic_heap_size = 0;
+                size_t dynamic_heap_size;
+            };
 
-            void InitializeDynamicHeapForBuildRomfs(ncm::ProgramId program_id) {
+            constexpr const ApplicationWithDynamicHeapInfo ApplicationsWithDynamicHeap[] = {
-                /* TODO */
+                { 0x01006F8002326000,  32_MB }, /* Animal Crossing: New Horizons. */
-                AMS_UNUSED(program_id);
+                { 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) {
+            /* Dynamic allocation globals. */
-                /* TODO */
+            constinit os::SdkMutex g_romfs_build_lock;
-                AMS_UNUSED(program_id);
+            constinit ncm::ProgramId g_dynamic_heap_program_id{};
+            constinit size_t g_dynamic_heap_size = 0;
+
+            constinit os::SdkMutex g_release_dynamic_heap_lock;
+            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) {
+                        /* NOTE: Lock not necessary, because this is the only location which do 0 -> non-zero. */
+
+                        /* 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 ReleaseDynamicHeap() {
+                std::scoped_lock lk(g_release_dynamic_heap_lock);
+
+                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));
+
+                    g_dynamic_heap_address = 0;
+                }
+            }
+
+            void FinalizeDynamicHeapForBuildRomfs() {
+                /* We are definitely no longer building out of dynamic heap. */
+                g_building_from_dynamic_heap = false;
+
+                if (g_dynamic_heap_outstanding_allocations == 0) {
+                    ReleaseDynamicHeap();
+                }
             }
 
         }
@@ -44,16 +109,32 @@ namespace ams::mitm::fs {
         void *AllocateTracked(AllocationType type, size_t size) {
             AMS_UNUSED(type);
 
-            /* TODO: Fancy dynamic allocation with memory stealing from application pool. */
+            if (g_building_from_dynamic_heap) {
-            return std::malloc(size);
+                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. */
+            if (IsAllocatedFromDynamicHeap(p)) {
-            return std::free(p);
+                --g_dynamic_heap_outstanding_allocations;
+                util::GetReference(g_dynamic_heap).Free(p);
+
+                if (!g_building_from_dynamic_heap && g_dynamic_heap_outstanding_allocations == 0) {
+                    ReleaseDynamicHeap();
+                }
+            } else {
+                std::free(p);
+            }
         }
 
         namespace {
@@ -147,21 +228,33 @@ namespace ams::mitm::fs {
                     size_t m_cache_idx;
                     u8 m_fallback_cache[FallbackCacheSize];
                 private:
-                    ALWAYS_INLINE void Read(size_t ofs, void *dst, size_t size) {
+                    ALWAYS_INLINE bool Read(size_t ofs, void *dst, size_t size) {
-                        R_ABORT_UNLESS(m_storage->Read(m_offset + ofs, dst, size));
+                        R_TRY_CATCH(m_storage->Read(m_offset + ofs, dst, size)) {
+                            R_CATCH(fs::ResultNcaExternalKeyNotFound) { return false; }
+                        } R_END_TRY_CATCH_WITH_ABORT_UNLESS;
+
+                        return true;
                     }
-                    ALWAYS_INLINE void ReloadCacheImpl(size_t idx) {
+                    ALWAYS_INLINE bool ReloadCacheImpl(size_t idx) {
                         const size_t rel_ofs = idx * this->GetCacheSize();
                         AMS_ABORT_UNLESS(rel_ofs < m_size);
                         const size_t new_cache_size = std::min(m_size - rel_ofs, this->GetCacheSize());
-                        this->Read(rel_ofs, m_cache, new_cache_size);
+                        if (!this->Read(rel_ofs, m_cache, new_cache_size)) {
+                            return false;
+                        }
+
                         m_cache_idx = idx;
+                        return true;
                     }
 
-                    ALWAYS_INLINE void ReloadCache(size_t idx) {
+                    ALWAYS_INLINE bool ReloadCache(size_t idx) {
                         if (m_cache_idx != idx) {
-                            this->ReloadCacheImpl(idx);
+                            if (!this->ReloadCacheImpl(idx)) {
+                                return false;
+                            }
                         }
+
+                        return true;
                     }
 
                     ALWAYS_INLINE size_t GetCacheIndex(u32 ofs) {
@@ -174,13 +267,18 @@ namespace ams::mitm::fs {
                     }
 
                     const Entry *GetEntry(u32 entry_offset) {
-                        this->ReloadCache(this->GetCacheIndex(entry_offset));
+                        if (!this->ReloadCache(this->GetCacheIndex(entry_offset))) {
+                            return nullptr;
+                        }
 
                         const size_t ofs = entry_offset % this->GetCacheSize();
 
                         const Entry *entry = reinterpret_cast<const Entry *>(reinterpret_cast<uintptr_t>(m_cache) + ofs);
                         if (AMS_UNLIKELY(this->GetCacheIndex(entry_offset) != this->GetCacheIndex(entry_offset + sizeof(Entry) + entry->name_size + sizeof(u32)))) {
-                            this->Read(entry_offset, m_fallback_cache, std::min(m_size - entry_offset, FallbackCacheSize));
+                            if (!this->Read(entry_offset, m_fallback_cache, std::min(m_size - entry_offset, FallbackCacheSize))) {
+                                return nullptr;
+                            }
+
                             entry = reinterpret_cast<const Entry *>(m_fallback_cache);
                         }
                         return entry;
@@ -342,7 +440,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();
@@ -447,10 +545,16 @@ namespace ams::mitm::fs {
 
         void Builder::VisitDirectory(BuildDirectoryContext *parent, u32 parent_offset, DirectoryTableReader &dir_table, FileTableReader &file_table) {
             const DirectoryEntry *parent_entry = dir_table.GetEntry(parent_offset);
+            if (AMS_UNLIKELY(parent_entry == nullptr)) {
+                return;
+            }
 
             u32 cur_file_offset = parent_entry->file;
             while (cur_file_offset != EmptyEntry) {
                 const FileEntry *cur_file = file_table.GetEntry(cur_file_offset);
+                if (AMS_UNLIKELY(cur_file == nullptr)) {
+                    return;
+                }
 
                 this->AddFile(parent, std::unique_ptr<BuildFileContext>(AllocateTyped<BuildFileContext>(AllocationType_BuildFileContext, cur_file->name, cur_file->name_size, cur_file->size, cur_file->offset, m_cur_source_type)));
 
@@ -463,6 +567,9 @@ namespace ams::mitm::fs {
                 u32 next_child_offset = 0;
                 {
                     const DirectoryEntry *cur_child = dir_table.GetEntry(cur_child_offset);
+                    if (AMS_UNLIKELY(cur_child == nullptr)) {
+                        return;
+                    }
 
                     this->AddDirectory(std::addressof(real_child), parent, std::unique_ptr<BuildDirectoryContext>(AllocateTyped<BuildDirectoryContext>(AllocationType_BuildDirContext, cur_child->name, cur_child->name_size)));
                     AMS_ABORT_UNLESS(real_child != nullptr);
@@ -801,6 +908,46 @@ 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) {
+                    ReleaseDynamicHeap();
+                }
+
+                /* 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));
+                if (mitm_boost_size > 0 || is_application) {
-                ON_RESULT_FAILURE_2 { R_ABORT_UNLESS(BoostSystemMemoryResourceLimitForMitm(0)); };
+                    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));