DiscIO/DirectoryBlob: Add callback to modify sys files before they're added to the disc.

Source/Core/Core/Boot/Boot.cpp

@@ -700,6 +700,10 @@ void AddRiivolutionPatches(BootParameters* boot_params,
   auto& disc = std::get<BootParameters::Disc>(boot_params->parameters);
   disc.volume = DiscIO::CreateDisc(DiscIO::DirectoryBlobReader::Create(
       std::move(disc.volume),
+      [&](std::vector<DiscIO::FSTBuilderNode>* fst) {
+        DiscIO::Riivolution::ApplyPatchesToFiles(
+            riivolution_patches, DiscIO::Riivolution::PatchIndex::DolphinSysFiles, fst, nullptr);
+      },
       [&](std::vector<DiscIO::FSTBuilderNode>* fst, DiscIO::FSTBuilderNode* dol_node) {
         DiscIO::Riivolution::ApplyPatchesToFiles(
             riivolution_patches, DiscIO::Riivolution::PatchIndex::FileSystem, fst, dol_node);

Source/Core/DiscIO/DirectoryBlob.cpp

@@ -362,6 +362,7 @@ std::unique_ptr<DirectoryBlobReader> DirectoryBlobReader::Create(const std::stri
 
 std::unique_ptr<DirectoryBlobReader> DirectoryBlobReader::Create(
     std::unique_ptr<DiscIO::VolumeDisc> volume,
+    const std::function<void(std::vector<FSTBuilderNode>* fst_nodes)>& sys_callback,
     const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
         fst_callback)
 {
@@ -369,7 +370,7 @@ std::unique_ptr<DirectoryBlobReader> DirectoryBlobReader::Create(
     return nullptr;
 
   return std::unique_ptr<DirectoryBlobReader>(
-      new DirectoryBlobReader(std::move(volume), fst_callback));
+      new DirectoryBlobReader(std::move(volume), sys_callback, fst_callback));
 }
 
 DirectoryBlobReader::DirectoryBlobReader(const std::string& game_partition_root,
@@ -418,12 +419,14 @@ DirectoryBlobReader::DirectoryBlobReader(const std::string& game_partition_root,
 
 DirectoryBlobReader::DirectoryBlobReader(
     std::unique_ptr<DiscIO::VolumeDisc> volume,
+    const std::function<void(std::vector<FSTBuilderNode>* fst_nodes)>& sys_callback,
     const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
         fst_callback)
     : m_encryption_cache(this), m_wrapped_volume(std::move(volume))
 {
-  DirectoryBlobPartition game_partition(
+  DirectoryBlobPartition game_partition(m_wrapped_volume.get(),
-      m_wrapped_volume.get(), m_wrapped_volume->GetGamePartition(), std::nullopt, fst_callback);
+                                        m_wrapped_volume->GetGamePartition(), std::nullopt,
+                                        sys_callback, fst_callback);
   m_is_wii = game_partition.IsWii();
 
   if (!m_is_wii)
@@ -463,7 +466,7 @@ DirectoryBlobReader::DirectoryBlobReader(
       if (type)
       {
         partitions.emplace_back(
-            DirectoryBlobPartition(m_wrapped_volume.get(), partition, m_is_wii, nullptr),
+            DirectoryBlobPartition(m_wrapped_volume.get(), partition, m_is_wii, nullptr, nullptr),
             static_cast<PartitionType>(*type));
       }
     }
@@ -820,32 +823,80 @@ DirectoryBlobPartition::DirectoryBlobPartition(const std::string& root_directory
   BuildFSTFromFolder(m_root_directory + "files/", fst_address);
 }
 
+static void FillSingleFileNode(FSTBuilderNode* node, std::vector<u8> data)
+{
+  std::vector<BuilderContentSource> contents;
+  const size_t size = data.size();
+  contents.emplace_back(BuilderContentSource{0, size, ContentByteVector{std::move(data)}});
+  node->m_size = size;
+  node->m_content = std::move(contents);
+}
+
+static FSTBuilderNode BuildSingleFileNode(std::string filename, std::vector<u8> data,
+                                          void* userdata)
+{
+  FSTBuilderNode node{std::move(filename), 0, {}, userdata};
+  FillSingleFileNode(&node, std::move(data));
+  return node;
+}
+
+static std::vector<u8> ExtractNodeToVector(std::vector<FSTBuilderNode>* nodes, void* userdata)
+{
+  std::vector<u8> data;
+  const auto it =
+      std::find_if(nodes->begin(), nodes->end(), [&userdata](const FSTBuilderNode& node) {
+        return node.m_user_data == userdata;
+      });
+  if (it == nodes->end() || !it->IsFile())
+    return data;
+
+  DiscContentContainer tmp;
+  for (auto& content : it->GetFileContent())
+    tmp.Add(content.m_offset, content.m_size, std::move(content.m_source));
+  data.resize(it->m_size);
+  tmp.Read(0, it->m_size, data.data());
+  return data;
+}
+
 DirectoryBlobPartition::DirectoryBlobPartition(
     DiscIO::VolumeDisc* volume, const DiscIO::Partition& partition, std::optional<bool> is_wii,
+    const std::function<void(std::vector<FSTBuilderNode>* fst_nodes)>& sys_callback,
     const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
         fst_callback)
     : m_wrapped_partition(partition)
 {
+  std::vector<FSTBuilderNode> sys_nodes;
+
   std::vector<u8> disc_header(DISCHEADER_SIZE);
   if (!volume->Read(DISCHEADER_ADDRESS, DISCHEADER_SIZE, disc_header.data(), partition))
     disc_header.clear();
-  SetDiscHeader(std::move(disc_header));
+  sys_nodes.emplace_back(BuildSingleFileNode("boot.bin", std::move(disc_header), &m_disc_header));
-  SetDiscType(is_wii);
 
   std::vector<u8> bi2(BI2_SIZE);
   if (!volume->Read(BI2_ADDRESS, BI2_SIZE, bi2.data(), partition))
     bi2.clear();
-  SetBI2(std::move(bi2));
+  sys_nodes.emplace_back(BuildSingleFileNode("bi2.bin", std::move(bi2), &m_bi2));
 
   std::vector<u8> apploader;
   const auto apploader_size = GetApploaderSize(*volume, partition);
+  auto& apploader_node =
+      sys_nodes.emplace_back(FSTBuilderNode{"apploader.img", 0, {}, &m_apploader});
   if (apploader_size)
   {
     apploader.resize(*apploader_size);
     if (!volume->Read(APPLOADER_ADDRESS, *apploader_size, apploader.data(), partition))
       apploader.clear();
+    FillSingleFileNode(&apploader_node, std::move(apploader));
   }
-  const u64 new_dol_address = SetApploader(apploader, "apploader");
+
+  if (sys_callback)
+    sys_callback(&sys_nodes);
+
+  SetDiscHeader(ExtractNodeToVector(&sys_nodes, &m_disc_header));
+  SetDiscType(is_wii);
+  SetBI2(ExtractNodeToVector(&sys_nodes, &m_bi2));
+  const u64 new_dol_address =
+      SetApploader(ExtractNodeToVector(&sys_nodes, &m_apploader), "apploader");
 
   FSTBuilderNode dol_node{"main.dol", 0, {}};
   const auto dol_offset = GetBootDOLOffset(*volume, partition);

Source/Core/DiscIO/DirectoryBlob.h

@@ -190,10 +190,11 @@ class DirectoryBlobPartition
 public:
   DirectoryBlobPartition() = default;
   DirectoryBlobPartition(const std::string& root_directory, std::optional<bool> is_wii);
-  DirectoryBlobPartition(DiscIO::VolumeDisc* volume, const DiscIO::Partition& partition,
+  DirectoryBlobPartition(
-                         std::optional<bool> is_wii,
+      DiscIO::VolumeDisc* volume, const DiscIO::Partition& partition, std::optional<bool> is_wii,
-                         const std::function<void(std::vector<FSTBuilderNode>* fst_nodes,
+      const std::function<void(std::vector<FSTBuilderNode>* fst_nodes)>& sys_callback,
-                                                  FSTBuilderNode* dol_node)>& fst_callback);
+      const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
+          fst_callback);
 
   // We do not allow copying, because it might mess up the pointers inside DiscContents
   DirectoryBlobPartition(const DirectoryBlobPartition&) = delete;
@@ -265,6 +266,7 @@ public:
   static std::unique_ptr<DirectoryBlobReader> Create(const std::string& dol_path);
   static std::unique_ptr<DirectoryBlobReader> Create(
       std::unique_ptr<DiscIO::VolumeDisc> volume,
+      const std::function<void(std::vector<FSTBuilderNode>* fst_nodes)>& sys_callback,
       const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
           fst_callback);
 
@@ -303,9 +305,11 @@ private:
 
   explicit DirectoryBlobReader(const std::string& game_partition_root,
                                const std::string& true_root);
-  explicit DirectoryBlobReader(std::unique_ptr<DiscIO::VolumeDisc> volume,
+  explicit DirectoryBlobReader(
-                               const std::function<void(std::vector<FSTBuilderNode>* fst_nodes,
+      std::unique_ptr<DiscIO::VolumeDisc> volume,
-                                                        FSTBuilderNode* dol_node)>& fst_callback);
+      const std::function<void(std::vector<FSTBuilderNode>* fst_nodes)>& sys_callback,
+      const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
+          fst_callback);
 
   const DirectoryBlobPartition* GetPartition(u64 offset, u64 size, u64 partition_data_offset) const;