From e3f1de023f44b2cba8a2cd6d4a5faa4c665329dc Mon Sep 17 00:00:00 2001
From: "Admiral H. Curtiss" <pikachu025@gmail.com>
Date: Wed, 22 Sep 2021 04:51:16 +0200
Subject: DiscIO/DirectoryBlob: Add a callback that allows patching the
 main.dol and FST of the game partition during blob construction.

---
 Source/Core/DiscIO/DirectoryBlob.cpp | 43 ++++++++++++++++++++++--------------
 1 file changed, 27 insertions(+), 16 deletions(-)

(limited to 'Source/Core/DiscIO/DirectoryBlob.cpp')

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);
 }
 
-- 
cgit v1.2.3