6 files changed, 65 insertions, 18 deletions

diff --git a/src/core/file_sys/archive_selfncch.cpp b/src/core/file_sys/archive_selfncch.cpp
index 7dc91a405..a16941c70 100644
--- a/src/core/file_sys/archive_selfncch.cpp
+++ b/src/core/file_sys/archive_selfncch.cpp
@@ -3,12 +3,14 @@
 // Refer to the license.txt file included.
 
 #include <array>
+#include <cinttypes>
 #include "common/common_types.h"
 #include "common/logging/log.h"
 #include "common/swap.h"
 #include "core/file_sys/archive_selfncch.h"
 #include "core/file_sys/errors.h"
 #include "core/file_sys/ivfc_archive.h"
+#include "core/hle/kernel/process.h"
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // FileSys namespace
@@ -227,38 +229,57 @@ private:
     NCCHData ncch_data;
 };
 
-ArchiveFactory_SelfNCCH::ArchiveFactory_SelfNCCH(Loader::AppLoader& app_loader) {
-    std::shared_ptr<FileUtil::IOFile> romfs_file;
+void ArchiveFactory_SelfNCCH::Register(Loader::AppLoader& app_loader) {
+    u64 program_id = 0;
+    if (app_loader.ReadProgramId(program_id) != Loader::ResultStatus::Success) {
+        LOG_WARNING(
+            Service_FS,
+            "Could not read program id when registering with SelfNCCH, this might be a 3dsx file");
+    }
+
+    LOG_DEBUG(Service_FS, "Registering program %016" PRIX64 " with the SelfNCCH archive factory",
+              program_id);
+
+    if (ncch_data.find(program_id) != ncch_data.end()) {
+        LOG_WARNING(Service_FS, "Registering program %016" PRIX64
+                                " with SelfNCCH will override existing mapping",
+                    program_id);
+    }
+
+    NCCHData& data = ncch_data[program_id];
+
+    std::shared_ptr<FileUtil::IOFile> romfs_file_;
     if (Loader::ResultStatus::Success ==
-        app_loader.ReadRomFS(romfs_file, ncch_data.romfs_offset, ncch_data.romfs_size)) {
+        app_loader.ReadRomFS(romfs_file_, data.romfs_offset, data.romfs_size)) {
 
-        ncch_data.romfs_file = std::move(romfs_file);
+        data.romfs_file = std::move(romfs_file_);
     }
 
     std::shared_ptr<FileUtil::IOFile> update_romfs_file;
     if (Loader::ResultStatus::Success ==
-        app_loader.ReadUpdateRomFS(update_romfs_file, ncch_data.update_romfs_offset,
-                                   ncch_data.update_romfs_size)) {
+        app_loader.ReadUpdateRomFS(update_romfs_file, data.update_romfs_offset,
+                                   data.update_romfs_size)) {
 
-        ncch_data.update_romfs_file = std::move(update_romfs_file);
+        data.update_romfs_file = std::move(update_romfs_file);
     }
 
     std::vector<u8> buffer;
 
     if (Loader::ResultStatus::Success == app_loader.ReadIcon(buffer))
-        ncch_data.icon = std::make_shared<std::vector<u8>>(std::move(buffer));
+        data.icon = std::make_shared<std::vector<u8>>(std::move(buffer));
 
     buffer.clear();
     if (Loader::ResultStatus::Success == app_loader.ReadLogo(buffer))
-        ncch_data.logo = std::make_shared<std::vector<u8>>(std::move(buffer));
+        data.logo = std::make_shared<std::vector<u8>>(std::move(buffer));
 
     buffer.clear();
     if (Loader::ResultStatus::Success == app_loader.ReadBanner(buffer))
-        ncch_data.banner = std::make_shared<std::vector<u8>>(std::move(buffer));
+        data.banner = std::make_shared<std::vector<u8>>(std::move(buffer));
 }
 
 ResultVal<std::unique_ptr<ArchiveBackend>> ArchiveFactory_SelfNCCH::Open(const Path& path) {
-    auto archive = std::make_unique<SelfNCCHArchive>(ncch_data);
+    auto archive = std::make_unique<SelfNCCHArchive>(
+        ncch_data[Kernel::g_current_process->codeset->program_id]);
     return MakeResult<std::unique_ptr<ArchiveBackend>>(std::move(archive));
 }
 
diff --git a/src/core/file_sys/archive_selfncch.h b/src/core/file_sys/archive_selfncch.h
index f1c659948..0d6d6766e 100644
--- a/src/core/file_sys/archive_selfncch.h
+++ b/src/core/file_sys/archive_selfncch.h
@@ -6,6 +6,7 @@
 
 #include <memory>
 #include <string>
+#include <unordered_map>
 #include <vector>
 #include "common/common_types.h"
 #include "core/file_sys/archive_backend.h"
@@ -33,7 +34,10 @@ struct NCCHData {
 /// File system interface to the SelfNCCH archive
 class ArchiveFactory_SelfNCCH final : public ArchiveFactory {
 public:
-    explicit ArchiveFactory_SelfNCCH(Loader::AppLoader& app_loader);
+    ArchiveFactory_SelfNCCH() = default;
+
+    /// Registers a loaded application so that we can open its SelfNCCH archive when requested.
+    void Register(Loader::AppLoader& app_loader);
 
     std::string GetName() const override {
         return "SelfNCCH";
@@ -43,7 +47,8 @@ public:
     ResultVal<ArchiveFormatInfo> GetFormatInfo(const Path& path) const override;
 
 private:
-    NCCHData ncch_data;
+    /// Mapping of ProgramId -> NCCHData
+    std::unordered_map<u64, NCCHData> ncch_data;
 };
 
 } // namespace FileSys
diff --git a/src/core/hle/service/fs/archive.cpp b/src/core/hle/service/fs/archive.cpp
index 033fbc9aa..4ccb3cd32 100644
--- a/src/core/hle/service/fs/archive.cpp
+++ b/src/core/hle/service/fs/archive.cpp
@@ -20,6 +20,7 @@
 #include "core/file_sys/archive_savedata.h"
 #include "core/file_sys/archive_sdmc.h"
 #include "core/file_sys/archive_sdmcwriteonly.h"
+#include "core/file_sys/archive_selfncch.h"
 #include "core/file_sys/archive_systemsavedata.h"
 #include "core/file_sys/directory_backend.h"
 #include "core/file_sys/errors.h"
@@ -48,7 +49,7 @@ struct hash<Service::FS::ArchiveIdCode> {
         return std::hash<Type>()(static_cast<Type>(id_code));
     }
 };
-}
+} // namespace std
 
 static constexpr Kernel::Handle INVALID_HANDLE{};
 
