4 files changed, 57 insertions, 46 deletions

diff --git a/src/core/file_sys/card_image.cpp b/src/core/file_sys/card_image.cpp
index db2f6a955..755d3303e 100644
--- a/src/core/file_sys/card_image.cpp
+++ b/src/core/file_sys/card_image.cpp
@@ -29,7 +29,7 @@ constexpr std::array partition_names{
     "logo",
 };
 
-XCI::XCI(VirtualFile file_, std::size_t program_index)
+XCI::XCI(VirtualFile file_, u64 program_id, size_t program_index)
     : file(std::move(file_)), program_nca_status{Loader::ResultStatus::ErrorXCIMissingProgramNCA},
       partitions(partition_names.size()),
       partitions_raw(partition_names.size()), keys{Core::Crypto::KeyManager::Instance()} {
@@ -63,12 +63,12 @@ XCI::XCI(VirtualFile file_, std::size_t program_index)
 
     secure_partition = std::make_shared<NSP>(
         main_hfs.GetFile(partition_names[static_cast<std::size_t>(XCIPartition::Secure)]),
-        program_index);
+        program_id, program_index);
 
     ncas = secure_partition->GetNCAsCollapsed();
     program =
         secure_partition->GetNCA(secure_partition->GetProgramTitleID(), ContentRecordType::Program);
-    program_nca_status = secure_partition->GetProgramStatus(secure_partition->GetProgramTitleID());
+    program_nca_status = secure_partition->GetProgramStatus();
     if (program_nca_status == Loader::ResultStatus::ErrorNSPMissingProgramNCA) {
         program_nca_status = Loader::ResultStatus::ErrorXCIMissingProgramNCA;
     }
@@ -174,6 +174,10 @@ u64 XCI::GetProgramTitleID() const {
     return secure_partition->GetProgramTitleID();
 }
 
