DiscIO/DirectoryBlob: Add a callback that allows patching the main.dol and FST of the game partition during blob construction.

2 files changed, 38 insertions, 19 deletions

diff --git a/Source/Core/DiscIO/DirectoryBlob.cpp b/Source/Core/DiscIO/DirectoryBlob.cpp
index b41a00cf66..80bd173ae0 100644
--- a/Source/Core/DiscIO/DirectoryBlob.cpp
+++ b/Source/Core/DiscIO/DirectoryBlob.cpp
@@ -356,13 +356,16 @@ std::unique_ptr<DirectoryBlobReader> DirectoryBlobReader::Create(const std::stri
   return std::unique_ptr<DirectoryBlobReader>(new DirectoryBlobReader(partition_root, true_root));
 }
 
-std::unique_ptr<DirectoryBlobReader>
-DirectoryBlobReader::Create(std::unique_ptr<DiscIO::VolumeDisc> volume)
+std::unique_ptr<DirectoryBlobReader> DirectoryBlobReader::Create(
+    std::unique_ptr<DiscIO::VolumeDisc> volume,
+    const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
+        fst_callback)
 {
   if (!volume)
     return nullptr;
 
-  return std::unique_ptr<DirectoryBlobReader>(new DirectoryBlobReader(std::move(volume)));
+  return std::unique_ptr<DirectoryBlobReader>(
+      new DirectoryBlobReader(std::move(volume), fst_callback));
 }
 
 DirectoryBlobReader::DirectoryBlobReader(const std::string& game_partition_root,
@@ -409,11 +412,14 @@ DirectoryBlobReader::DirectoryBlobReader(const std::string& game_partition_root,
   }
 }
 
-DirectoryBlobReader::DirectoryBlobReader(std::unique_ptr<DiscIO::VolumeDisc> volume)
+DirectoryBlobReader::DirectoryBlobReader(
+    std::unique_ptr<DiscIO::VolumeDisc> volume,
+    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(m_wrapped_volume.get(),
-                                        m_wrapped_volume->GetGamePartition(), std::nullopt);
+  DirectoryBlobPartition game_partition(
+      m_wrapped_volume.get(), m_wrapped_volume->GetGamePartition(), std::nullopt, fst_callback);
   m_is_wii = game_partition.IsWii();
 
   if (!m_is_wii)
@@ -452,8 +458,9 @@ DirectoryBlobReader::DirectoryBlobReader(std::unique_ptr<DiscIO::VolumeDisc> vol
       auto type = m_wrapped_volume->GetPartitionType(partition);
       if (type)
       {
-        partitions.emplace_back(DirectoryBlobPartition(m_wrapped_volume.get(), partition, m_is_wii),
-                                static_cast<PartitionType>(*type));
+        partitions.emplace_back(
+            DirectoryBlobPartition(m_wrapped_volume.get(), partition, m_is_wii, nullptr),
+            static_cast<PartitionType>(*type));
       }
     }
 
@@ -809,9 +816,10 @@ DirectoryBlobPartition::DirectoryBlobPartition(const std::string& root_directory
   BuildFSTFromFolder(m_root_directory + "files/", fst_address);
 }
 
-DirectoryBlobPartition::DirectoryBlobPartition(DiscIO::VolumeDisc* volume,
-                                               const DiscIO::Partition& partition,
-                                               std::optional<bool> is_wii)
+DirectoryBlobPartition::DirectoryBlobPartition(
+    DiscIO::VolumeDisc* volume, const DiscIO::Partition& partition, std::optional<bool> is_wii,
+    const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
+        fst_callback)
     : m_wrapped_partition(partition)
 {
   std::vector<u8> disc_header(DISCHEADER_SIZE);
@@ -849,14 +857,17 @@ DirectoryBlobPartition::DirectoryBlobPartition(DiscIO::VolumeDisc* volume,
       dol_node.m_content = std::move(dol_contents);
     }
   }
-  const u64 new_fst_address = SetDOL(std::move(dol_node), new_dol_address);
+
+  std::vector<FSTBuilderNode> nodes;
 
   const FileSystem* fs = volume->GetFileSystem(partition);
-  if (!fs || !fs->IsValid())
-    return;
+  if (fs && fs->IsValid())
+    GenerateBuilderNodesFromFileSystem(*volume, partition, &nodes, fs->GetRoot());
 
-  std::vector<FSTBuilderNode> nodes;
-  GenerateBuilderNodesFromFileSystem(*volume, partition, &nodes, fs->GetRoot());
+  if (fst_callback)
+    fst_callback(&nodes, &dol_node);
+
+  const u64 new_fst_address = SetDOL(std::move(dol_node), new_dol_address);
   BuildFST(std::move(nodes), new_fst_address);
 }
 
diff --git a/Source/Core/DiscIO/DirectoryBlob.h b/Source/Core/DiscIO/DirectoryBlob.h
index db88367d8c..b8ec90b4f7 100644
--- a/Source/Core/DiscIO/DirectoryBlob.h
+++ b/Source/Core/DiscIO/DirectoryBlob.h
@@ -5,6 +5,7 @@
 
 #include <array>
 #include <cstddef>
+#include <functional>
 #include <map>
 #include <memory>
 #include <optional>
@@ -182,7 +183,9 @@ public:
   DirectoryBlobPartition() = default;
   DirectoryBlobPartition(const std::string& root_directory, std::optional<bool> is_wii);
   DirectoryBlobPartition(DiscIO::VolumeDisc* volume, const DiscIO::Partition& partition,
-                         std::optional<bool> is_wii);
+                         std::optional<bool> is_wii,
+                         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;
@@ -252,7 +255,10 @@ class DirectoryBlobReader : public BlobReader
 
 public:
   static std::unique_ptr<DirectoryBlobReader> Create(const std::string& dol_path);
-  static std::unique_ptr<DirectoryBlobReader> Create(std::unique_ptr<DiscIO::VolumeDisc> volume);
+  static std::unique_ptr<DirectoryBlobReader> Create(
+      std::unique_ptr<DiscIO::VolumeDisc> volume,
+      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
   DirectoryBlobReader(const DirectoryBlobReader&) = delete;
@@ -288,7 +294,9 @@ private:
 
   explicit DirectoryBlobReader(const std::string& game_partition_root,
                                const std::string& true_root);
-  explicit DirectoryBlobReader(std::unique_ptr<DiscIO::VolumeDisc> volume);
+  explicit DirectoryBlobReader(std::unique_ptr<DiscIO::VolumeDisc> volume,
+                               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;