1 files changed, 55 insertions, 78 deletions
diff --git a/src/Ryujinx.HLE/Loaders/Processes/Extensions/PartitionFileSystemExtensions.cs b/src/Ryujinx.HLE/Loaders/Processes/Extensions/PartitionFileSystemExtensions.cs
index e95b1b05..bee2572a 100644
--- a/src/Ryujinx.HLE/Loaders/Processes/Extensions/PartitionFileSystemExtensions.cs
+++ b/src/Ryujinx.HLE/Loaders/Processes/Extensions/PartitionFileSystemExtensions.cs
@@ -1,26 +1,58 @@
 using LibHac.Common;
+using LibHac.Common.Keys;
 using LibHac.Fs;
 using LibHac.Fs.Fsa;
 using LibHac.FsSystem;
+using LibHac.Ncm;
 using LibHac.Tools.Fs;
 using LibHac.Tools.FsSystem;
 using LibHac.Tools.FsSystem.NcaUtils;
+using LibHac.Tools.Ncm;
 using Ryujinx.Common.Configuration;
 using Ryujinx.Common.Logging;
 using Ryujinx.Common.Utilities;
+using Ryujinx.HLE.FileSystem;
 using System;
 using System.Collections.Generic;
-using System.Globalization;
 using System.IO;