+std::vector<u64> XCI::GetProgramTitleIDs() const {
+    return secure_partition->GetProgramTitleIDs();
+}
+
 u32 XCI::GetSystemUpdateVersion() {
     const auto update = GetPartition(XCIPartition::Update);
     if (update == nullptr) {
@@ -229,9 +233,11 @@ const std::vector<std::shared_ptr<NCA>>& XCI::GetNCAs() const {
 }
 
 std::shared_ptr<NCA> XCI::GetNCAByType(NCAContentType type) const {
-    const auto iter =
-        std::find_if(ncas.begin(), ncas.end(),
-                     [type](const std::shared_ptr<NCA>& nca) { return nca->GetType() == type; });
+    const auto program_id = secure_partition->GetProgramTitleID();
+    const auto iter = std::find_if(
+        ncas.begin(), ncas.end(), [this, type, program_id](const std::shared_ptr<NCA>& nca) {
+            return nca->GetType() == type && nca->GetTitleId() == program_id;
+        });
     return iter == ncas.end() ? nullptr : *iter;
 }
 
diff --git a/src/core/file_sys/card_image.h b/src/core/file_sys/card_image.h
index 4960e90fe..0fd9fa87c 100644
--- a/src/core/file_sys/card_image.h
+++ b/src/core/file_sys/card_image.h
@@ -78,7 +78,7 @@ enum class XCIPartition : u8 { Update, Normal, Secure, Logo };
 
 class XCI : public ReadOnlyVfsDirectory {
 public:
-    explicit XCI(VirtualFile file, std::size_t program_index = 0);
+    explicit XCI(VirtualFile file, u64 program_id = 0, size_t program_index = 0);
     ~XCI() override;
 
     Loader::ResultStatus GetStatus() const;
@@ -104,6 +104,7 @@ public:
     VirtualFile GetLogoPartitionRaw() const;
 
     u64 GetProgramTitleID() const;
+    std::vector<u64> GetProgramTitleIDs() const;
     u32 GetSystemUpdateVersion();
     u64 GetSystemUpdateTitleID() const;
 
diff --git a/src/core/file_sys/submission_package.cpp b/src/core/file_sys/submission_package.cpp
index d51d469e3..f192dffa5 100644
--- a/src/core/file_sys/submission_package.cpp
+++ b/src/core/file_sys/submission_package.cpp
@@ -20,8 +20,9 @@
 
 namespace FileSys {
 
-NSP::NSP(VirtualFile file_, std::size_t program_index_)
-    : file(std::move(file_)), program_index(program_index_), status{Loader::ResultStatus::Success},
+NSP::NSP(VirtualFile file_, u64 title_id_, std::size_t program_index_)
+    : file(std::move(file_)), expected_program_id(title_id_),
+      program_index(program_index_), status{Loader::ResultStatus::Success},
       pfs(std::make_shared<PartitionFilesystem>(file)), keys{Core::Crypto::KeyManager::Instance()} {
     if (pfs->GetStatus() != Loader::ResultStatus::Success) {
         status = pfs->GetStatus();
@@ -46,60 +47,59 @@ Loader::ResultStatus NSP::GetStatus() const {
     return status;
 }
 
-Loader::ResultStatus NSP::GetProgramStatus(u64 title_id) const {
+Loader::ResultStatus NSP::GetProgramStatus() const {
     if (IsExtractedType() && GetExeFS() != nullptr && FileSys::IsDirectoryExeFS(GetExeFS())) {
         return Loader::ResultStatus::Success;
     }
 
-    const auto iter = program_status.find(title_id);
+    const auto iter = program_status.find(GetProgramTitleID());
     if (iter == program_status.end())
         return Loader::ResultStatus::ErrorNSPMissingProgramNCA;
     return iter->second;
 }
 
-u64 NSP::GetFirstTitleID() const {
-    if (IsExtractedType()) {
-        return GetProgramTitleID();
-    }
-
-    if (program_status.empty())
-        return 0;
-    return program_status.begin()->first;
-}
-
 u64 NSP::GetProgramTitleID() const {
     if (IsExtractedType()) {
-        if (GetExeFS() == nullptr || !IsDirectoryExeFS(GetExeFS())) {
-            return 0;
-        }
+        return GetExtractedTitleID() + program_index;
+    }
 
-        ProgramMetadata meta;
-        if (meta.Load(GetExeFS()->GetFile("main.npdm")) == Loader::ResultStatus::Success) {
-            return meta.GetTitleID();
-        } else {
-            return 0;
+    auto program_id = expected_program_id;
+    if (program_id == 0) {
+        if (!program_status.empty()) {
+            program_id = program_status.begin()->first;
         }
     }
 
-    const auto out = GetFirstTitleID();
-    if ((out & 0x800) == 0)
-        return out;
+    program_id = program_id + program_index;
+    if (program_status.find(program_id) != program_status.end()) {
+        return program_id;
+    }
 
-    const auto ids = GetTitleIDs();
+    const auto ids = GetProgramTitleIDs();
     const auto iter =
         std::find_if(ids.begin(), ids.end(), [](u64 tid) { return (tid & 0x800) == 0; });
-    return iter == ids.end() ? out : *iter;
+    return iter == ids.end() ? 0 : *iter;
+}
+
+u64 NSP::GetExtractedTitleID() const {
+    if (GetExeFS() == nullptr || !IsDirectoryExeFS(GetExeFS())) {
+        return 0;
+    }
+
+    ProgramMetadata meta;
+    if (meta.Load(GetExeFS()->GetFile("main.npdm")) == Loader::ResultStatus::Success) {
+        return meta.GetTitleID();
+    } else {
+        return 0;
+    }
 }
 
-std::vector<u64> NSP::GetTitleIDs() const {
+std::vector<u64> NSP::GetProgramTitleIDs() const {
     if (IsExtractedType()) {
-        return {GetProgramTitleID()};
+        return {GetExtractedTitleID()};
     }
 
-    std::vector<u64> out;
-    out.reserve(ncas.size());
-    for (const auto& kv : ncas)
-        out.push_back(kv.first);
+    std::vector<u64> out{program_ids.cbegin(), program_ids.cend()};
     return out;
 }
 
@@ -146,7 +146,7 @@ std::shared_ptr<NCA> NSP::GetNCA(u64 title_id, ContentRecordType type, TitleType
     if (extracted)
         LOG_WARNING(Service_FS, "called on an NSP that is of type extracted.");
 
-    const auto title_id_iter = ncas.find(title_id + program_index);
+    const auto title_id_iter = ncas.find(title_id);
     if (title_id_iter == ncas.end())
         return nullptr;
 
@@ -160,7 +160,7 @@ std::shared_ptr<NCA> NSP::GetNCA(u64 title_id, ContentRecordType type, TitleType
 VirtualFile NSP::GetNCAFile(u64 title_id, ContentRecordType type, TitleType title_type) const {
     if (extracted)
         LOG_WARNING(Service_FS, "called on an NSP that is of type extracted.");
-    const auto nca = GetNCA(title_id, type);
+    const auto nca = GetNCA(title_id, type, title_type);
     if (nca != nullptr)
         return nca->GetBaseFile();
     return nullptr;
@@ -286,6 +286,7 @@ void NSP::ReadNCAs(const std::vector<VirtualFile>& files) {
 
                 if (next_nca->GetType() == NCAContentType::Program) {
                     program_status[next_nca->GetTitleId()] = next_nca->GetStatus();
+                    program_ids.insert(next_nca->GetTitleId() & 0xFFFFFFFFFFFFF000);
                 }
 
                 if (next_nca->GetStatus() != Loader::ResultStatus::Success &&
diff --git a/src/core/file_sys/submission_package.h b/src/core/file_sys/submission_package.h
index ecb3b6f15..030f36c09 100644
--- a/src/core/file_sys/submission_package.h
+++ b/src/core/file_sys/submission_package.h
@@ -6,6 +6,7 @@
 
 #include <map>
 #include <memory>
+#include <set>
 #include <vector>
 #include "common/common_types.h"
 #include "core/file_sys/vfs.h"
@@ -27,15 +28,15 @@ enum class ContentRecordType : u8;
 
 class NSP : public ReadOnlyVfsDirectory {
 public:
-    explicit NSP(VirtualFile file_, std::size_t program_index_ = 0);
+    explicit NSP(VirtualFile file_, u64 title_id = 0, std::size_t program_index_ = 0);
     ~NSP() override;
 
     Loader::ResultStatus GetStatus() const;
-    Loader::ResultStatus GetProgramStatus(u64 title_id) const;
+    Loader::ResultStatus GetProgramStatus() const;
     // Should only be used when one title id can be assured.
-    u64 GetFirstTitleID() const;
     u64 GetProgramTitleID() const;
-    std::vector<u64> GetTitleIDs() const;
+    u64 GetExtractedTitleID() const;
+    std::vector<u64> GetProgramTitleIDs() const;
 
     bool IsExtractedType() const;
 
@@ -69,6 +70,7 @@ private:
 
     VirtualFile file;
 
+    const u64 expected_program_id;
     const std::size_t program_index;
 
     bool extracted = false;
@@ -78,6 +80,7 @@ private:
     std::shared_ptr<PartitionFilesystem> pfs;
     // Map title id -> {map type -> NCA}
     std::map<u64, std::map<std::pair<TitleType, ContentRecordType>, std::shared_ptr<NCA>>> ncas;
+    std::set<u64> program_ids;
     std::vector<VirtualFile> ticket_files;
 
     Core::Crypto::KeyManager& keys;