@@ -564,6 +565,21 @@ void RegisterArchiveTypes() {
     auto systemsavedata_factory =
         std::make_unique<FileSys::ArchiveFactory_SystemSaveData>(nand_directory);
     RegisterArchiveType(std::move(systemsavedata_factory), ArchiveIdCode::SystemSaveData);
+
+    auto selfncch_factory = std::make_unique<FileSys::ArchiveFactory_SelfNCCH>();
+    RegisterArchiveType(std::move(selfncch_factory), ArchiveIdCode::SelfNCCH);
+}
+
+void RegisterSelfNCCH(Loader::AppLoader& app_loader) {
+    auto itr = id_code_map.find(ArchiveIdCode::SelfNCCH);
+    if (itr == id_code_map.end()) {
+        LOG_ERROR(Service_FS,
+                  "Could not register a new NCCH because the SelfNCCH archive hasn't been created");
+        return;
+    }
+
+    auto* factory = static_cast<FileSys::ArchiveFactory_SelfNCCH*>(itr->second.get());
+    factory->Register(app_loader);
 }
 
 void UnregisterArchiveTypes() {
diff --git a/src/core/hle/service/fs/archive.h b/src/core/hle/service/fs/archive.h
index 3a3371c88..e3c8fc2ef 100644
--- a/src/core/hle/service/fs/archive.h
+++ b/src/core/hle/service/fs/archive.h
@@ -21,6 +21,10 @@ static constexpr char SYSTEM_ID[]{"00000000000000000000000000000000"};
 /// The scrambled SD card CID, also known as ID1
 static constexpr char SDCARD_ID[]{"00000000000000000000000000000000"};
 
+namespace Loader {
+class AppLoader;
+}
+
 namespace Service {
 namespace FS {
 
@@ -259,6 +263,9 @@ void ArchiveInit();
 /// Shutdown archives
 void ArchiveShutdown();
 
+/// Registers a new NCCH file with the SelfNCCH archive factory
+void RegisterSelfNCCH(Loader::AppLoader& app_loader);
+
 /// Register all archive types
 void RegisterArchiveTypes();
 
diff --git a/src/core/loader/3dsx.cpp b/src/core/loader/3dsx.cpp
index a03515e6e..5ad5c5287 100644
--- a/src/core/loader/3dsx.cpp
+++ b/src/core/loader/3dsx.cpp
@@ -278,8 +278,7 @@ ResultStatus AppLoader_THREEDSX::Load() {
 
     Kernel::g_current_process->Run(48, Kernel::DEFAULT_STACK_SIZE);
 
-    Service::FS::RegisterArchiveType(std::make_unique<FileSys::ArchiveFactory_SelfNCCH>(*this),
-                                     Service::FS::ArchiveIdCode::SelfNCCH);
+    Service::FS::RegisterSelfNCCH(*this);
 
     is_loaded = true;
     return ResultStatus::Success;
diff --git a/src/core/loader/ncch.cpp b/src/core/loader/ncch.cpp
index c46d7cfc6..5107135f9 100644
--- a/src/core/loader/ncch.cpp
+++ b/src/core/loader/ncch.cpp
@@ -187,8 +187,7 @@ ResultStatus AppLoader_NCCH::Load() {
     if (ResultStatus::Success != result)
         return result;
 
-    Service::FS::RegisterArchiveType(std::make_unique<FileSys::ArchiveFactory_SelfNCCH>(*this),
-                                     Service::FS::ArchiveIdCode::SelfNCCH);
+    Service::FS::RegisterSelfNCCH(*this);
 
     ParseRegionLockoutInfo();