+using ContentType = LibHac.Ncm.ContentType;
 
 namespace Ryujinx.HLE.Loaders.Processes.Extensions
 {
     public static class PartitionFileSystemExtensions
     {
         private static readonly DownloadableContentJsonSerializerContext _contentSerializerContext = new(JsonHelper.GetDefaultSerializerOptions());
-        private static readonly TitleUpdateMetadataJsonSerializerContext _titleSerializerContext = new(JsonHelper.GetDefaultSerializerOptions());
 
-        internal static (bool, ProcessResult) TryLoad<TMetaData, TFormat, THeader, TEntry>(this PartitionFileSystemCore<TMetaData, TFormat, THeader, TEntry> partitionFileSystem, Switch device, string path, out string errorMessage)
+        public static Dictionary<ulong, ContentMetaData> GetContentData(this IFileSystem partitionFileSystem,
+            ContentMetaType contentType, VirtualFileSystem fileSystem, IntegrityCheckLevel checkLevel)
+        {
+            fileSystem.ImportTickets(partitionFileSystem);
+
+            var programs = new Dictionary<ulong, ContentMetaData>();
+
+            foreach (DirectoryEntryEx fileEntry in partitionFileSystem.EnumerateEntries("/", "*.cnmt.nca"))
+            {
+                Cnmt cnmt = partitionFileSystem.GetNca(fileSystem.KeySet, fileEntry.FullPath).GetCnmt(checkLevel, contentType);
+
+                if (cnmt == null)
+                {
+                    continue;
+                }
+
+                ContentMetaData content = new(partitionFileSystem, cnmt);
+
+                if (content.Type != contentType)
+                {
+                    continue;
+                }
+
+                programs.TryAdd(content.ApplicationId, content);
+            }
+
+            return programs;
+        }
+
+        internal static (bool, ProcessResult) TryLoad<TMetaData, TFormat, THeader, TEntry>(this PartitionFileSystemCore<TMetaData, TFormat, THeader, TEntry> partitionFileSystem, Switch device, string path, ulong applicationId, out string errorMessage)
             where TMetaData : PartitionFileSystemMetaCore<TFormat, THeader, TEntry>, new()
             where TFormat : IPartitionFileSystemFormat
             where THeader : unmanaged, IPartitionFileSystemHeader
@@ -35,31 +67,22 @@ namespace Ryujinx.HLE.Loaders.Processes.Extensions
 
             try
             {
-                device.Configuration.VirtualFileSystem.ImportTickets(partitionFileSystem);
+                Dictionary<ulong, ContentMetaData> applications = partitionFileSystem.GetContentData(ContentMetaType.Application, device.FileSystem, device.System.FsIntegrityCheckLevel);
 
-                // TODO: To support multi-games container, this should use CNMT NCA instead.
-                foreach (DirectoryEntryEx fileEntry in partitionFileSystem.EnumerateEntries("/", "*.nca"))
+                if (applicationId == 0)
                 {
-                    Nca nca = partitionFileSystem.GetNca(device, fileEntry.FullPath);
-
-                    if (nca.GetProgramIndex() != device.Configuration.UserChannelPersistence.Index)
-                    {
-                        continue;
-                    }
-
-                    if (nca.IsPatch())
+                    foreach ((ulong _, ContentMetaData content) in applications)
                     {
-                        patchNca = nca;
-                    }
-                    else if (nca.IsProgram())
-                    {
-                        mainNca = nca;
-                    }
-                    else if (nca.IsControl())
-                    {
-                        controlNca = nca;
+                        mainNca = content.GetNcaByType(device.FileSystem.KeySet, ContentType.Program, device.Configuration.UserChannelPersistence.Index);
+                        controlNca = content.GetNcaByType(device.FileSystem.KeySet, ContentType.Control, device.Configuration.UserChannelPersistence.Index);
+                        break;
                     }
                 }
+                else if (applications.TryGetValue(applicationId, out ContentMetaData content))
+                {
+                    mainNca = content.GetNcaByType(device.FileSystem.KeySet, ContentType.Program, device.Configuration.UserChannelPersistence.Index);
+                    controlNca = content.GetNcaByType(device.FileSystem.KeySet, ContentType.Control, device.Configuration.UserChannelPersistence.Index);
+                }
 
                 ProcessLoaderHelper.RegisterProgramMapInfo(device, partitionFileSystem).ThrowIfFailure();
             }
@@ -79,54 +102,7 @@ namespace Ryujinx.HLE.Loaders.Processes.Extensions
                     return (false, ProcessResult.Failed);
                 }
 
-                // Load Update NCAs.
-                Nca updatePatchNca = null;
-                Nca updateControlNca = null;
-
-                if (ulong.TryParse(mainNca.Header.TitleId.ToString("x16"), NumberStyles.HexNumber, CultureInfo.InvariantCulture, out ulong titleIdBase))
-                {
-                    // Clear the program index part.
-                    titleIdBase &= ~0xFUL;
-
-                    // Load update information if exists.
-                    string titleUpdateMetadataPath = System.IO.Path.Combine(AppDataManager.GamesDirPath, titleIdBase.ToString("x16"), "updates.json");
-                    if (File.Exists(titleUpdateMetadataPath))
-                    {
-                        string updatePath = JsonHelper.DeserializeFromFile(titleUpdateMetadataPath, _titleSerializerContext.TitleUpdateMetadata).Selected;
-                        if (File.Exists(updatePath))
-                        {
-                            PartitionFileSystem updatePartitionFileSystem = new();
-                            updatePartitionFileSystem.Initialize(new FileStream(updatePath, FileMode.Open, FileAccess.Read).AsStorage()).ThrowIfFailure();
-
-                            device.Configuration.VirtualFileSystem.ImportTickets(updatePartitionFileSystem);
-
-                            // TODO: This should use CNMT NCA instead.
-                            foreach (DirectoryEntryEx fileEntry in updatePartitionFileSystem.EnumerateEntries("/", "*.nca"))
-                            {
-                                Nca nca = updatePartitionFileSystem.GetNca(device, fileEntry.FullPath);
-
-                                if (nca.GetProgramIndex() != device.Configuration.UserChannelPersistence.Index)
-                                {
-                                    continue;
-                                }
-
-                                if ($"{nca.Header.TitleId.ToString("x16")[..^3]}000" != titleIdBase.ToString("x16"))
-                                {
-                                    break;
-                                }
-
-                                if (nca.IsProgram())
-                                {
-                                    updatePatchNca = nca;
-                                }
-                                else if (nca.IsControl())
-                                {
-                                    updateControlNca = nca;
-                                }
-                            }
-                        }
-                    }
-                }
+                (Nca updatePatchNca, Nca updateControlNca) = mainNca.GetUpdateData(device.FileSystem, device.System.FsIntegrityCheckLevel, device.Configuration.UserChannelPersistence.Index, out string _);
 
                 if (updatePatchNca != null)
                 {
@@ -138,10 +114,8 @@ namespace Ryujinx.HLE.Loaders.Processes.Extensions
                     controlNca = updateControlNca;
                 }
 
-                // Load contained DownloadableContents.
                 // TODO: If we want to support multi-processes in future, we shouldn't clear AddOnContent data here.
                 device.Configuration.ContentManager.ClearAocData();
-                device.Configuration.ContentManager.AddAocData(partitionFileSystem, path, mainNca.Header.TitleId, device.Configuration.FsIntegrityCheckLevel);
 
                 // Load DownloadableContents.
                 string addOnContentMetadataPath = System.IO.Path.Combine(AppDataManager.GamesDirPath, mainNca.Header.TitleId.ToString("x16"), "dlc.json");
@@ -153,9 +127,12 @@ namespace Ryujinx.HLE.Loaders.Processes.Extensions
                     {
                         foreach (DownloadableContentNca downloadableContentNca in downloadableContentContainer.DownloadableContentNcaList)
                         {
-                            if (File.Exists(downloadableContentContainer.ContainerPath) && downloadableContentNca.Enabled)
+                            if (File.Exists(downloadableContentContainer.ContainerPath))
                             {
-                                device.Configuration.ContentManager.AddAocItem(downloadableContentNca.TitleId, downloadableContentContainer.ContainerPath, downloadableContentNca.FullPath);
+                                if (downloadableContentNca.Enabled)
+                                {
+                                    device.Configuration.ContentManager.AddAocItem(downloadableContentNca.TitleId, downloadableContentContainer.ContainerPath, downloadableContentNca.FullPath);
+                                }
                             }
                             else
                             {
@@ -168,18 +145,18 @@ namespace Ryujinx.HLE.Loaders.Processes.Extensions
                 return (true, mainNca.Load(device, patchNca, controlNca));
             }
 
-            errorMessage = "Unable to load: Could not find Main NCA";
+            errorMessage = $"Unable to load: Could not find Main NCA for title \"{applicationId:X16}\"";
 
             return (false, ProcessResult.Failed);
         }
 
-        public static Nca GetNca(this IFileSystem fileSystem, Switch device, string path)
+        public static Nca GetNca(this IFileSystem fileSystem, KeySet keySet, string path)
         {
             using var ncaFile = new UniqueRef<IFile>();
 
             fileSystem.OpenFile(ref ncaFile.Ref, path.ToU8Span(), OpenMode.Read).ThrowIfFailure();
 
-            return new Nca(device.Configuration.VirtualFileSystem.KeySet, ncaFile.Release().AsStorage());
+            return new Nca(keySet, ncaFile.Release().AsStorage());
         }
     }
 }