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

This commit is contained in:
Admiral H. Curtiss
2021-09-22 04:51:16 +02:00
parent 00ef9f2b4f
commit e3f1de023f
2 changed files with 39 additions and 20 deletions

View File

@ -356,13 +356,16 @@ std::unique_ptr<DirectoryBlobReader> DirectoryBlobReader::Create(const std::stri
return std::unique_ptr<DirectoryBlobReader>(new DirectoryBlobReader(partition_root, true_root)); return std::unique_ptr<DirectoryBlobReader>(new DirectoryBlobReader(partition_root, true_root));
} }
std::unique_ptr<DirectoryBlobReader> std::unique_ptr<DirectoryBlobReader> DirectoryBlobReader::Create(
DirectoryBlobReader::Create(std::unique_ptr<DiscIO::VolumeDisc> volume) std::unique_ptr<DiscIO::VolumeDisc> volume,
const std::function<void(std::vector<FSTBuilderNode>* fst_nodes, FSTBuilderNode* dol_node)>&
fst_callback)
{ {
if (!volume) if (!volume)
return nullptr; 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, 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)) : m_encryption_cache(this), m_wrapped_volume(std::move(volume))
{ {
DirectoryBlobPartition game_partition(m_wrapped_volume.get(), DirectoryBlobPartition game_partition(
m_wrapped_volume->GetGamePartition(), std::nullopt); m_wrapped_volume.get(), m_wrapped_volume->GetGamePartition(), std::nullopt, fst_callback);
m_is_wii = game_partition.IsWii(); m_is_wii = game_partition.IsWii();
if (!m_is_wii) if (!m_is_wii)
@ -452,8 +458,9 @@ DirectoryBlobReader::DirectoryBlobReader(std::unique_ptr<DiscIO::VolumeDisc> vol
auto type = m_wrapped_volume->GetPartitionType(partition); auto type = m_wrapped_volume->GetPartitionType(partition);
if (type) if (type)
{ {
partitions.emplace_back(DirectoryBlobPartition(m_wrapped_volume.get(), partition, m_is_wii), partitions.emplace_back(
static_cast<PartitionType>(*type)); 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); BuildFSTFromFolder(m_root_directory + "files/", fst_address);
} }
DirectoryBlobPartition::DirectoryBlobPartition(DiscIO::VolumeDisc* volume, DirectoryBlobPartition::DirectoryBlobPartition(
const DiscIO::Partition& partition, 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)
: m_wrapped_partition(partition) : m_wrapped_partition(partition)
{ {
std::vector<u8> disc_header(DISCHEADER_SIZE); std::vector<u8> disc_header(DISCHEADER_SIZE);
@ -849,14 +857,17 @@ DirectoryBlobPartition::DirectoryBlobPartition(DiscIO::VolumeDisc* volume,
dol_node.m_content = std::move(dol_contents); dol_node.m_content = std::move(dol_contents);
} }
} }
const u64 new_fst_address = SetDOL(std::move(dol_node), new_dol_address);
const FileSystem* fs = volume->GetFileSystem(partition);
if (!fs || !fs->IsValid())
return;
std::vector<FSTBuilderNode> nodes; std::vector<FSTBuilderNode> nodes;
GenerateBuilderNodesFromFileSystem(*volume, partition, &nodes, fs->GetRoot());
const FileSystem* fs = volume->GetFileSystem(partition);
if (fs && fs->IsValid())
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); BuildFST(std::move(nodes), new_fst_address);
} }

View File

@ -5,6 +5,7 @@
#include <array> #include <array>
#include <cstddef> #include <cstddef>
#include <functional>
#include <map> #include <map>
#include <memory> #include <memory>
#include <optional> #include <optional>
@ -182,7 +183,9 @@ public:
DirectoryBlobPartition() = default; DirectoryBlobPartition() = default;
DirectoryBlobPartition(const std::string& root_directory, std::optional<bool> is_wii); DirectoryBlobPartition(const std::string& root_directory, std::optional<bool> is_wii);
DirectoryBlobPartition(DiscIO::VolumeDisc* volume, const DiscIO::Partition& partition, 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 // We do not allow copying, because it might mess up the pointers inside DiscContents
DirectoryBlobPartition(const DirectoryBlobPartition&) = delete; DirectoryBlobPartition(const DirectoryBlobPartition&) = delete;
@ -252,7 +255,10 @@ class DirectoryBlobReader : public BlobReader
public: public:
static std::unique_ptr<DirectoryBlobReader> Create(const std::string& dol_path); 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 // We do not allow copying, because it might mess up the pointers inside DiscContents
DirectoryBlobReader(const DirectoryBlobReader&) = delete; DirectoryBlobReader(const DirectoryBlobReader&) = delete;
@ -288,7 +294,9 @@ private:
explicit DirectoryBlobReader(const std::string& game_partition_root, explicit DirectoryBlobReader(const std::string& game_partition_root,
const std::string& true_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; const DirectoryBlobPartition* GetPartition(u64 offset, u64 size, u64 partition_data_offset) const;