nod/lib/DiscBase.cpp
Minty-Meeo 3c25647b6e The Encoding Update
While Nintendo's own documents claim GameCube and Wii disc file symbol tables only support 7-bit ASCII, this is far from the truth.  Indeed, even some first-party Nintendo games shipped with Shift-JIS encoded file symbol tables.  My guess?  The locale of whatever Windows machine mastered a GameCube or Wii disc influenced how wide character strings (UCS-2) were converted to narrow character strings.  To account for all possibilites, this update adds extensible multi-byte character set options to NOD-Tool.

A rundown of notable changes:
 - "-c XXXXX" option added to set the encoding of the GameCube / Wii ISO(s) being processed.
 - "SystemStringConv" renamed to "DiscLocToSystemConv"
 - "SystemUTF8Conv" renamed to "SystemToDiscLocConv"
 - Help message updated with new info.
 - Bugfix: AddBuildName had a logic error wherein the length of the SystemString was being used instead of length of the disc locale string.  This would corrupt the File Symbol Table if the disc locale string's length was greater than the SystemString's length.
 - Bugfix: recursiveMergeFST was not keeping track of parent indexes at all, meaning nested folders and their contents would be corrupted.  I simply copied the way recursiveBuildFST did things to fix this.
 - Bugfix (Windows): On Windows, for some reason, Sstat was a typedef for _stat (32-bit) instead of _stat64 (64-bit).  This is confounding, because untrimmed Wii ISOs will always be larger than the unsigned 32-bit integer limit (4,699,979,776 bytes vs 4,294,967,295 bytes), meaning the MergeWii errand has never worked for untrimmed ISOs on Windows.  Was this never tested??
 - Bugfix (Windows): Did you know Windows Command Prompt fully supports Unicode?  Stdio streams are now in _O_U16TEXT mode for Windows only.  Previously, attempting to print any character that could not be narrowed to your locale's encoding would either silently fail (std functions), or throw an exception (fmt functions).  As a minor drawback, narrow character print functions can no longer be used when stdio is in _O_U16TEXT mode, necessitating my PR for Logvisor here: (AxioDL/logvisor#7)
 - ExtractionContext::progressCB now uses SystemStringView because widechar printing works correctly on Windows now.
 - progFunc lambda no longer throws exceptions when printing unicode because widechar printing works correctly on Windows now.
 - Top-level constructors and functions with a Codepage_t parameter have also signatures that default to the US-ASCII codepage.
    - DiscGCN constructor
    - DiscBuilderGCN constructor
    - DiscBuilderGCN::CalculateTotalSizeRequired
    - DiscMergerGCN constructor
    - DiscMergerGCN::CalculateTotalSizeRequired
    - DiscWii constructor
    - DiscBuilderWii constructor
    - DiscBuilderWii::CalculateTotalSizeRequired
    - DiscMergerWii constructor
    - DiscMergerWii::CalculateTotalSizeRequired
    - OpenDiscFromImage
 - Conversion between system encoding and disc locale encoding has checks in place to warn the user if string conversion goes awry.
2021-06-27 03:26:20 -05:00

859 lines
30 KiB
C++

