From ed4ee26fe6d6ec113efc54d188c5864f44c1e14d Mon Sep 17 00:00:00 2001 From: KiritoDv Date: Fri, 6 Dec 2024 14:44:36 -0600 Subject: [PATCH] Implemented rom preprocessing system and comptool support --- src/Companion.cpp | 40 ++++++++++++++++- src/preprocess/CompTool.cpp | 86 +++++++++++++++++++++++++++++++++++++ src/preprocess/CompTool.h | 18 ++++++++ src/utils/Decompressor.cpp | 6 +-- src/utils/Decompressor.h | 2 +- 5 files changed, 147 insertions(+), 5 deletions(-) create mode 100644 src/preprocess/CompTool.cpp create mode 100644 src/preprocess/CompTool.h diff --git a/src/Companion.cpp b/src/Companion.cpp index 3fca4e38..199d6490 100644 --- a/src/Companion.cpp +++ b/src/Companion.cpp @@ -79,6 +79,8 @@ #include "factories/naudio/v1/BookFactory.h" #include "factories/naudio/v1/SequenceFactory.h" +#include "preprocess/CompTool.h" + using namespace std::chrono; namespace fs = std::filesystem; @@ -1013,8 +1015,44 @@ void Companion::Process() { } auto rom = !isDirectoryMode ? config[this->gCartridge->GetHash()] : config; - auto cfg = rom["config"]; + if(rom["preprocess"]) { + auto preprocess = rom["preprocess"]; + for(auto job = preprocess.begin(); job != preprocess.end(); job++) { + auto name = job->first.as(); + auto item = job->second; + auto method = GetSafeNode(item, "method"); + if (method == "mio0-comptool") { + auto type = GetSafeNode(item, "type"); + auto target = GetSafeNode(item, "target"); + auto restart = GetSafeNode(item, "restart"); + + if (type == "decompress") { + this->gRomData = CompTool::Decompress(this->gRomData); + this->gCartridge = std::make_shared(this->gRomData); + this->gCartridge->Initialize(); + + auto hash = this->gCartridge->GetHash(); + + SPDLOG_INFO("ROM decompressed to {}", hash); + + if (hash != target) { + throw std::runtime_error("Hash mismatch"); + } + + if(restart){ + rom = config[this->gCartridge->GetHash()]; + } + } else { + throw std::runtime_error("Only decompression is supported"); + } + } else { + throw std::runtime_error("Invalid preprocess method"); + } + } + } + + auto cfg = rom["config"]; if(!cfg) { SPDLOG_ERROR("No config found for {}", !isDirectoryMode ? this->gCartridge->GetHash() : GetSafeNode(config, "folder")); return; diff --git a/src/preprocess/CompTool.cpp b/src/preprocess/CompTool.cpp new file mode 100644 index 00000000..9f585ef3 --- /dev/null +++ b/src/preprocess/CompTool.cpp @@ -0,0 +1,86 @@ +#include "CompTool.h" +#include "utils/Decompressor.h" +#include "lib/binarytools/BinaryWriter.h" +#include "lib/binarytools/BinaryReader.h" +#include + +uint32_t CompTool::FindFileTable(std::vector& rom) { + uint8_t query_one[] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x50, 0x00, 0x00, 0x00, 0x00 }; + uint8_t query_two[] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x60, 0x00, 0x00, 0x00, 0x00 }; + + for(size_t i = 0; i < rom.size() - sizeof(query_one); i++) { + if(memcmp(rom.data() + i, query_one, sizeof(query_one)) == 0){ + return i; + } + + if(memcmp(rom.data() + i, query_two, sizeof(query_two)) == 0){ + return i; + } + } + + throw std::runtime_error("Failed to find file table"); +} + +std::vector CompTool::Decompress(std::vector rom){ + LUS::BinaryReader basefile((char*) rom.data(), rom.size()); + basefile.SetEndianness(Torch::Endianness::Big); + + LUS::BinaryWriter decompfile; + decompfile.SetEndianness(Torch::Endianness::Big); + decompfile.Write((uint8_t)0x80); + + uint32_t table = CompTool::FindFileTable(rom); + uint32_t count = 0; + while (true){ + auto entry = table + 0x10 * count; + basefile.Seek(entry, LUS::SeekOffsetType::Start); + + auto v_begin = basefile.ReadInt32(); + auto p_begin = basefile.ReadInt32(); + auto p_end = basefile.ReadInt32(); + auto comp_flag = basefile.ReadInt32(); + + auto p_size = p_end - p_begin; + auto v_size = (int32_t) 0; + DataChunk* decoded = nullptr; + + if(v_begin == 0 && p_end == 0){ + break; + } + + basefile.Seek(p_begin, LUS::SeekOffsetType::Start); + + auto bytes = new uint8_t[p_size]; + basefile.Read((char*) bytes, p_size); + + switch ((CompType) comp_flag) { + case CompType::UNCOMPRESSED: + v_size = p_size; + break; + case CompType::COMPRESSED: + decoded = Decompressor::Decode(std::vector(bytes, bytes + p_size), 0, CompressionType::MIO0, true); + bytes = decoded->data; + v_size = decoded->size; + break; + default: + throw std::runtime_error("Invalid compression flag. There may be a problem with your ROM."); + } + + decompfile.Seek(v_begin, LUS::SeekOffsetType::Start); + decompfile.Write((char*) bytes, v_size); + auto v_end = v_begin + v_size; + + decompfile.Seek(entry + 4, LUS::SeekOffsetType::Start); + decompfile.Write(v_begin); + decompfile.Write(v_end); + decompfile.Write((uint32_t) CompType::UNCOMPRESSED); + count++; + } + + decompfile.Seek(0x10, LUS::SeekOffsetType::Start); + decompfile.Write(0xA7D5F194); // CRC1 + decompfile.Write(0xFE3DF761); // CRC2 + + auto result = decompfile.ToVector(); + return { (uint8_t*) result.data(), (uint8_t*) result.data() + result.size() }; +} \ No newline at end of file diff --git a/src/preprocess/CompTool.h b/src/preprocess/CompTool.h new file mode 100644 index 00000000..15b7db5a --- /dev/null +++ b/src/preprocess/CompTool.h @@ -0,0 +1,18 @@ +#pragma once + +#include +#include +#include + +enum class CompType { + UNCOMPRESSED, + COMPRESSED, + UNKNOWN +}; + +class CompTool { +public: + static std::vector Decompress(std::vector rom); +private: + static uint32_t FindFileTable(std::vector& rom); +}; \ No newline at end of file diff --git a/src/utils/Decompressor.cpp b/src/utils/Decompressor.cpp index 95af7b1b..1fee3820 100644 --- a/src/utils/Decompressor.cpp +++ b/src/utils/Decompressor.cpp @@ -12,9 +12,9 @@ extern "C" { std::unordered_map gCachedChunks; -DataChunk* Decompressor::Decode(const std::vector& buffer, const uint32_t offset, const CompressionType type) { +DataChunk* Decompressor::Decode(const std::vector& buffer, const uint32_t offset, const CompressionType type, bool ignoreCache) { - if(gCachedChunks.contains(offset)){ + if(!ignoreCache && gCachedChunks.contains(offset)){ return gCachedChunks[offset]; } @@ -208,4 +208,4 @@ void Decompressor::ClearCache() { delete value->data; } gCachedChunks.clear(); -} +} \ No newline at end of file diff --git a/src/utils/Decompressor.h b/src/utils/Decompressor.h index 4d2056aa..4aff9429 100644 --- a/src/utils/Decompressor.h +++ b/src/utils/Decompressor.h @@ -32,7 +32,7 @@ struct DecompressedData { class Decompressor { public: - static DataChunk* Decode(const std::vector& buffer, uint32_t offset, CompressionType type); + static DataChunk* Decode(const std::vector& buffer, uint32_t offset, CompressionType type, bool ignoreCache = false); static DataChunk* DecodeTKMK00(const std::vector& buffer, const uint32_t offset, const uint32_t size, const uint32_t alpha); static DecompressedData AutoDecode(YAML::Node& node, std::vector& buffer, std::optional size = std::nullopt); static DecompressedData AutoDecode(uint32_t offset, std::optional size, std::vector& buffer);