1 files changed, 74 insertions, 79 deletions
diff --git a/src/Ryujinx.HLE/Loaders/Processes/Extensions/PartitionFileSystemExtensions.cs b/src/Ryujinx.HLE/Loaders/Processes/Extensions/PartitionFileSystemExtensions.cs
index 5f45cd45..50f7d585 100644
--- a/src/Ryujinx.HLE/Loaders/Processes/Extensions/PartitionFileSystemExtensions.cs
+++ b/src/Ryujinx.HLE/Loaders/Processes/Extensions/PartitionFileSystemExtensions.cs
@@ -1,87 +1,26 @@
 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());
 
-        public static Dictionary<ulong, ContentCollection> GetApplicationData(this IFileSystem partitionFileSystem,
-            VirtualFileSystem fileSystem, IntegrityCheckLevel checkLevel)
-        {
-            fileSystem.ImportTickets(partitionFileSystem);
-
-            var programs = new Dictionary<ulong, ContentCollection>();
-
-            foreach (DirectoryEntryEx fileEntry in partitionFileSystem.EnumerateEntries("/", "*.cnmt.nca"))
-            {
-                Cnmt cnmt = partitionFileSystem.GetNca(fileSystem.KeySet, fileEntry.FullPath).GetCnmt(checkLevel, ContentMetaType.Application);
-
-                if (cnmt == null)
-                {
-                    continue;
-                }
-
-                ContentCollection content = new(partitionFileSystem, cnmt);
-
-                if (content.Type != ContentMetaType.Application)
-                {
-                    continue;
-                }
-
-                programs.TryAdd(content.ApplicationId, content);
-            }
-
-            return programs;
-        }
-
-        public static Dictionary<ulong, ContentCollection> GetUpdateData(this IFileSystem partitionFileSystem,
-            VirtualFileSystem fileSystem, IntegrityCheckLevel checkLevel)
-        {
-            fileSystem.ImportTickets(partitionFileSystem);
-
-            var programs = new Dictionary<ulong, ContentCollection>();
-
-            foreach (DirectoryEntryEx fileEntry in partitionFileSystem.EnumerateEntries("/", "*.cnmt.nca"))
-            {
-                Cnmt cnmt = partitionFileSystem.GetNca(fileSystem.KeySet, fileEntry.FullPath).GetCnmt(checkLevel, ContentMetaType.Patch);
-
-                if (cnmt == null)
-                {
-                    continue;
-                }
-
-                ContentCollection content = new(partitionFileSystem, cnmt);
-
-                if (content.Type != ContentMetaType.Patch)
-                {
-                    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 titleId, out string errorMessage)
+        internal static (bool, ProcessResult) TryLoad<TMetaData, TFormat, THeader, TEntry>(this PartitionFileSystemCore<TMetaData, TFormat, THeader, TEntry> partitionFileSystem, Switch device, string path, out string errorMessage)
             where TMetaData : PartitionFileSystemMetaCore<TFormat, THeader, TEntry>, new()
             where TFormat : IPartitionFileSystemFormat
             where THeader : unmanaged, IPartitionFileSystemHeader
@@ -96,21 +35,30 @@ namespace Ryujinx.HLE.Loaders.Processes.Extensions
 
             try
             {
-                Dictionary<ulong, ContentCollection> applications = partitionFileSystem.GetApplicationData(device.FileSystem, device.System.FsIntegrityCheckLevel);
+                device.Configuration.VirtualFileSystem.ImportTickets(partitionFileSystem);
 
-                if (titleId == 0)
+                // TODO: To support multi-games container, this should use CNMT NCA instead.
+                foreach (DirectoryEntryEx fileEntry in partitionFileSystem.EnumerateEntries("/", "*.nca"))
                 {
-                    foreach ((ulong _, ContentCollection content) in applications)
+                    Nca nca = partitionFileSystem.GetNca(device, fileEntry.FullPath);
+
+                    if (nca.GetProgramIndex() != device.Configuration.UserChannelPersistence.Index)
                     {
-                        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;
+                        continue;
+                    }
+
+                    if (nca.IsPatch())
+                    {
+                        patchNca = nca;
+                    }
+                    else if (nca.IsProgram())
+                    {
+                        mainNca = nca;
+                    }
+                    else if (nca.IsControl())
+                    {
+                        controlNca = nca;
                     }
-                }
-                else if (applications.TryGetValue(titleId, out ContentCollection 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();
@@ -131,7 +79,54 @@ namespace Ryujinx.HLE.Loaders.Processes.Extensions
                     return (false, ProcessResult.Failed);
                 }
 
-                (Nca updatePatchNca, Nca updateControlNca) = mainNca.GetUpdateData(device.FileSystem, device.System.FsIntegrityCheckLevel, device.Configuration.UserChannelPersistence.Index, out string _);
+                // 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;
+                                }
+                            }
+                        }
+                    }
+                }
 
                 if (updatePatchNca != null)
                 {
@@ -173,18 +168,18 @@ namespace Ryujinx.HLE.Loaders.Processes.Extensions
                 return (true, mainNca.Load(device, patchNca, controlNca));
             }
 
-            errorMessage = $"Unable to load: Could not find Main NCA for title \"{titleId:X16}\"";
+            errorMessage = "Unable to load: Could not find Main NCA";
 
             return (false, ProcessResult.Failed);
         }
 
-        public static Nca GetNca(this IFileSystem fileSystem, KeySet keySet, string path)
+        public static Nca GetNca(this IFileSystem fileSystem, Switch device, string path)
         {
             using var ncaFile = new UniqueRef<IFile>();
 
             fileSystem.OpenFile(ref ncaFile.Ref, path.ToU8Span(), OpenMode.Read).ThrowIfFailure();
 
-            return new Nca(keySet, ncaFile.Release().AsStorage());
+            return new Nca(device.Configuration.VirtualFileSystem.KeySet, ncaFile.Release().AsStorage());
         }
     }
 }