#include "nod/DiscBase.hpp"
#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <unordered_map>
#include "nod/DirectoryEnumerator.hpp"
#include "nod/IFileIO.hpp"
#include "nod/nod.hpp"
#include "nod/Util.hpp"
#ifndef _WIN32
#include <unistd.h>
#else
#include <logvisor/logvisor.hpp>
static void* memmem(const void* haystack, size_t hlen, const void* needle, size_t nlen) {
int needle_first;
const uint8_t* p = static_cast<const uint8_t*>(haystack);
size_t plen = hlen;
if (!nlen)
return NULL;
needle_first = *(unsigned char*)needle;
while (plen >= nlen && (p = static_cast<const uint8_t*>(memchr(p, needle_first, plen - nlen + 1)))) {
if (!memcmp(p, needle, nlen))
return (void*)p;
p++;
plen = hlen - (p - static_cast<const uint8_t*>(haystack));
}
return NULL;
}
#endif
#include <algorithm>
namespace nod {
const SystemChar* getKindString(PartitionKind kind) {
switch (kind) {
case PartitionKind::Data:
return _SYS_STR("DATA");
case PartitionKind::Update:
return _SYS_STR("UPDATE");
case PartitionKind::Channel:
return _SYS_STR("CHANNEL");
default:
return nullptr;
}
}
void IPartition::parseFST(IPartReadStream& s) {
std::unique_ptr<uint8_t[]> fst(new uint8_t[m_fstSz]);
s.seek(m_fstOff);
s.read(fst.get(), m_fstSz);
const FSTNode* nodes = (FSTNode*)fst.get();
/* Root node indicates the count of all contained nodes */
uint32_t nodeCount = nodes[0].getLength();
const char* names = (char*)fst.get() + 12 * nodeCount;
m_nodes.clear();
m_nodes.reserve(nodeCount);
/* Construct nodes */
for (uint32_t n = 0; n < nodeCount; ++n) {
const FSTNode& node = nodes[n];
m_nodes.emplace_back(*this, node, n ? names + node.getNameOffset() : "");
}
/* Setup dir-child iterators */
for (std::vector<Node>::iterator it = m_nodes.begin(); it != m_nodes.end(); ++it) {
Node& node = *it;
if (node.m_kind == Node::Kind::Directory) {
node.m_childrenBegin = it + 1;
node.m_childrenEnd = m_nodes.begin() + node.m_discLength;
}
}
}
void IPartition::parseDOL(IPartReadStream& s) {
/* Read Dol header */
DOLHeader dolHeader;
s.read(&dolHeader, sizeof(DOLHeader));
/* Calculate Dol size */
uint32_t dolSize = SBig(dolHeader.textOff[0]);
for (uint32_t i = 0; i < 7; i++)
dolSize += SBig(dolHeader.textSizes[i]);
for (uint32_t i = 0; i < 11; i++)
dolSize += SBig(dolHeader.dataSizes[i]);
m_dolSz = dolSize;
}
Node::Node(const IPartition& parent, const FSTNode& node, std::string_view name)
: m_parent(parent)
, m_kind(node.isDir() ? Kind::Directory : Kind::File)
, m_discOffset(parent.normalizeOffset(node.getOffset()))
, m_discLength(node.getLength())
, m_name(name) {}
std::unique_ptr<IPartReadStream> Node::beginReadStream(uint64_t offset) const {
if (m_kind != Kind::File) {
LogModule.report(logvisor::Error, FMT_STRING("unable to stream a non-file {}"), m_name);
return std::unique_ptr<IPartReadStream>();
}
return m_parent.beginReadStream(m_discOffset + offset);
}
std::unique_ptr<uint8_t[]> Node::getBuf() const {
if (m_kind != Kind::File) {
LogModule.report(logvisor::Error, FMT_STRING("unable to buffer a non-file {}"), m_name);
return nullptr;
}
auto buf = std::unique_ptr<uint8_t[]>(new uint8_t[m_discLength]);
beginReadStream()->read(buf.get(), m_discLength);
return buf;
}
bool Node::extractToDirectory(SystemStringView basePath, const ExtractionContext& ctx, Codepage_t codepage) const {
DiscLocToSystemConv nameView(getName(), codepage);
SystemString path = SystemString(basePath) + _SYS_STR('/') + nameView.sys_str().data();
if (m_kind == Kind::Directory) {
++m_parent.m_curNodeIdx;
if (ctx.progressCB && !getName().empty())
ctx.progressCB(nameView.sys_str(), m_parent.m_curNodeIdx / float(m_parent.getNodeCount()));
if (Mkdir(path.c_str(), 0755) && errno != EEXIST) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("unable to mkdir '{}'")), path);
return false;
}
for (Node& subnode : *this)
if (!subnode.extractToDirectory(path, ctx, codepage))
return false;
} else if (m_kind == Kind::File) {
Sstat theStat;
if (ctx.progressCB)
ctx.progressCB(nameView.sys_str(), m_parent.m_curNodeIdx / float(m_parent.getNodeCount()));
if (ctx.force || Stat(path.c_str(), &theStat)) {
std::unique_ptr<IPartReadStream> rs = beginReadStream();
std::unique_ptr<IFileIO::IWriteStream> ws = NewFileIO(path)->beginWriteStream();
if (!rs || !ws)
return false;
ws->copyFromDisc(*rs, m_discLength, [&](float prog) {
if (ctx.progressCB)
ctx.progressCB(nameView.sys_str(), (m_parent.m_curNodeIdx + prog) / float(m_parent.getNodeCount()));
});
}
++m_parent.m_curNodeIdx;
}
return true;
}
bool IPartition::extractToDirectory(SystemStringView path, const ExtractionContext& ctx) {
m_curNodeIdx = 0;
if (Mkdir(path.data(), 0755) && errno != EEXIST) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("unable to mkdir '{}'")), path);
return false;
}
SystemString basePath = m_isWii ? SystemString(path) + _SYS_STR("/") + getKindString(m_kind) : SystemString(path);
if (m_isWii) {
if (Mkdir(basePath.c_str(), 0755) && errno != EEXIST) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("unable to mkdir '{}'")), basePath);
return false;
}
}
/* Extract Disc Files */
if (!m_parent.extractDiscHeaderFiles(basePath, ctx))
return false;
/* Extract Crypto Files */
if (!extractCryptoFiles(basePath, ctx))
return false;
if (!extractSysFiles(basePath, ctx))
return false;
/* Extract Filesystem */
SystemString fsPath = basePath + _SYS_STR("/files");
if (Mkdir(fsPath.c_str(), 0755) && errno != EEXIST) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("unable to mkdir '{}'")), fsPath);
return false;
}
return m_nodes[0].extractToDirectory(fsPath, ctx, m_codepage);
}
bool IPartition::extractSysFiles(SystemStringView basePath, const ExtractionContext& ctx) const {
SystemString basePathStr(basePath);
if (Mkdir((basePathStr + _SYS_STR("/sys")).c_str(), 0755) && errno != EEXIST) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("unable to mkdir '{}/sys'")), basePath);
return false;
}
Sstat theStat;
/* Extract Apploader */
SystemString apploaderPath = basePathStr + _SYS_STR("/sys/apploader.img");
if (ctx.force || Stat(apploaderPath.c_str(), &theStat)) {
if (ctx.progressCB)
ctx.progressCB(_SYS_STR("apploader.bin"), 0.f);
std::unique_ptr<uint8_t[]> buf = getApploaderBuf();
auto ws = NewFileIO(apploaderPath)->beginWriteStream();
if (!ws)
return false;
ws->write(buf.get(), m_apploaderSz);
}
/* Extract Dol */
SystemString dolPath = basePathStr + _SYS_STR("/sys/main.dol");
if (ctx.force || Stat(dolPath.c_str(), &theStat)) {
if (ctx.progressCB)
ctx.progressCB(_SYS_STR("main.dol"), 0.f);
std::unique_ptr<uint8_t[]> buf = getDOLBuf();
auto ws = NewFileIO(dolPath)->beginWriteStream();
if (!ws)
return false;
ws->write(buf.get(), m_dolSz);
}
/* Extract Boot info */
SystemString bootPath = basePathStr + _SYS_STR("/sys/boot.bin");
if (ctx.force || Stat(bootPath.c_str(), &theStat)) {
if (ctx.progressCB)
ctx.progressCB(_SYS_STR("boot.bin"), 0.f);
auto ws = NewFileIO(bootPath)->beginWriteStream();
if (!ws)
return false;
m_header.write(*ws.get());
}
/* Extract BI2 info */
SystemString bi2Path = basePathStr + _SYS_STR("/sys/bi2.bin");
if (ctx.force || Stat(bi2Path.c_str(), &theStat)) {
if (ctx.progressCB)
ctx.progressCB(_SYS_STR("bi2.bin"), 0.f);
auto ws = NewFileIO(bi2Path)->beginWriteStream();
if (!ws)
return false;
m_bi2Header.write(*ws);
}
return true;
}
static bool IsSystemFile(SystemStringView name, bool& isDol) {
isDol = false;
if (name.size() < 4)
return false;
if (!StrCaseCmp((&*(name.cend() - 4)), _SYS_STR(".dol"))) {
isDol = true;
return true;
}
if (!StrCaseCmp((&*(name.cend() - 4)), _SYS_STR(".rel")))
return true;
if (!StrCaseCmp((&*(name.cend() - 4)), _SYS_STR(".rso")))
return true;
if (!StrCaseCmp((&*(name.cend() - 4)), _SYS_STR(".sel")))
return true;
if (!StrCaseCmp((&*(name.cend() - 4)), _SYS_STR(".bnr")))
return true;
if (!StrCaseCmp((&*(name.cend() - 4)), _SYS_STR(".elf")))
return true;
if (!StrCaseCmp((&*(name.cend() - 4)), _SYS_STR(".wad")))
return true;
return false;
}
/** Patches out pesky #001 integrity check performed by game's OSInit.
* This is required for multi-DOL games, but doesn't harm functionality otherwise */
static void PatchDOL(std::unique_ptr<uint8_t[]>& buf, size_t sz, bool& patched) {
patched = false;
uint8_t* found = static_cast<uint8_t*>(memmem(buf.get(), sz,
"\x3C\x03\xF8\x00\x28\x00\x00\x00\x40\x82\x00\x0C"
"\x38\x60\x00\x01\x48\x00\x02\x44\x38\x61\x00\x18\x48",
25));
if (found) {
found[11] = '\x04';
patched = true;
}
}
static size_t PatchDOL(IFileIO::IReadStream& in, IPartWriteStream& out, size_t sz, bool& patched) {
std::unique_ptr<uint8_t[]> buf(new uint8_t[sz]);
sz = in.read(buf.get(), sz);
PatchDOL(buf, sz, patched);
return out.write(buf.get(), sz);
}
void DiscBuilderBase::PartitionBuilderBase::recursiveBuildNodesPre(SystemStringView filesIn) {
DirectoryEnumerator dEnum(filesIn, DirectoryEnumerator::Mode::DirsThenFilesSorted, false, false, true);
for (const DirectoryEnumerator::Entry& e : dEnum) {
if (e.m_isDir)
recursiveBuildNodesPre(e.m_path.c_str());
else
++m_parent.m_progressTotal;
}
}
bool DiscBuilderBase::PartitionBuilderBase::recursiveBuildNodes(IPartWriteStream& ws, bool system,
SystemStringView filesIn) {
DirectoryEnumerator dEnum(filesIn, DirectoryEnumerator::Mode::DirsThenFilesSorted, false, false, true);
for (const DirectoryEnumerator::Entry& e : dEnum) {
if (e.m_isDir) {
if (!recursiveBuildNodes(ws, system, e.m_path.c_str()))
return false;
} else {
bool isDol;
bool isSys = IsSystemFile(e.m_name, isDol);
if (system ^ isSys)
continue;
size_t fileSz = ROUND_UP_32(e.m_fileSz);
uint64_t fileOff = userAllocate(fileSz, ws);
if (fileOff == UINT64_MAX)
return false;
m_fileOffsetsSizes[e.m_path] = std::make_pair(fileOff, fileSz);
std::unique_ptr<IFileIO::IReadStream> rs = NewFileIO(e.m_path)->beginReadStream();
if (!rs)
return false;
size_t xferSz = 0;
if (isDol) {
bool patched;
xferSz = PatchDOL(*rs, ws, e.m_fileSz, patched);
m_parent.m_progressCB(m_parent.getProgressFactor(),
e.m_name + (patched ? _SYS_STR(" [PATCHED]") : _SYS_STR("")), xferSz);
++m_parent.m_progressIdx;
} else {
char buf[0x8000];
while (xferSz < e.m_fileSz) {
size_t rdSz = rs->read(buf, nod::min(size_t(0x8000ul), e.m_fileSz - xferSz));
if (!rdSz)
break;
ws.write(buf, rdSz);
xferSz += rdSz;
m_parent.m_progressCB(m_parent.getProgressFactorMidFile(xferSz, e.m_fileSz), e.m_name, xferSz);
}
++m_parent.m_progressIdx;
}
for (size_t i = 0; i < fileSz - xferSz; ++i)
ws.write("\xff", 1);
}
}
return true;
}
bool DiscBuilderBase::PartitionBuilderBase::recursiveBuildFST(SystemStringView filesIn,
std::function<void(void)> incParents,
size_t parentDirIdx) {
DirectoryEnumerator dEnum(filesIn, DirectoryEnumerator::Mode::DirsThenFilesSorted, false, false, true);
for (const DirectoryEnumerator::Entry& e : dEnum) {
if (e.m_isDir) {
size_t dirNodeIdx = m_buildNodes.size();
m_buildNodes.emplace_back(true, m_buildNameOff, parentDirIdx, dirNodeIdx + 1);
addBuildName(e.m_name);
incParents();
if (!recursiveBuildFST(e.m_path.c_str(),
[&]() {
m_buildNodes[dirNodeIdx].incrementLength();
incParents();
},
dirNodeIdx))
return false;
} else {
std::pair<uint64_t, uint64_t> fileOffSz = m_fileOffsetsSizes.at(e.m_path);
m_buildNodes.emplace_back(false, m_buildNameOff, packOffset(fileOffSz.first), fileOffSz.second);
addBuildName(e.m_name);
incParents();
}
}
return true;
}
void DiscBuilderBase::PartitionBuilderBase::recursiveMergeNodesPre(const Node* nodeIn, SystemStringView dirIn) {
/* Build map of existing nodes to write-through later */
std::unordered_map<std::string, const Node*> fileNodes;
std::unordered_map<std::string, const Node*> dirNodes;
if (nodeIn) {
fileNodes.reserve(nodeIn->size());
dirNodes.reserve(nodeIn->size());
for (const Node& ch : *nodeIn) {
if (ch.getKind() == Node::Kind::File)
fileNodes.insert(std::make_pair(ch.getName(), &ch));
else if (ch.getKind() == Node::Kind::Directory)
dirNodes.insert(std::make_pair(ch.getName(), &ch));
}
}
/* Merge this directory's files */
if (!dirIn.empty()) {
DirectoryEnumerator dEnum(dirIn, DirectoryEnumerator::Mode::DirsThenFilesSorted, false, false, true);
for (const DirectoryEnumerator::Entry& e : dEnum) {
SystemToDiscLocConv nameView(e.m_name, m_codepage);
if (e.m_isDir) {
auto search = dirNodes.find(nameView.disc_str().data());
if (search != dirNodes.cend()) {
recursiveMergeNodesPre(search->second, e.m_path.c_str());
dirNodes.erase(search);
} else {
recursiveMergeNodesPre(nullptr, e.m_path.c_str());
}
} else {
fileNodes.erase(nameView.disc_str().data());
++m_parent.m_progressTotal;
}
}
}
/* Write-through remaining dir nodes */
for (const auto& p : dirNodes) {
recursiveMergeNodesPre(p.second, {});
}
/* Write-through remaining file nodes */
m_parent.m_progressTotal += fileNodes.size();
}
bool DiscBuilderBase::PartitionBuilderBase::recursiveMergeNodes(IPartWriteStream& ws, bool system, const Node* nodeIn,
SystemStringView dirIn, SystemStringView keyPath) {
/* Build map of existing nodes to write-through later */
std::unordered_map<std::string, const Node*> fileNodes;
std::unordered_map<std::string, const Node*> dirNodes;
if (nodeIn) {
fileNodes.reserve(nodeIn->size());
dirNodes.reserve(nodeIn->size());
for (const Node& ch : *nodeIn) {
if (ch.getKind() == Node::Kind::File)
fileNodes.insert(std::make_pair(ch.getName(), &ch));
else if (ch.getKind() == Node::Kind::Directory)
dirNodes.insert(std::make_pair(ch.getName(), &ch));
}
}
/* Merge this directory's files */
if (!dirIn.empty()) {
DirectoryEnumerator dEnum(dirIn, DirectoryEnumerator::Mode::DirsThenFilesSorted, false, false, true);
for (const DirectoryEnumerator::Entry& e : dEnum) {
SystemToDiscLocConv nameView(e.m_name, m_codepage);
SystemString chKeyPath = SystemString(keyPath) + _SYS_STR('/') + e.m_name;
if (e.m_isDir) {
auto search = dirNodes.find(nameView.disc_str().data());
if (search != dirNodes.cend()) {
if (!recursiveMergeNodes(ws, system, search->second, e.m_path.c_str(), chKeyPath))
return false;
dirNodes.erase(search);
} else {
if (!recursiveMergeNodes(ws, system, nullptr, e.m_path.c_str(), chKeyPath))
return false;
}
} else {
bool isDol;
bool isSys = IsSystemFile(e.m_name, isDol);
if (system ^ isSys)
continue;
fileNodes.erase(nameView.disc_str().data());
size_t fileSz = ROUND_UP_32(e.m_fileSz);
uint64_t fileOff = userAllocate(fileSz, ws);
if (fileOff == UINT64_MAX)
return false;
m_fileOffsetsSizes[chKeyPath] = std::make_pair(fileOff, fileSz);
std::unique_ptr<IFileIO::IReadStream> rs = NewFileIO(e.m_path)->beginReadStream();
if (!rs)
return false;
size_t xferSz = 0;
if (isDol) {
bool patched;
xferSz = PatchDOL(*rs, ws, e.m_fileSz, patched);
m_parent.m_progressCB(m_parent.getProgressFactor(),
e.m_name + (patched ? _SYS_STR(" [PATCHED]") : _SYS_STR("")), xferSz);
++m_parent.m_progressIdx;
} else {
char buf[0x8000];
while (xferSz < e.m_fileSz) {
size_t rdSz = rs->read(buf, nod::min(size_t(0x8000ul), e.m_fileSz - xferSz));
if (!rdSz)
break;
ws.write(buf, rdSz);
xferSz += rdSz;
m_parent.m_progressCB(m_parent.getProgressFactorMidFile(xferSz, e.m_fileSz), e.m_name, xferSz);
}
++m_parent.m_progressIdx;
}
for (size_t i = 0; i < fileSz - xferSz; ++i)
ws.write("\xff", 1);
}
}
}
/* Write-through remaining dir nodes */
for (const auto& p : dirNodes) {
DiscLocToSystemConv sysName(p.second->getName(), m_codepage);
SystemString chKeyPath = SystemString(keyPath) + _SYS_STR('/') + sysName.c_str();
if (!recursiveMergeNodes(ws, system, p.second, {}, chKeyPath))
return false;
}
/* Write-through remaining file nodes */
for (const auto& p : fileNodes) {
const Node& ch = *p.second;
DiscLocToSystemConv sysName(ch.getName(), m_codepage);
SystemString chKeyPath = SystemString(keyPath) + _SYS_STR('/') + sysName.c_str();
bool isDol;
bool isSys = IsSystemFile(sysName.sys_str(), isDol);
if (system ^ isSys)
continue;
size_t fileSz = ROUND_UP_32(ch.size());
uint64_t fileOff = userAllocate(fileSz, ws);
if (fileOff == UINT64_MAX)
return false;
m_fileOffsetsSizes[chKeyPath] = std::make_pair(fileOff, fileSz);
std::unique_ptr<IPartReadStream> rs = ch.beginReadStream();
if (!rs)
return false;
size_t xferSz = 0;
if (isDol) {
xferSz = ch.size();
std::unique_ptr<uint8_t[]> dolBuf = ch.getBuf();
bool patched;
PatchDOL(dolBuf, xferSz, patched);
ws.write(dolBuf.get(), xferSz);
m_parent.m_progressCB(m_parent.getProgressFactor(),
SystemString(sysName.sys_str()) + (patched ? _SYS_STR(" [PATCHED]") : _SYS_STR("")),
xferSz);
++m_parent.m_progressIdx;
} else {
char buf[0x8000];
while (xferSz < ch.size()) {
size_t rdSz = rs->read(buf, nod::min(size_t(0x8000), size_t(ch.size() - xferSz)));
if (!rdSz)
break;
ws.write(buf, rdSz);
xferSz += rdSz;
m_parent.m_progressCB(m_parent.getProgressFactorMidFile(xferSz, ch.size()), sysName.sys_str(), xferSz);
}
++m_parent.m_progressIdx;
}
for (size_t i = 0; i < fileSz - xferSz; ++i)
ws.write("\xff", 1);
}
return true;
}
bool DiscBuilderBase::PartitionBuilderBase::recursiveMergeFST(const Node* nodeIn, SystemStringView dirIn,
std::function<void(void)> incParents,
size_t parentDirIdx, SystemStringView keyPath) {
/* Build map of existing nodes to write-through later */
std::unordered_map<std::string, const Node*> fileNodes;
std::unordered_map<std::string, const Node*> dirNodes;
if (nodeIn) {
fileNodes.reserve(nodeIn->size());
dirNodes.reserve(nodeIn->size());
for (const Node& ch : *nodeIn) {
if (ch.getKind() == Node::Kind::File)
fileNodes.insert(std::make_pair(ch.getName(), &ch));
else if (ch.getKind() == Node::Kind::Directory)
dirNodes.insert(std::make_pair(ch.getName(), &ch));
}
}
/* Merge this directory's files */
if (!dirIn.empty()) {
DirectoryEnumerator dEnum(dirIn, DirectoryEnumerator::Mode::DirsThenFilesSorted, false, false, true);
for (const DirectoryEnumerator::Entry& e : dEnum) {
SystemToDiscLocConv nameView(e.m_name, m_codepage);
SystemString chKeyPath = SystemString(keyPath) + _SYS_STR('/') + e.m_name;
if (e.m_isDir) {
size_t dirNodeIdx = m_buildNodes.size();
m_buildNodes.emplace_back(true, m_buildNameOff, parentDirIdx, dirNodeIdx + 1);
addBuildName(e.m_name);
incParents();
auto search = dirNodes.find(nameView.disc_str().data());
if (search != dirNodes.cend()) {
if (!recursiveMergeFST(search->second, e.m_path.c_str(),
[&]() {
m_buildNodes[dirNodeIdx].incrementLength();
incParents();
},
dirNodeIdx, chKeyPath))
return false;
dirNodes.erase(search);
} else {
if (!recursiveMergeFST(nullptr, e.m_path.c_str(),
[&]() {
m_buildNodes[dirNodeIdx].incrementLength();
incParents();
},
dirNodeIdx, chKeyPath))
return false;
}
} else {
fileNodes.erase(nameView.disc_str().data());
std::pair<uint64_t, uint64_t> fileOffSz = m_fileOffsetsSizes.at(chKeyPath);
m_buildNodes.emplace_back(false, m_buildNameOff, packOffset(fileOffSz.first), fileOffSz.second);
addBuildName(e.m_name);
incParents();
}
}
}
/* Write-through remaining dir nodes */
for (const auto& p : dirNodes) {
DiscLocToSystemConv sysName(p.second->getName(), m_codepage);
SystemString chKeyPath = SystemString(keyPath) + _SYS_STR('/') + sysName.sys_str().data();
size_t dirNodeIdx = m_buildNodes.size();
m_buildNodes.emplace_back(true, m_buildNameOff, parentDirIdx, dirNodeIdx + 1);
addBuildName(sysName.sys_str());
incParents();
if (!recursiveMergeFST(p.second, {},
[&]() {
m_buildNodes[dirNodeIdx].incrementLength();
incParents();
},
dirNodeIdx, chKeyPath))
return false;
}
/* Write-through remaining file nodes */
for (const auto& p : fileNodes) {
const Node& ch = *p.second;
DiscLocToSystemConv sysName(ch.getName(), m_codepage);
SystemString chKeyPath = SystemString(keyPath) + _SYS_STR('/') + sysName.sys_str().data();
std::pair<uint64_t, uint64_t> fileOffSz = m_fileOffsetsSizes.at(chKeyPath);
m_buildNodes.emplace_back(false, m_buildNameOff, packOffset(fileOffSz.first), fileOffSz.second);
addBuildName(sysName.sys_str());
incParents();
}
return true;
}
bool DiscBuilderBase::PartitionBuilderBase::RecursiveCalculateTotalSize(uint64_t& totalSz, const Node* nodeIn,
SystemStringView dirIn, Codepage_t codepage) {
/* Build map of existing nodes to write-through later */
std::unordered_map<std::string, const Node*> fileNodes;
std::unordered_map<std::string, const Node*> dirNodes;
if (nodeIn) {
fileNodes.reserve(nodeIn->size());
dirNodes.reserve(nodeIn->size());
for (const Node& ch : *nodeIn) {
if (ch.getKind() == Node::Kind::File)
fileNodes.insert(std::make_pair(ch.getName(), &ch));
else if (ch.getKind() == Node::Kind::Directory)
dirNodes.insert(std::make_pair(ch.getName(), &ch));
}
}
/* Merge this directory's files */
if (!dirIn.empty()) {
DirectoryEnumerator dEnum(dirIn, DirectoryEnumerator::Mode::DirsThenFilesSorted, false, false, true);
for (const DirectoryEnumerator::Entry& e : dEnum) {
SystemToDiscLocConv nameView(e.m_name, codepage);
if (e.m_isDir) {
auto search = dirNodes.find(nameView.disc_str().data());
if (search != dirNodes.cend()) {
if (!RecursiveCalculateTotalSize(totalSz, search->second, e.m_path.c_str(), codepage))
return false;
dirNodes.erase(search);
} else {
if (!RecursiveCalculateTotalSize(totalSz, nullptr, e.m_path.c_str(), codepage))
return false;
}
} else {
fileNodes.erase(nameView.disc_str().data());
totalSz += ROUND_UP_32(e.m_fileSz);
}
}
}
/* Write-through remaining dir nodes */
for (const auto& p : dirNodes) {
if (!RecursiveCalculateTotalSize(totalSz, p.second, {}, codepage))
return false;
}
/* Write-through remaining file nodes */
for (const auto& p : fileNodes) {
const Node& ch = *p.second;
totalSz += ROUND_UP_32(ch.size());
}
return true;
}
bool DiscBuilderBase::PartitionBuilderBase::buildFromDirectory(IPartWriteStream& ws, SystemStringView dirIn) {
if (dirIn.empty()) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("all arguments must be supplied to buildFromDirectory()")));
return false;
}
SystemString dirStr(dirIn);
SystemString basePath = m_isWii ? dirStr + _SYS_STR("/") + getKindString(m_kind) : dirStr;
SystemString dolIn = basePath + _SYS_STR("/sys/main.dol");
SystemString filesIn = basePath + _SYS_STR("/files");
/* 1st pass - Tally up total progress steps */
m_parent.m_progressTotal += 2; /* Prep and DOL */
recursiveBuildNodesPre(filesIn.c_str());
/* Clear file */
m_parent.m_progressCB(m_parent.getProgressFactor(), _SYS_STR("Preparing output image"), -1);
++m_parent.m_progressIdx;
/* Add root node */
m_buildNodes.emplace_back(true, m_buildNameOff, 0, 1);
addBuildName(_SYS_STR("<root>"));
/* Write Boot DOL first (first thing seeked to after Apploader) */
{
Sstat dolStat;
if (Stat(dolIn.c_str(), &dolStat)) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("unable to stat {}")), dolIn);
return false;
}
size_t fileSz = ROUND_UP_32(dolStat.st_size);
uint64_t fileOff = userAllocate(fileSz, ws);
if (fileOff == UINT64_MAX)
return false;
m_dolOffset = fileOff;
m_dolSize = fileSz;
std::unique_ptr<IFileIO::IReadStream> rs = NewFileIO(dolIn.c_str())->beginReadStream();
if (!rs)
return false;
bool patched;
size_t xferSz = PatchDOL(*rs, ws, dolStat.st_size, patched);
m_parent.m_progressCB(m_parent.getProgressFactor(), dolIn + (patched ? _SYS_STR(" [PATCHED]") : _SYS_STR("")),
xferSz);
++m_parent.m_progressIdx;
for (size_t i = 0; i < fileSz - xferSz; ++i)
ws.write("\xff", 1);
}
/* Gather files in root directory */
if (!recursiveBuildNodes(ws, true, filesIn.c_str()))
return false;
if (!recursiveBuildNodes(ws, false, filesIn.c_str()))
return false;
if (!recursiveBuildFST(filesIn.c_str(), [&]() { m_buildNodes[0].incrementLength(); }, 0))
return false;
return true;
}
std::optional<uint64_t> DiscBuilderBase::PartitionBuilderBase::CalculateTotalSizeBuild(SystemStringView dirIn,
PartitionKind kind, bool isWii, Codepage_t codepage) {
SystemString dirStr(dirIn);
SystemString basePath = isWii ? dirStr + _SYS_STR("/") + getKindString(kind) : dirStr;
SystemString dolIn = basePath + _SYS_STR("/sys/main.dol");
SystemString filesIn = basePath + _SYS_STR("/files");
Sstat dolStat;
if (Stat(dolIn.c_str(), &dolStat)) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("unable to stat {}")), dolIn);
return std::nullopt;
}
uint64_t totalSz = ROUND_UP_32(dolStat.st_size);
if (!RecursiveCalculateTotalSize(totalSz, nullptr, filesIn.c_str(), codepage))
return std::nullopt;
return totalSz;
}
bool DiscBuilderBase::PartitionBuilderBase::mergeFromDirectory(IPartWriteStream& ws, const IPartition* partIn,
SystemStringView dirIn) {
if (dirIn.empty()) {
LogModule.report(logvisor::Error, FMT_STRING(_SYS_STR("all arguments must be supplied to mergeFromDirectory()")));
return false;
}
SystemString dirStr(dirIn);
SystemString basePath = m_isWii ? dirStr + _SYS_STR("/") + getKindString(m_kind) : dirStr;
SystemString filesIn = basePath + _SYS_STR("/files");
/* 1st pass - Tally up total progress steps */
m_parent.m_progressTotal += 2; /* Prep and DOL */
recursiveMergeNodesPre(&partIn->getFSTRoot(), filesIn.c_str());
/* Clear file */
m_parent.m_progressCB(m_parent.getProgressFactor(), _SYS_STR("Preparing output image"), -1);
++m_parent.m_progressIdx;
/* Add root node */
m_buildNodes.emplace_back(true, m_buildNameOff, 0, 1);
addBuildName(_SYS_STR("<root>"));
/* Write Boot DOL first (first thing seeked to after Apploader) */
{
size_t xferSz = partIn->getDOLSize();
size_t fileSz = ROUND_UP_32(xferSz);
uint64_t fileOff = userAllocate(fileSz, ws);
if (fileOff == UINT64_MAX)
return false;
m_dolOffset = fileOff;
m_dolSize = fileSz;
std::unique_ptr<uint8_t[]> dolBuf = partIn->getDOLBuf();
bool patched;
PatchDOL(dolBuf, xferSz, patched);
ws.write(dolBuf.get(), xferSz);
m_parent.m_progressCB(m_parent.getProgressFactor(),
SystemString(_SYS_STR("<boot-dol>")) + (patched ? _SYS_STR(" [PATCHED]") : _SYS_STR("")),
xferSz);
++m_parent.m_progressIdx;
for (size_t i = 0; i < fileSz - xferSz; ++i)
ws.write("\xff", 1);
}
/* Gather files in root directory */
SystemString keyPath;
if (!recursiveMergeNodes(ws, true, &partIn->getFSTRoot(), filesIn.c_str(), keyPath))
return false;
if (!recursiveMergeNodes(ws, false, &partIn->getFSTRoot(), filesIn.c_str(), keyPath))
return false;
if (!recursiveMergeFST(&partIn->getFSTRoot(), filesIn.c_str(), [&]() { m_buildNodes[0].incrementLength(); }, 0, keyPath))
return false;
return true;
}
std::optional<uint64_t> DiscBuilderBase::PartitionBuilderBase::CalculateTotalSizeMerge(const IPartition* partIn,
SystemStringView dirIn, Codepage_t codepage) {
SystemString dirStr(dirIn);
SystemString basePath = partIn->isWii() ? dirStr + _SYS_STR("/") + getKindString(partIn->getKind()) : dirStr;
SystemString filesIn = basePath + _SYS_STR("/files");
uint64_t totalSz = ROUND_UP_32(partIn->getDOLSize());
if (!RecursiveCalculateTotalSize(totalSz, &partIn->getFSTRoot(), filesIn.c_str(), codepage))
return std::nullopt;
return totalSz;
}
} // namespace nod