diff options
| author | Tilka <tilkax@gmail.com> | 2020-06-21 11:40:58 +0100 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2020-06-21 11:40:58 +0100 |
| commit | 99822518994d62def4e336f6746d6a2e2d108fb2 (patch) | |
| tree | 3c3ba54bb6df32a32fa72de0b9e46e09b36b14ff /Source/Core | |
| parent | 03e0d2c820a8aed809799e86321024244aa6791a (diff) | |
| parent | 660d81a10b4a8ff9a61631e69fb93e9010bccbc0 (diff) | |
Merge pull request #8538 from JosJuice/wia
Add support for the WIA and RVZ disc image formats
Diffstat (limited to 'Source/Core')
23 files changed, 4086 insertions, 97 deletions
diff --git a/Source/Core/Core/Boot/Boot.cpp b/Source/Core/Core/Boot/Boot.cpp index a4f85208e3..a1ab46ed3c 100644 --- a/Source/Core/Core/Boot/Boot.cpp +++ b/Source/Core/Core/Boot/Boot.cpp @@ -159,7 +159,7 @@ BootParameters::GenerateFromFile(std::vector<std::string> paths, paths.clear(); static const std::unordered_set<std::string> disc_image_extensions = { - {".gcm", ".iso", ".tgc", ".wbfs", ".ciso", ".gcz", ".dol", ".elf"}}; + {".gcm", ".iso", ".tgc", ".wbfs", ".ciso", ".gcz", ".wia", ".rvz", ".dol", ".elf"}}; if (disc_image_extensions.find(extension) != disc_image_extensions.end() || is_drive) { std::unique_ptr<DiscIO::VolumeDisc> disc = DiscIO::CreateDisc(path); diff --git a/Source/Core/DiscIO/Blob.cpp b/Source/Core/DiscIO/Blob.cpp index 20fe23e9ca..50146c9c91 100644 --- a/Source/Core/DiscIO/Blob.cpp +++ b/Source/Core/DiscIO/Blob.cpp @@ -20,6 +20,7 @@ #include "DiscIO/DriveBlob.h" #include "DiscIO/FileBlob.h" #include "DiscIO/TGCBlob.h" +#include "DiscIO/WIABlob.h" #include "DiscIO/WbfsBlob.h" namespace DiscIO @@ -205,6 +206,10 @@ std::unique_ptr<BlobReader> CreateBlobReader(const std::string& filename) return TGCFileReader::Create(std::move(file)); case WBFS_MAGIC: return WbfsFileReader::Create(std::move(file), filename); + case WIA_MAGIC: + return WIAFileReader::Create(std::move(file), filename); + case RVZ_MAGIC: + return RVZFileReader::Create(std::move(file), filename); default: if (auto directory_blob = DirectoryBlobReader::Create(filename)) return std::move(directory_blob); diff --git a/Source/Core/DiscIO/Blob.h b/Source/Core/DiscIO/Blob.h index e8846753d1..e2a1c3c9e9 100644 --- a/Source/Core/DiscIO/Blob.h +++ b/Source/Core/DiscIO/Blob.h @@ -25,6 +25,8 @@ namespace DiscIO { +enum class WIARVZCompressionType : u32; + // Increment CACHE_REVISION (GameFileCache.cpp) if the enum below is modified enum class BlobType { @@ -34,7 +36,9 @@ enum class BlobType GCZ, CISO, WBFS, - TGC + TGC, + WIA, + RVZ, }; class BlobReader @@ -172,5 +176,9 @@ bool ConvertToGCZ(BlobReader* infile, const std::string& infile_path, bool ConvertToPlain(BlobReader* infile, const std::string& infile_path, const std::string& outfile_path, CompressCB callback = nullptr, void* arg = nullptr); +bool ConvertToWIAOrRVZ(BlobReader* infile, const std::string& infile_path, + const std::string& outfile_path, bool rvz, + WIARVZCompressionType compression_type, int compression_level, + int chunk_size, CompressCB callback = nullptr, void* arg = nullptr); } // namespace DiscIO diff --git a/Source/Core/DiscIO/CMakeLists.txt b/Source/Core/DiscIO/CMakeLists.txt index 6437bf35d1..4c0f04f78a 100644 --- a/Source/Core/DiscIO/CMakeLists.txt +++ b/Source/Core/DiscIO/CMakeLists.txt @@ -21,6 +21,8 @@ add_library(discio FileSystemGCWii.h Filesystem.cpp Filesystem.h + LaggedFibonacciGenerator.cpp + LaggedFibonacciGenerator.h MultithreadedCompressor.h NANDImporter.cpp NANDImporter.h @@ -42,6 +44,10 @@ add_library(discio VolumeWii.h WbfsBlob.cpp WbfsBlob.h + WIABlob.cpp + WIABlob.h + WIACompression.cpp + WIACompression.h WiiEncryptionCache.cpp WiiEncryptionCache.h WiiSaveBanner.cpp @@ -49,6 +55,11 @@ add_library(discio ) target_link_libraries(discio +PUBLIC + BZip2::BZip2 + LibLZMA::LibLZMA + zstd + PRIVATE minizip pugixml diff --git a/Source/Core/DiscIO/DiscIO.vcxproj b/Source/Core/DiscIO/DiscIO.vcxproj index ee1d2ecd36..10f870fe69 100644 --- a/Source/Core/DiscIO/DiscIO.vcxproj +++ b/Source/Core/DiscIO/DiscIO.vcxproj @@ -55,6 +55,7 @@ <ClCompile Include="FileBlob.cpp" /> <ClCompile Include="Filesystem.cpp" /> <ClCompile Include="FileSystemGCWii.cpp" /> + <ClCompile Include="LaggedFibonacciGenerator.cpp" /> <ClCompile Include="NANDImporter.cpp" /> <ClCompile Include="ScrubbedBlob.cpp" /> <ClCompile Include="TGCBlob.cpp" /> @@ -65,6 +66,8 @@ <ClCompile Include="VolumeWad.cpp" /> <ClCompile Include="VolumeWii.cpp" /> <ClCompile Include="WbfsBlob.cpp" /> + <ClCompile Include="WIABlob.cpp" /> + <ClCompile Include="WIACompression.cpp" /> <ClCompile Include="WiiEncryptionCache.cpp" /> <ClCompile Include="WiiSaveBanner.cpp" /> </ItemGroup> @@ -80,6 +83,7 @@ <ClInclude Include="FileBlob.h" /> <ClInclude Include="Filesystem.h" /> <ClInclude Include="FileSystemGCWii.h" /> + <ClInclude Include="LaggedFibonacciGenerator.h" /> <ClInclude Include="MultithreadedCompressor.h" /> <ClInclude Include="NANDImporter.h" /> <ClInclude Include="ScrubbedBlob.h" /> @@ -91,6 +95,8 @@ <ClInclude Include="VolumeWad.h" /> <ClInclude Include="VolumeWii.h" /> <ClInclude Include="WbfsBlob.h" /> + <ClInclude Include="WIABlob.h" /> + <ClInclude Include="WIACompression.h" /> <ClInclude Include="WiiEncryptionCache.h" /> <ClInclude Include="WiiSaveBanner.h" /> </ItemGroup> @@ -110,6 +116,15 @@ <ProjectReference Include="$(ExternalsDir)pugixml\pugixml.vcxproj"> <Project>{38fee76f-f347-484b-949c-b4649381cffb}</Project> </ProjectReference> + <ProjectReference Include="$(ExternalsDir)bzip2\bzip2.vcxproj"> + <Project>{055a775f-b4f5-4970-9240-f6cf7661f37b}</Project> + </ProjectReference> + <ProjectReference Include="$(ExternalsDir)liblzma\liblzma.vcxproj"> + <Project>{1d8c51d2-ffa4-418e-b183-9f42b6a6717e}</Project> + </ProjectReference> + <ProjectReference Include="$(ExternalsDir)zstd\zstd.vcxproj"> + <Project>{1bea10f3-80ce-4bc4-9331-5769372cdf99}</Project> + </ProjectReference> </ItemGroup> <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" /> <ImportGroup Label="ExtensionTargets"> diff --git a/Source/Core/DiscIO/DiscIO.vcxproj.filters b/Source/Core/DiscIO/DiscIO.vcxproj.filters index 43cf82fa03..fd67a4487f 100644 --- a/Source/Core/DiscIO/DiscIO.vcxproj.filters +++ b/Source/Core/DiscIO/DiscIO.vcxproj.filters @@ -90,6 +90,15 @@ <ClCompile Include="ScrubbedBlob.cpp"> <Filter>Volume\Blob</Filter> </ClCompile> + <ClCompile Include="WIABlob.cpp"> + <Filter>Volume\Blob</Filter> + </ClCompile> + <ClCompile Include="LaggedFibonacciGenerator.cpp"> + <Filter>Volume\Blob</Filter> + </ClCompile> + <ClCompile Include="WIACompression.cpp"> + <Filter>Volume\Blob</Filter> + </ClCompile> </ItemGroup> <ItemGroup> <ClInclude Include="DiscScrubber.h"> @@ -164,6 +173,15 @@ <ClInclude Include="MultithreadedCompressor.h"> <Filter>Volume\Blob</Filter> </ClInclude> + <ClInclude Include="WIABlob.h"> + <Filter>Volume\Blob</Filter> + </ClInclude> + <ClInclude Include="LaggedFibonacciGenerator.h"> + <Filter>Volume\Blob</Filter> + </ClInclude> + <ClInclude Include="WIACompression.h"> + <Filter>Volume\Blob</Filter> + </ClInclude> </ItemGroup> <ItemGroup> <Text Include="CMakeLists.txt" /> diff --git a/Source/Core/DiscIO/LaggedFibonacciGenerator.cpp b/Source/Core/DiscIO/LaggedFibonacciGenerator.cpp new file mode 100644 index 0000000000..e5539b334d --- /dev/null +++ b/Source/Core/DiscIO/LaggedFibonacciGenerator.cpp @@ -0,0 +1,212 @@ +// This file is under the public domain. + +#include "DiscIO/LaggedFibonacciGenerator.h" + +#include <algorithm> +#include <cstddef> +#include <cstring> + +#include "Common/Align.h" +#include "Common/Assert.h" +#include "Common/CommonTypes.h" +#include "Common/Swap.h" + +namespace DiscIO +{ +void LaggedFibonacciGenerator::SetSeed(const u32 seed[SEED_SIZE]) +{ + SetSeed(reinterpret_cast<const u8*>(seed)); +} + +void LaggedFibonacciGenerator::SetSeed(const u8 seed[SEED_SIZE * sizeof(u32)]) +{ + m_position_bytes = 0; + + for (size_t i = 0; i < SEED_SIZE; ++i) + m_buffer[i] = Common::swap32(seed + i * sizeof(u32)); + + Initialize(false); +} + +size_t LaggedFibonacciGenerator::GetSeed(const u8* data, size_t size, size_t data_offset, + u32 seed_out[SEED_SIZE]) +{ + if ((reinterpret_cast<uintptr_t>(data) - data_offset) % alignof(u32) != 0) + { + ASSERT(false); + return 0; + } + + // For code simplicity, only include whole u32 words when regenerating the seed. It would be + // possible to get rid of this restriction and use a few additional bytes, but it's probably more + // effort than it's worth considering that junk data often starts or ends on 4-byte offsets. + const size_t bytes_to_skip = Common::AlignUp(data_offset, sizeof(u32)) - data_offset; + const u32* u32_data = reinterpret_cast<const u32*>(data + bytes_to_skip); + const size_t u32_size = (size - bytes_to_skip) / sizeof(u32); + const size_t u32_data_offset = (data_offset + bytes_to_skip) / sizeof(u32); + + LaggedFibonacciGenerator lfg; + if (!GetSeed(u32_data, u32_size, u32_data_offset, &lfg, seed_out)) + return false; + + lfg.m_position_bytes = data_offset % (LFG_K * sizeof(u32)); + + const u8* end = data + size; + size_t reconstructed_bytes = 0; + while (data < end && lfg.GetByte() == *data) + { + ++reconstructed_bytes; + ++data; + } + return reconstructed_bytes; +} + +bool LaggedFibonacciGenerator::GetSeed(const u32* data, size_t size, size_t data_offset, + LaggedFibonacciGenerator* lfg, u32 seed_out[SEED_SIZE]) +{ + if (size < LFG_K) + return false; + + // If the data doesn't look like something we can regenerate, return early to save time + if (!std::all_of(data, data + LFG_K, [](u32 x) { + return (Common::swap32(x) & 0x00C00000) == (Common::swap32(x) >> 2 & 0x00C00000); + })) + { + return false; + } + + const size_t data_offset_mod_k = data_offset % LFG_K; + const size_t data_offset_div_k = data_offset / LFG_K; + + std::copy(data, data + LFG_K - data_offset_mod_k, lfg->m_buffer.data() + data_offset_mod_k); + std::copy(data + LFG_K - data_offset_mod_k, data + LFG_K, lfg->m_buffer.data()); + + lfg->Backward(0, data_offset_mod_k); + + for (size_t i = 0; i < data_offset_div_k; ++i) + lfg->Backward(); + + if (!lfg->Reinitialize(seed_out)) + return false; + + for (size_t i = 0; i < data_offset_div_k; ++i) + lfg->Forward(); + + return true; +} + +void LaggedFibonacciGenerator::GetBytes(size_t count, u8* out) +{ + while (count > 0) + { + const size_t length = std::min(count, LFG_K * sizeof(u32) - m_position_bytes); + + std::memcpy(out, reinterpret_cast<u8*>(m_buffer.data()) + m_position_bytes, length); + + m_position_bytes += length; + count -= length; + out += length; + + if (m_position_bytes == LFG_K * sizeof(u32)) + { + Forward(); + m_position_bytes = 0; + } + } +} + +u8 LaggedFibonacciGenerator::GetByte() +{ + const u8 result = reinterpret_cast<u8*>(m_buffer.data())[m_position_bytes]; + + ++m_position_bytes; + + if (m_position_bytes == LFG_K * sizeof(u32)) + { + Forward(); + m_position_bytes = 0; + } + + return result; +} + +void LaggedFibonacciGenerator::Forward(size_t count) +{ + m_position_bytes += count; + while (m_position_bytes >= LFG_K * sizeof(u32)) + { + Forward(); + m_position_bytes -= LFG_K * sizeof(u32); + } +} + +void LaggedFibonacciGenerator::Forward() +{ + for (size_t i = 0; i < LFG_J; ++i) + m_buffer[i] ^= m_buffer[i + LFG_K - LFG_J]; + + for (size_t i = LFG_J; i < LFG_K; ++i) + m_buffer[i] ^= m_buffer[i - LFG_J]; +} + +void LaggedFibonacciGenerator::Backward(size_t start_word, size_t end_word) +{ + const size_t loop_end = std::max(LFG_J, start_word); + for (size_t i = std::min(end_word, LFG_K); i > loop_end; --i) + m_buffer[i - 1] ^= m_buffer[i - 1 - LFG_J]; + + for (size_t i = std::min(end_word, LFG_J); i > start_word; --i) + m_buffer[i - 1] ^= m_buffer[i - 1 + LFG_K - LFG_J]; +} + +bool LaggedFibonacciGenerator::Reinitialize(u32 seed_out[SEED_SIZE]) +{ + for (size_t i = 0; i < 4; ++i) + Backward(); + + for (u32& x : m_buffer) + x = Common::swap32(x); + + // Reconstruct the bits which are missing due to the output code shifting by 18 instead of 16. + // Unfortunately we can't reconstruct bits 16 and 17 (counting LSB as 0) for the first word, + // but the observable result (when shifting by 18 instead of 16) is not affected by this. + for (size_t i = 0; i < SEED_SIZE; ++i) + { + m_buffer[i] = (m_buffer[i] & 0xFF00FFFF) | (m_buffer[i] << 2 & 0x00FC0000) | + ((m_buffer[i + 16] ^ m_buffer[i + 15]) << 9 & 0x00030000); + } + + for (size_t i = 0; i < SEED_SIZE; ++i) + seed_out[i] = Common::swap32(m_buffer[i]); + + return Initialize(true); +} + +bool LaggedFibonacciGenerator::Initialize(bool check_existing_data) +{ + for (size_t i = SEED_SIZE; i < LFG_K; ++i) + { + const u32 calculated = (m_buffer[i - 17] << 23) ^ (m_buffer[i - 16] >> 9) ^ m_buffer[i - 1]; + + if (check_existing_data) + { + const u32 actual = (m_buffer[i] & 0xFF00FFFF) | (m_buffer[i] << 2 & 0x00FC0000); + if ((calculated & 0xFFFCFFFF) != actual) + return false; + } + + m_buffer[i] = calculated; + } + + // Instead of doing the "shift by 18 instead of 16" oddity when actually outputting the data, + // we can do the shifting (and byteswapping) at this point to make the output code simpler. + for (u32& x : m_buffer) + x = Common::swap32((x & 0xFF00FFFF) | ((x >> 2) & 0x00FF0000)); + + for (size_t i = 0; i < 4; ++i) + Forward(); + + return true; +} + +} // namespace DiscIO diff --git a/Source/Core/DiscIO/LaggedFibonacciGenerator.h b/Source/Core/DiscIO/LaggedFibonacciGenerator.h new file mode 100644 index 0000000000..9520700d05 --- /dev/null +++ b/Source/Core/DiscIO/LaggedFibonacciGenerator.h @@ -0,0 +1,51 @@ +// This file is under the public domain. + +#pragma once + +#include <array> +#include <cstddef> + +#include "Common/CommonTypes.h" + +namespace DiscIO +{ +class LaggedFibonacciGenerator +{ +public: + static constexpr size_t SEED_SIZE = 17; + + // Reconstructs a seed and writes it to seed_out, then returns the number of bytes which can + // be reconstructed using that seed. Can return any number between 0 and size, inclusive. + // data - data_offset must be 4-byte aligned. + static size_t GetSeed(const u8* data, size_t size, size_t data_offset, u32 seed_out[SEED_SIZE]); + + // SetSeed must be called before using the functions below + void SetSeed(const u32 seed[SEED_SIZE]); + void SetSeed(const u8 seed[SEED_SIZE * sizeof(u32)]); + + // Outputs a number of bytes and advances the internal state by the same amount. + void GetBytes(size_t count, u8* out); + u8 GetByte(); + + // Advances the internal state like GetBytes, but without outputting data. O(N), like GetBytes. + void Forward(size_t count); + +private: + static bool GetSeed(const u32* data, size_t size, size_t data_offset, + LaggedFibonacciGenerator* lfg, u32 seed_out[SEED_SIZE]); + + void Forward(); + void Backward(size_t start_word = 0, size_t end_word = LFG_K); + + bool Reinitialize(u32 seed_out[SEED_SIZE]); + bool Initialize(bool check_existing_data); + + static constexpr size_t LFG_K = 521; + static constexpr size_t LFG_J = 32; + + std::array<u32, LFG_K> m_buffer; + + size_t m_position_bytes = 0; +}; + +} // namespace DiscIO diff --git a/Source/Core/DiscIO/VolumeWii.cpp b/Source/Core/DiscIO/VolumeWii.cpp index 3bbe347211..6e2b40041f 100644 --- a/Source/Core/DiscIO/VolumeWii.cpp +++ b/Source/Core/DiscIO/VolumeWii.cpp @@ -201,18 +201,9 @@ bool VolumeWii::Read(u64 offset, u64 length, u8* buffer, const Partition& partit if (!m_reader->Read(block_offset_on_disc, BLOCK_TOTAL_SIZE, read_buffer.data())) return false; - // Decrypt the block's data. - // 0x3D0 - 0x3DF in read_buffer will be overwritten, - // but that won't affect anything, because we won't - // use the content of read_buffer anymore after this - mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, BLOCK_DATA_SIZE, &read_buffer[0x3D0], - &read_buffer[BLOCK_HEADER_SIZE], m_last_decrypted_block_data); + // Decrypt the block's data + DecryptBlockData(read_buffer.data(), m_last_decrypted_block_data, aes_context); m_last_decrypted_block = block_offset_on_disc; - - // The only thing we currently use from the 0x000 - 0x3FF part - // of the block is the IV (at 0x3D0), but it also contains SHA-1 - // hashes that IOS uses to check that discs aren't tampered with. - // http://wiibrew.org/wiki/Wii_Disc#Encrypted } // Copy the decrypted data @@ -482,14 +473,10 @@ bool VolumeWii::CheckBlockIntegrity(u64 block_index, const std::vector<u8>& encr return false; HashBlock hashes; - u8 iv[16] = {0}; - mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, sizeof(HashBlock), iv, - encrypted_data.data(), reinterpret_cast<u8*>(&hashes)); + DecryptBlockHashes(encrypted_data.data(), &hashes, aes_context); u8 cluster_data[BLOCK_DATA_SIZE]; - std::memcpy(iv, encrypted_data.data() + 0x3D0, 16); - mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, sizeof(cluster_data), iv, - encrypted_data.data() + sizeof(HashBlock), cluster_data); + DecryptBlockData(encrypted_data.data(), cluster_data, aes_context); for (u32 hash_index = 0; hash_index < 31; ++hash_index) { @@ -532,54 +519,33 @@ bool VolumeWii::CheckBlockIntegrity(u64 block_index, const Partition& partition) return CheckBlockIntegrity(block_index, cluster, partition); } -bool VolumeWii::EncryptGroup(u64 offset, u64 partition_data_offset, - u64 partition_data_decrypted_size, - const std::array<u8, AES_KEY_SIZE>& key, BlobReader* blob, - std::array<u8, GROUP_TOTAL_SIZE>* out) +bool VolumeWii::HashGroup(const std::array<u8, BLOCK_DATA_SIZE> in[BLOCKS_PER_GROUP], + HashBlock out[BLOCKS_PER_GROUP], + const std::function<bool(size_t block)>& read_function) { - std::vector<std::array<u8, BLOCK_DATA_SIZE>> unencrypted_data(BLOCKS_PER_GROUP); - std::vector<HashBlock> unencrypted_hashes(BLOCKS_PER_GROUP); - std::array<std::future<void>, BLOCKS_PER_GROUP> hash_futures; - bool error_occurred = false; + bool success = true; for (size_t i = 0; i < BLOCKS_PER_GROUP; ++i) { - if (!error_occurred) - { - if (offset + (i + 1) * BLOCK_DATA_SIZE <= partition_data_decrypted_size) - { - if (!blob->ReadWiiDecrypted(offset + i * BLOCK_DATA_SIZE, BLOCK_DATA_SIZE, - unencrypted_data[i].data(), partition_data_offset)) - { - error_occurred = true; - } - } - else - { - unencrypted_data[i].fill(0); - } - } + if (read_function && success) + success = read_function(i); - hash_futures[i] = std::async(std::launch::async, [&unencrypted_data, &unencrypted_hashes, - &hash_futures, error_occurred, i]() { + hash_futures[i] = std::async(std::launch::async, [&in, &out, &hash_futures, success, i]() { const size_t h1_base = Common::AlignDown(i, 8); - if (!error_occurred) + if (success) { // H0 hashes for (size_t j = 0; j < 31; ++j) - { - mbedtls_sha1_ret(unencrypted_data[i].data() + j * 0x400, 0x400, - unencrypted_hashes[i].h0[j]); - } + mbedtls_sha1_ret(in[i].data() + j * 0x400, 0x400, out[i].h0[j]); // H0 padding - std::memset(unencrypted_hashes[i].padding_0, 0, sizeof(HashBlock::padding_0)); + std::memset(out[i].padding_0, 0, sizeof(HashBlock::padding_0)); // H1 hash - mbedtls_sha1_ret(reinterpret_cast<u8*>(unencrypted_hashes[i].h0), sizeof(HashBlock::h0), - unencrypted_hashes[h1_base].h1[i - h1_base]); + mbedtls_sha1_ret(reinterpret_cast<u8*>(out[i].h0), sizeof(HashBlock::h0), + out[h1_base].h1[i - h1_base]); } if (i % 8 == 7) @@ -587,21 +553,18 @@ bool VolumeWii::EncryptGroup(u64 offset, u64 partition_data_offset, for (size_t j = 0; j < 7; ++j) hash_futures[h1_base + j].get(); - if (!error_occurred) + if (success) { // H1 padding - std::memset(unencrypted_hashes[h1_base].padding_1, 0, sizeof(HashBlock::padding_1)); + std::memset(out[h1_base].padding_1, 0, sizeof(HashBlock::padding_1)); // H1 copies for (size_t j = 1; j < 8; ++j) - { - std::memcpy(unencrypted_hashes[h1_base + j].h1, unencrypted_hashes[h1_base].h1, - sizeof(HashBlock::h1)); - } + std::memcpy(out[h1_base + j].h1, out[h1_base].h1, sizeof(HashBlock::h1)); // H2 hash - mbedtls_sha1_ret(reinterpret_cast<u8*>(unencrypted_hashes[i].h1), sizeof(HashBlock::h1), - unencrypted_hashes[0].h2[h1_base / 8]); + mbedtls_sha1_ret(reinterpret_cast<u8*>(out[i].h1), sizeof(HashBlock::h1), + out[0].h2[h1_base / 8]); } if (i == BLOCKS_PER_GROUP - 1) @@ -609,17 +572,14 @@ bool VolumeWii::EncryptGroup(u64 offset, u64 partition_data_offset, for (size_t j = 0; j < 7; ++j) hash_futures[j * 8 + 7].get(); - if (!error_occurred) + if (success) { // H2 padding - std::memset(unencrypted_hashes[0].padding_2, 0, sizeof(HashBlock::padding_2)); + std::memset(out[0].padding_2, 0, sizeof(HashBlock::padding_2)); // H2 copies for (size_t j = 1; j < BLOCKS_PER_GROUP; ++j) - { - std::memcpy(unencrypted_hashes[j].h2, unencrypted_hashes[0].h2, - sizeof(HashBlock::h2)); - } + std::memcpy(out[j].h2, out[0].h2, sizeof(HashBlock::h2)); } } } @@ -629,9 +589,41 @@ bool VolumeWii::EncryptGroup(u64 offset, u64 partition_data_offset, // Wait for all the async tasks to finish hash_futures.back().get(); - if (error_occurred) + return success; +} + +bool VolumeWii::EncryptGroup( + u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size, + const std::array<u8, AES_KEY_SIZE>& key, BlobReader* blob, + std::array<u8, GROUP_TOTAL_SIZE>* out, + const std::function<void(HashBlock hash_blocks[BLOCKS_PER_GROUP])>& hash_exception_callback) +{ + std::vector<std::array<u8, BLOCK_DATA_SIZE>> unencrypted_data(BLOCKS_PER_GROUP); + std::vector<HashBlock> unencrypted_hashes(BLOCKS_PER_GROUP); + + const bool success = + HashGroup(unencrypted_data.data(), unencrypted_hashes.data(), [&](size_t block) { + if (offset + (block + 1) * BLOCK_DATA_SIZE <= partition_data_decrypted_size) + { + if (!blob->ReadWiiDecrypted(offset + block * BLOCK_DATA_SIZE, BLOCK_DATA_SIZE, + unencrypted_data[block].data(), partition_data_offset)) + { + return false; + } + } + else + { + unencrypted_data[block].fill(0); + } + return true; + }); + + if (!success) return false; + if (hash_exception_callback) + hash_exception_callback(unencrypted_hashes.data()); + const unsigned int threads = std::min(BLOCKS_PER_GROUP, std::max<unsigned int>(1, std::thread::hardware_concurrency())); @@ -667,4 +659,20 @@ bool VolumeWii::EncryptGroup(u64 offset, u64 partition_data_offset, return true; } +void VolumeWii::DecryptBlockHashes(const u8* in, HashBlock* out, mbedtls_aes_context* aes_context) +{ + std::array<u8, 16> iv; + iv.fill(0); + mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, sizeof(HashBlock), iv.data(), in, + reinterpret_cast<u8*>(out)); +} + +void VolumeWii::DecryptBlockData(const u8* in, u8* out, mbedtls_aes_context* aes_context) +{ + std::array<u8, 16> iv; + std::copy(&in[0x3d0], &in[0x3e0], iv.data()); + mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, BLOCK_DATA_SIZE, iv.data(), + &in[BLOCK_HEADER_SIZE], out); +} + } // namespace DiscIO diff --git a/Source/Core/DiscIO/VolumeWii.h b/Source/Core/DiscIO/VolumeWii.h index 4def1a9393..0f2bb43d01 100644 --- a/Source/Core/DiscIO/VolumeWii.h +++ b/Source/Core/DiscIO/VolumeWii.h @@ -5,6 +5,7 @@ #pragma once #include <array> +#include <functional> #include <map> #include <memory> #include <optional> @@ -97,9 +98,22 @@ public: u64 GetRawSize() const override; const BlobReader& GetBlobReader() const; + // The in parameter can either contain all the data to begin with, + // or read_function can write data into the in parameter when called. + // The latter lets reading run in parallel with hashing. + // This function returns false iff read_function returns false. + static bool HashGroup(const std::array<u8, BLOCK_DATA_SIZE> in[BLOCKS_PER_GROUP], + HashBlock out[BLOCKS_PER_GROUP], + const std::function<bool(size_t block)>& read_function = {}); + static bool EncryptGroup(u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size, const std::array<u8, AES_KEY_SIZE>& key, BlobReader* blob, - std::array<u8, GROUP_TOTAL_SIZE>* out); + std::array<u8, GROUP_TOTAL_SIZE>* out, + const std::function<void(HashBlock hash_blocks[BLOCKS_PER_GROUP])>& + hash_exception_callback = {}); + + static void DecryptBlockHashes(const u8* in, HashBlock* out, mbedtls_aes_context* aes_context); + static void DecryptBlockData(const u8* in, u8* out, mbedtls_aes_context* aes_context); protected: u32 GetOffsetShift() const override { return 2; } diff --git a/Source/Core/DiscIO/WIABlob.cpp b/Source/Core/DiscIO/WIABlob.cpp new file mode 100644 index 0000000000..6f64b3ba01 --- /dev/null +++ b/Source/Core/DiscIO/WIABlob.cpp @@ -0,0 +1,2023 @@ +// Copyright 2018 Dolphin Emulator Project +// Licensed under GPLv2+ +// Refer to the license.txt file included. + +#include "DiscIO/WIABlob.h" + +#include <algorithm> +#include <array> +#include <cinttypes> +#include <cstring> +#include <limits> +#include <map> +#include <memory> +#include <mutex> +#include <optional> +#include <type_traits> +#include <utility> + +#include <mbedtls/sha1.h> +#include <zstd.h> + +#include "Common/Align.h" +#include "Common/Assert.h" +#include "Common/CommonTypes.h" +#include "Common/File.h" +#include "Common/FileUtil.h" +#include "Common/Logging/Log.h" +#include "Common/MsgHandler.h" +#include "Common/ScopeGuard.h" +#include "Common/StringUtil.h" +#include "Common/Swap.h" + +#include "DiscIO/Blob.h" +#include "DiscIO/DiscExtractor.h" +#include "DiscIO/Filesystem.h" +#include "DiscIO/LaggedFibonacciGenerator.h" +#include "DiscIO/MultithreadedCompressor.h" +#include "DiscIO/Volume.h" +#include "DiscIO/VolumeWii.h" +#include "DiscIO/WIACompression.h" +#include "DiscIO/WiiEncryptionCache.h" + +namespace DiscIO +{ +static void PushBack(std::vector<u8>* vector, const u8* begin, const u8* end) +{ + const size_t offset_in_vector = vector->size(); + vector->resize(offset_in_vector + (end - begin)); + std::copy(begin, end, vector->data() + offset_in_vector); +} + +template <typename T> +static void PushBack(std::vector<u8>* vector, const T& x) +{ + static_assert(std::is_trivially_copyable_v<T>); + + const u8* x_ptr = reinterpret_cast<const u8*>(&x); + PushBack(vector, x_ptr, x_ptr + sizeof(T)); +} + +std::pair<int, int> GetAllowedCompressionLevels(WIARVZCompressionType compression_type) +{ + switch (compression_type) + { + case WIARVZCompressionType::Bzip2: + case WIARVZCompressionType::LZMA: + case WIARVZCompressionType::LZMA2: + return {1, 9}; + case WIARVZCompressionType::Zstd: + // The actual minimum level can be gotten by calling ZSTD_minCLevel(). However, returning that + // would make the UI rather weird, because it is a negative number with very large magnitude. + // Note: Level 0 is a special number which means "default level" (level 3 as of this writing). + return {1, ZSTD_maxCLevel()}; + default: + return {0, -1}; + } +} + +template <bool RVZ> +WIARVZFileReader<RVZ>::WIARVZFileReader(File::IOFile file, const std::string& path) + : m_file(std::move(file)), m_encryption_cache(this) +{ + m_valid = Initialize(path); +} + +template <bool RVZ> +WIARVZFileReader<RVZ>::~WIARVZFileReader() = default; + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::Initialize(const std::string& path) +{ + if (!m_file.Seek(0, SEEK_SET) || !m_file.ReadArray(&m_header_1, 1)) + return false; + + if ((!RVZ && m_header_1.magic != WIA_MAGIC) || (RVZ && m_header_1.magic != RVZ_MAGIC)) + return false; + + const u32 version = RVZ ? RVZ_VERSION : WIA_VERSION; + const u32 version_read_compatible = + RVZ ? RVZ_VERSION_READ_COMPATIBLE : WIA_VERSION_READ_COMPATIBLE; + + const u32 file_version = Common::swap32(m_header_1.version); + const u32 file_version_compatible = Common::swap32(m_header_1.version_compatible); + + if (version < file_version_compatible || version_read_compatible > file_version) + { + ERROR_LOG(DISCIO, "Unsupported version %s in %s", VersionToString(file_version).c_str(), + path.c_str()); + return false; + } + + SHA1 header_1_actual_hash; + mbedtls_sha1_ret(reinterpret_cast<const u8*>(&m_header_1), sizeof(m_header_1) - sizeof(SHA1), + header_1_actual_hash.data()); + if (m_header_1.header_1_hash != header_1_actual_hash) + return false; + + if (Common::swap64(m_header_1.wia_file_size) != m_file.GetSize()) + { + ERROR_LOG(DISCIO, "File size is incorrect for %s", path.c_str()); + return false; + } + + const u32 header_2_size = Common::swap32(m_header_1.header_2_size); + const u32 header_2_min_size = sizeof(WIAHeader2) - sizeof(WIAHeader2::compressor_data); + if (header_2_size < header_2_min_size) + return false; + + std::vector<u8> header_2(header_2_size); + if (!m_file.ReadBytes(header_2.data(), header_2.size())) + return false; + + SHA1 header_2_actual_hash; + mbedtls_sha1_ret(header_2.data(), header_2.size(), header_2_actual_hash.data()); + if (m_header_1.header_2_hash != header_2_actual_hash) + return false; + + std::memcpy(&m_header_2, header_2.data(), std::min(header_2.size(), sizeof(WIAHeader2))); + + if (m_header_2.compressor_data_size > sizeof(WIAHeader2::compressor_data) || + header_2_size < header_2_min_size + m_header_2.compressor_data_size) + { + return false; + } + + const u32 chunk_size = Common::swap32(m_header_2.chunk_size); + const auto is_power_of_two = [](u32 x) { return (x & (x - 1)) == 0; }; + if ((!RVZ || chunk_size < VolumeWii::BLOCK_TOTAL_SIZE || !is_power_of_two(chunk_size)) && + chunk_size % VolumeWii::GROUP_TOTAL_SIZE != 0) + { + return false; + } + + const u32 compression_type = Common::swap32(m_header_2.compression_type); + m_compression_type = static_cast<WIARVZCompressionType>(compression_type); + if (m_compression_type > (RVZ ? WIARVZCompressionType::Zstd : WIARVZCompressionType::LZMA2) || + (RVZ && m_compression_type == WIARVZCompressionType::Purge)) + { + ERROR_LOG(DISCIO, "Unsupported compression type %u in %s", compression_type, path.c_str()); + return false; + } + + const size_t number_of_partition_entries = Common::swap32(m_header_2.number_of_partition_entries); + const size_t partition_entry_size = Common::swap32(m_header_2.partition_entry_size); + std::vector<u8> partition_entries(partition_entry_size * number_of_partition_entries); + if (!m_file.Seek(Common::swap64(m_header_2.partition_entries_offset), SEEK_SET)) + return false; + if (!m_file.ReadBytes(partition_entries.data(), partition_entries.size())) + return false; + + SHA1 partition_entries_actual_hash; + mbedtls_sha1_ret(reinterpret_cast<const u8*>(partition_entries.data()), partition_entries.size(), + partition_entries_actual_hash.data()); + if (m_header_2.partition_entries_hash != partition_entries_actual_hash) + return false; + + const size_t copy_length = std::min(partition_entry_size, sizeof(PartitionEntry)); + const size_t memset_length = sizeof(PartitionEntry) - copy_length; + u8* ptr = partition_entries.data(); + m_partition_entries.resize(number_of_partition_entries); + for (size_t i = 0; i < number_of_partition_entries; ++i, ptr += partition_entry_size) + { + std::memcpy(&m_partition_entries[i], ptr, copy_length); + std::memset(reinterpret_cast<u8*>(&m_partition_entries[i]) + copy_length, 0, memset_length); + } + + for (size_t i = 0; i < m_partition_entries.size(); ++i) + { + const std::array<PartitionDataEntry, 2>& entries = m_partition_entries[i].data_entries; + + size_t non_empty_entries = 0; + for (size_t j = 0; j < entries.size(); ++j) + { + const u32 number_of_sectors = Common::swap32(entries[j].number_of_sectors); + if (number_of_sectors != 0) + { + ++non_empty_entries; + + const u32 last_sector = Common::swap32(entries[j].first_sector) + number_of_sectors; + m_data_entries.emplace(last_sector * VolumeWii::BLOCK_TOTAL_SIZE, DataEntry(i, j)); + } + } + + if (non_empty_entries > 1) + { + if (Common::swap32(entries[0].first_sector) > Common::swap32(entries[1].first_sector)) + return false; + } + } + + const u32 number_of_raw_data_entries = Common::swap32(m_header_2.number_of_raw_data_entries); + m_raw_data_entries.resize(number_of_raw_data_entries); + Chunk& raw_data_entries = + ReadCompressedData(Common::swap64(m_header_2.raw_data_entries_offset), + Common::swap32(m_header_2.raw_data_entries_size), + number_of_raw_data_entries * sizeof(RawDataEntry), m_compression_type); + if (!raw_data_entries.ReadAll(&m_raw_data_entries)) + return false; + + for (size_t i = 0; i < m_raw_data_entries.size(); ++i) + { + const RawDataEntry& entry = m_raw_data_entries[i]; + const u64 data_size = Common::swap64(entry.data_size); + if (data_size != 0) + m_data_entries.emplace(Common::swap64(entry.data_offset) + data_size, DataEntry(i)); + } + + const u32 number_of_group_entries = Common::swap32(m_header_2.number_of_group_entries); + m_group_entries.resize(number_of_group_entries); + Chunk& group_entries = + ReadCompressedData(Common::swap64(m_header_2.group_entries_offset), + Common::swap32(m_header_2.group_entries_size), + number_of_group_entries * sizeof(GroupEntry), m_compression_type); + if (!group_entries.ReadAll(&m_group_entries)) + return false; + + if (HasDataOverlap()) + return false; + + return true; +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::HasDataOverlap() const +{ + for (size_t i = 0; i < m_partition_entries.size(); ++i) + { + const std::array<PartitionDataEntry, 2>& entries = m_partition_entries[i].data_entries; + for (size_t j = 0; j < entries.size(); ++j) + { + if (Common::swap32(entries[j].number_of_sectors) == 0) + continue; + + const u64 data_offset = Common::swap32(entries[j].first_sector) * VolumeWii::BLOCK_TOTAL_SIZE; + const auto it = m_data_entries.upper_bound(data_offset); + if (it == m_data_entries.end()) + return true; // Not an overlap, but an error nonetheless + if (!it->second.is_partition || it->second.index != i || it->second.partition_data_index != j) + return true; // Overlap + } + } + + for (size_t i = 0; i < m_raw_data_entries.size(); ++i) + { + if (Common::swap64(m_raw_data_entries[i].data_size) == 0) + continue; + + const u64 data_offset = Common::swap64(m_raw_data_entries[i].data_offset); + const auto it = m_data_entries.upper_bound(data_offset); + if (it == m_data_entries.end()) + return true; // Not an overlap, but an error nonetheless + if (it->second.is_partition || it->second.index != i) + return true; // Overlap + } + + return false; +} + +template <bool RVZ> +std::unique_ptr<WIARVZFileReader<RVZ>> WIARVZFileReader<RVZ>::Create(File::IOFile file, + const std::string& path) +{ + std::unique_ptr<WIARVZFileReader> blob(new WIARVZFileReader(std::move(file), path)); + return blob->m_valid ? std::move(blob) : nullptr; +} + +template <bool RVZ> +BlobType WIARVZFileReader<RVZ>::GetBlobType() const +{ + return RVZ ? BlobType::RVZ : BlobType::WIA; +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::Read(u64 offset, u64 size, u8* out_ptr) +{ + if (offset + size > Common::swap64(m_header_1.iso_file_size)) + return false; + + if (offset < sizeof(WIAHeader2::disc_header)) + { + const u64 bytes_to_read = std::min(sizeof(WIAHeader2::disc_header) - offset, size); + std::memcpy(out_ptr, m_header_2.disc_header.data() + offset, bytes_to_read); + offset += bytes_to_read; + size -= bytes_to_read; + out_ptr += bytes_to_read; + } + + const u32 chunk_size = Common::swap32(m_header_2.chunk_size); + while (size > 0) + { + const auto it = m_data_entries.upper_bound(offset); + if (it == m_data_entries.end()) + return false; + + const DataEntry& data = it->second; + if (data.is_partition) + { + const PartitionEntry& partition = m_partition_entries[it->second.index]; + + const u32 partition_first_sector = Common::swap32(partition.data_entries[0].first_sector); + const u64 partition_data_offset = partition_first_sector * VolumeWii::BLOCK_TOTAL_SIZE; + + const u32 second_number_of_sectors = + Common::swap32(partition.data_entries[1].number_of_sectors); + const u32 partition_total_sectors = + second_number_of_sectors ? Common::swap32(partition.data_entries[1].first_sector) - + partition_first_sector + second_number_of_sectors : + Common::swap32(partition.data_entries[0].number_of_sectors); + + for (const PartitionDataEntry& partition_data : partition.data_entries) + { + if (size == 0) + return true; + + const u32 first_sector = Common::swap32(partition_data.first_sector); + const u32 number_of_sectors = Common::swap32(partition_data.number_of_sectors); + + const u64 data_offset = first_sector * VolumeWii::BLOCK_TOTAL_SIZE; + const u64 data_size = number_of_sectors * VolumeWii::BLOCK_TOTAL_SIZE; + + if (data_size == 0) + continue; + + if (data_offset + data_size <= offset) + continue; + + if (offset < data_offset) + return false; + + const u64 bytes_to_read = std::min(data_size - (offset - data_offset), size); + + m_exception_list.clear(); + m_write_to_exception_list = true; + m_exception_list_last_group_index = std::numeric_limits<u64>::max(); + Common::ScopeGuard guard([this] { m_write_to_exception_list = false; }); + + bool hash_exception_error = false; + if (!m_encryption_cache.EncryptGroups( + offset - partition_data_offset, bytes_to_read, out_ptr, partition_data_offset, + partition_total_sectors * VolumeWii::BLOCK_DATA_SIZE, partition.partition_key, + [this, &hash_exception_error]( + VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP], u64 offset) { + // EncryptGroups calls ReadWiiDecrypted, which calls ReadFromGroups, + // which populates m_exception_list when m_write_to_exception_list == true + if (!ApplyHashExceptions(m_exception_list, hash_blocks)) + hash_exception_error = true; + })) + { + return false; + } + if (hash_exception_error) + return false; + + offset += bytes_to_read; + size -= bytes_to_read; + out_ptr += bytes_to_read; + } + } + else + { + const RawDataEntry& raw_data = m_raw_data_entries[data.index]; + if (!ReadFromGroups(&offset, &size, &out_ptr, chunk_size, VolumeWii::BLOCK_TOTAL_SIZE, + Common::swap64(raw_data.data_offset), Common::swap64(raw_data.data_size), + Common::swap32(raw_data.group_index), + Common::swap32(raw_data.number_of_groups), 0)) + { + return false; + } + } + } + + return true; +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::SupportsReadWiiDecrypted() const +{ + return !m_partition_entries.empty(); +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::ReadWiiDecrypted(u64 offset, u64 size, u8* out_ptr, + u64 partition_data_offset) +{ + const u64 chunk_size = Common::swap32(m_header_2.chunk_size) * VolumeWii::BLOCK_DATA_SIZE / + VolumeWii::BLOCK_TOTAL_SIZE; + + const auto it = m_data_entries.upper_bound(partition_data_offset); + if (it == m_data_entries.end() || !it->second.is_partition) + return false; + + const PartitionEntry& partition = m_partition_entries[it->second.index]; + const u32 partition_first_sector = Common::swap32(partition.data_entries[0].first_sector); + if (partition_data_offset != partition_first_sector * VolumeWii::BLOCK_TOTAL_SIZE) + return false; + + for (const PartitionDataEntry& data : partition.data_entries) + { + if (size == 0) + return true; + + const u64 data_offset = + (Common::swap32(data.first_sector) - partition_first_sector) * VolumeWii::BLOCK_DATA_SIZE; + const u64 data_size = Common::swap32(data.number_of_sectors) * VolumeWii::BLOCK_DATA_SIZE; + + if (!ReadFromGroups(&offset, &size, &out_ptr, chunk_size, VolumeWii::BLOCK_DATA_SIZE, + data_offset, data_size, Common::swap32(data.group_index), + Common::swap32(data.number_of_groups), + std::max<u64>(1, chunk_size / VolumeWii::GROUP_DATA_SIZE))) + { + return false; + } + } + + return size == 0; +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::ReadFromGroups(u64* offset, u64* size, u8** out_ptr, u64 chunk_size, + u32 sector_size, u64 data_offset, u64 data_size, + u32 group_index, u32 number_of_groups, + u32 exception_lists) +{ + if (data_offset + data_size <= *offset) + return true; + + if (*offset < data_offset) + return false; + + const u64 skipped_data = data_offset % sector_size; + data_offset -= skipped_data; + data_size += skipped_data; + + const u64 start_group_index = (*offset - data_offset) / chunk_size; + for (u64 i = start_group_index; i < number_of_groups && (*size) > 0; ++i) + { + const u64 total_group_index = group_index + i; + if (total_group_index >= m_group_entries.size()) + return false; + + const GroupEntry group = m_group_entries[total_group_index]; + const u64 group_offset_in_data = i * chunk_size; + const u64 offset_in_group = *offset - group_offset_in_data - data_offset; + + chunk_size = std::min(chunk_size, data_size - group_offset_in_data); + + const u64 bytes_to_read = std::min(chunk_size - offset_in_group, *size); + u32 group_data_size = Common::swap32(group.data_size); + + WIARVZCompressionType compression_type = m_compression_type; + u32 rvz_packed_size = 0; + if constexpr (RVZ) + { + if ((group_data_size & 0x80000000) == 0) + compression_type = WIARVZCompressionType::None; + + group_data_size &= 0x7FFFFFFF; + + rvz_packed_size = Common::swap32(group.rvz_packed_size); + } + + if (group_data_size == 0) + { + std::memset(*out_ptr, 0, bytes_to_read); + } + else + { + const u64 group_offset_in_file = static_cast<u64>(Common::swap32(group.data_offset)) << 2; + + Chunk& chunk = + ReadCompressedData(group_offset_in_file, group_data_size, chunk_size, compression_type, + exception_lists, rvz_packed_size, group_offset_in_data); + + if (!chunk.Read(offset_in_group, bytes_to_read, *out_ptr)) + { + m_cached_chunk_offset = std::numeric_limits<u64>::max(); // Invalidate the cache + return false; + } + + if (m_write_to_exception_list && m_exception_list_last_group_index != total_group_index) + { + const u64 exception_list_index = offset_in_group / VolumeWii::GROUP_DATA_SIZE; + const u16 additional_offset = + static_cast<u16>(group_offset_in_data % VolumeWii::GROUP_DATA_SIZE / + VolumeWii::BLOCK_DATA_SIZE * VolumeWii::BLOCK_HEADER_SIZE); + chunk.GetHashExceptions(&m_exception_list, exception_list_index, additional_offset); + m_exception_list_last_group_index = total_group_index; + } + } + + *offset += bytes_to_read; + *size -= bytes_to_read; + *out_ptr += bytes_to_read; + } + + return true; +} + +template <bool RVZ> +typename WIARVZFileReader<RVZ>::Chunk& +WIARVZFileReader<RVZ>::ReadCompressedData(u64 offset_in_file, u64 compressed_size, + u64 decompressed_size, + WIARVZCompressionType compression_type, + u32 exception_lists, u32 rvz_packed_size, u64 data_offset) +{ + if (offset_in_file == m_cached_chunk_offset) + return m_cached_chunk; + + std::unique_ptr<Decompressor> decompressor; + switch (compression_type) + { + case WIARVZCompressionType::None: + decompressor = std::make_unique<NoneDecompressor>(); + break; + case WIARVZCompressionType::Purge: + decompressor = std::make_unique<PurgeDecompressor>(rvz_packed_size == 0 ? decompressed_size : + rvz_packed_size); + break; + case WIARVZCompressionType::Bzip2: + decompressor = std::make_unique<Bzip2Decompressor>(); + break; + case WIARVZCompressionType::LZMA: + decompressor = std::make_unique<LZMADecompressor>(false, m_header_2.compressor_data, + m_header_2.compressor_data_size); + break; + case WIARVZCompressionType::LZMA2: + decompressor = std::make_unique<LZMADecompressor>(true, m_header_2.compressor_data, + m_header_2.compressor_data_size); + break; + case WIARVZCompressionType::Zstd: + decompressor = std::make_unique<ZstdDecompressor>(); + break; + } + + const bool compressed_exception_lists = compression_type > WIARVZCompressionType::Purge; + + m_cached_chunk = + Chunk(&m_file, offset_in_file, compressed_size, decompressed_size, exception_lists, + compressed_exception_lists, rvz_packed_size, data_offset, std::move(decompressor)); + m_cached_chunk_offset = offset_in_file; + return m_cached_chunk; +} + +template <bool RVZ> +std::string WIARVZFileReader<RVZ>::VersionToString(u32 version) +{ + const u8 a = version >> 24; + const u8 b = (version >> 16) & 0xff; + const u8 c = (version >> 8) & 0xff; + const u8 d = version & 0xff; + + if (d == 0 || d == 0xff) + return StringFromFormat("%u.%02x.%02x", a, b, c); + else + return StringFromFormat("%u.%02x.%02x.beta%u", a, b, c, d); +} + +template <bool RVZ> +WIARVZFileReader<RVZ>::Chunk::Chunk() = default; + +template <bool RVZ> +WIARVZFileReader<RVZ>::Chunk::Chunk(File::IOFile* file, u64 offset_in_file, u64 compressed_size, + u64 decompressed_size, u32 exception_lists, + bool compressed_exception_lists, u32 rvz_packed_size, + u64 data_offset, std::unique_ptr<Decompressor> decompressor) + : m_file(file), m_offset_in_file(offset_in_file), m_exception_lists(exception_lists), + m_compressed_exception_lists(compressed_exception_lists), m_rvz_packed_size(rvz_packed_size), + m_data_offset(data_offset), m_decompressor(std::move(decompressor)) +{ + constexpr size_t MAX_SIZE_PER_EXCEPTION_LIST = + Common::AlignUp(VolumeWii::BLOCK_HEADER_SIZE, sizeof(SHA1)) / sizeof(SHA1) * + VolumeWii::BLOCKS_PER_GROUP * sizeof(HashExceptionEntry) + + sizeof(u16); + + m_out_bytes_allocated_for_exceptions = + m_compressed_exception_lists ? MAX_SIZE_PER_EXCEPTION_LIST * m_exception_lists : 0; + + m_in.data.resize(compressed_size); + m_out.data.resize(decompressed_size + m_out_bytes_allocated_for_exceptions); +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::Chunk::Read(u64 offset, u64 size, u8* out_ptr) +{ + if (!m_decompressor || !m_file || + offset + size > m_out.data.size() - m_out_bytes_allocated_for_exceptions) + { + return false; + } + + while (offset + size > m_out.bytes_written - m_out_bytes_used_for_exceptions) + { + u64 bytes_to_read; + if (offset + size == m_out.data.size()) + { + // Read all the remaining data. + bytes_to_read = m_in.data.size() - m_in.bytes_written; + } + else + { + // Pick a suitable amount of compressed data to read. The std::min line has to + // be as it is, but the rest is a bit arbitrary and can be changed if desired. + + // The compressed data is probably not much bigger than the decompressed data. + // Add a few bytes for possible compression overhead and for any hash exceptions. + bytes_to_read = + offset + size - (m_out.bytes_written - m_out_bytes_used_for_exceptions) + 0x100; + + // Align the access in an attempt to gain speed. But we don't actually know the + // block size of the underlying storage device, so we just use the Wii block size. + bytes_to_read = + Common::AlignUp(bytes_to_read + m_offset_in_file, VolumeWii::BLOCK_TOTAL_SIZE) - + m_offset_in_file; + + // Ensure we don't read too much. + bytes_to_read = std::min<u64>(m_in.data.size() - m_in.bytes_written, bytes_to_read); + } + + if (bytes_to_read == 0) + { + // Compressed size is larger than expected or decompressed size is smaller than expected + return false; + } + + if (!m_file->Seek(m_offset_in_file, SEEK_SET)) + return false; + if (!m_file->ReadBytes(m_in.data.data() + m_in.bytes_written, bytes_to_read)) + return false; + + m_offset_in_file += bytes_to_read; + m_in.bytes_written += bytes_to_read; + + if (m_exception_lists > 0 && !m_compressed_exception_lists) + { + if (!HandleExceptions(m_in.data.data(), m_in.data.size(), m_in.bytes_written, + &m_in_bytes_used_for_exceptions, true)) + { + return false; + } + + m_in_bytes_read = m_in_bytes_used_for_exceptions; + } + + if (m_exception_lists == 0 || m_compressed_exception_lists) + { + if (!Decompress()) + return false; + } + + if (m_exception_lists > 0 && m_compressed_exception_lists) + { + if (!HandleExceptions(m_out.data.data(), m_out_bytes_allocated_for_exceptions, + m_out.bytes_written, &m_out_bytes_used_for_exceptions, false)) + { + return false; + } + + if (m_rvz_packed_size != 0 && m_exception_lists == 0) + { + if (!Decompress()) + return false; + } + } + + if (m_exception_lists == 0) + { + const size_t expected_out_bytes = m_out.data.size() - m_out_bytes_allocated_for_exceptions + + m_out_bytes_used_for_exceptions; + + if (m_out.bytes_written > expected_out_bytes) + return false; // Decompressed size is larger than expected + + // The reason why we need the m_in.bytes_written == m_in.data.size() check as part of + // this conditional is because (for example) zstd can finish writing all data to m_out + // before becoming done if we've given it all input data except the checksum at the end. + if (m_out.bytes_written == expected_out_bytes && !m_decompressor->Done() && + m_in.bytes_written == m_in.data.size()) + { + return false; // Decompressed size is larger than expected + } + + if (m_decompressor->Done() && m_in_bytes_read != m_in.data.size()) + return false; // Compressed size is smaller than expected + } + } + + std::memcpy(out_ptr, m_out.data.data() + offset + m_out_bytes_used_for_exceptions, size); + return true; +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::Chunk::Decompress() +{ + if (m_rvz_packed_size != 0 && m_exception_lists == 0) + { + const size_t bytes_to_move = m_out.bytes_written - m_out_bytes_used_for_exceptions; + + DecompressionBuffer in{std::vector<u8>(bytes_to_move), bytes_to_move}; + std::memcpy(in.data.data(), m_out.data.data() + m_out_bytes_used_for_exceptions, bytes_to_move); + + m_out.bytes_written = m_out_bytes_used_for_exceptions; + + m_decompressor = std::make_unique<RVZPackDecompressor>(std::move(m_decompressor), std::move(in), + m_data_offset, m_rvz_packed_size); + + m_rvz_packed_size = 0; + } + + return m_decompressor->Decompress(m_in, &m_out, &m_in_bytes_read); +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::Chunk::HandleExceptions(const u8* data, size_t bytes_allocated, + size_t bytes_written, size_t* bytes_used, + bool align) +{ + while (m_exception_lists > 0) + { + if (sizeof(u16) + *bytes_used > bytes_allocated) + { + ERROR_LOG(DISCIO, "More hash exceptions than expected"); + return false; + } + if (sizeof(u16) + *bytes_used > bytes_written) + return true; + + const u16 exceptions = Common::swap16(data + *bytes_used); + + size_t exception_list_size = exceptions * sizeof(HashExceptionEntry) + sizeof(u16); + if (align && m_exception_lists == 1) + exception_list_size = Common::AlignUp(*bytes_used + exception_list_size, 4) - *bytes_used; + + if (exception_list_size + *bytes_used > bytes_allocated) + { + ERROR_LOG(DISCIO, "More hash exceptions than expected"); + return false; + } + if (exception_list_size + *bytes_used > bytes_written) + return true; + + *bytes_used += exception_list_size; + --m_exception_lists; + } + + return true; +} + +template <bool RVZ> +void WIARVZFileReader<RVZ>::Chunk::GetHashExceptions( + std::vector<HashExceptionEntry>* exception_list, u64 exception_list_index, + u16 additional_offset) const +{ + ASSERT(m_exception_lists == 0); + + const u8* data_start = m_compressed_exception_lists ? m_out.data.data() : m_in.data.data(); + const u8* data = data_start; + + for (u64 i = exception_list_index; i > 0; --i) + data += Common::swap16(data) * sizeof(HashExceptionEntry) + sizeof(u16); + + const u16 exceptions = Common::swap16(data); + data += sizeof(u16); + + for (size_t i = 0; i < exceptions; ++i) + { + std::memcpy(&exception_list->emplace_back(), data, sizeof(HashExceptionEntry)); + data += sizeof(HashExceptionEntry); + + u16& offset = exception_list->back().offset; + offset = Common::swap16(Common::swap16(offset) + additional_offset); + } + + ASSERT(data <= data_start + (m_compressed_exception_lists ? m_out_bytes_used_for_exceptions : + m_in_bytes_used_for_exceptions)); +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::ApplyHashExceptions( + const std::vector<HashExceptionEntry>& exception_list, + VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP]) +{ + for (const HashExceptionEntry& exception : exception_list) + { + const u16 offset = Common::swap16(exception.offset); + + const size_t block_index = offset / VolumeWii::BLOCK_HEADER_SIZE; + if (block_index > VolumeWii::BLOCKS_PER_GROUP) + return false; + + const size_t offset_in_block = offset % VolumeWii::BLOCK_HEADER_SIZE; + if (offset_in_block + sizeof(SHA1) > VolumeWii::BLOCK_HEADER_SIZE) + return false; + + std::memcpy(reinterpret_cast<u8*>(&hash_blocks[block_index]) + offset_in_block, &exception.hash, + sizeof(SHA1)); + } + + return true; +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::PadTo4(File::IOFile* file, u64* bytes_written) +{ + constexpr u32 ZEROES = 0; + const u64 bytes_to_write = Common::AlignUp(*bytes_written, 4) - *bytes_written; + if (bytes_to_write == 0) + return true; + + *bytes_written += bytes_to_write; + return file->WriteBytes(&ZEROES, bytes_to_write); +} + +template <bool RVZ> +void WIARVZFileReader<RVZ>::AddRawDataEntry(u64 offset, u64 size, int chunk_size, u32* total_groups, + std::vector<RawDataEntry>* raw_data_entries, + std::vector<DataEntry>* data_entries) +{ + constexpr size_t SKIP_SIZE = sizeof(WIAHeader2::disc_header); + const u64 skip = offset < SKIP_SIZE ? std::min(SKIP_SIZE - offset, size) : 0; + + offset += skip; + size -= skip; + + if (size == 0) + return; + + const u32 group_index = *total_groups; + const u32 groups = static_cast<u32>(Common::AlignUp(size, chunk_size) / chunk_size); + *total_groups += groups; + + data_entries->emplace_back(raw_data_entries->size()); + raw_data_entries->emplace_back(RawDataEntry{Common::swap64(offset), Common::swap64(size), + Common::swap32(group_index), Common::swap32(groups)}); +} + +template <bool RVZ> +typename WIARVZFileReader<RVZ>::PartitionDataEntry WIARVZFileReader<RVZ>::CreatePartitionDataEntry( + u64 offset, u64 size, u32 index, int chunk_size, u32* total_groups, + const std::vector<PartitionEntry>& partition_entries, std::vector<DataEntry>* data_entries) +{ + const u32 group_index = *total_groups; + const u64 rounded_size = Common::AlignDown(size, VolumeWii::BLOCK_TOTAL_SIZE); + const u32 groups = static_cast<u32>(Common::AlignUp(rounded_size, chunk_size) / chunk_size); + *total_groups += groups; + + data_entries->emplace_back(partition_entries.size(), index); + return PartitionDataEntry{Common::swap32(offset / VolumeWii::BLOCK_TOTAL_SIZE), + Common::swap32(size / VolumeWii::BLOCK_TOTAL_SIZE), + Common::swap32(group_index), Common::swap32(groups)}; +} + +template <bool RVZ> +ConversionResultCode WIARVZFileReader<RVZ>::SetUpDataEntriesForWriting( + const VolumeDisc* volume, int chunk_size, u64 iso_size, u32* total_groups, + std::vector<PartitionEntry>* partition_entries, std::vector<RawDataEntry>* raw_data_entries, + std::vector<DataEntry>* data_entries, std::vector<const FileSystem*>* partition_file_systems) +{ + std::vector<Partition> partitions; + if (volume && volume->IsEncryptedAndHashed()) + partitions = volume->GetPartitions(); + + std::sort(partitions.begin(), partitions.end(), + [](const Partition& a, const Partition& b) { return a.offset < b.offset; }); + + *total_groups = 0; + + u64 last_partition_end_offset = 0; + + const auto add_raw_data_entry = [&](u64 offset, u64 size) { + return AddRawDataEntry(offset, size, chunk_size, total_groups, raw_data_entries, data_entries); + }; + + const auto create_partition_data_entry = [&](u64 offset, u64 size, u32 index) { + return CreatePartitionDataEntry(offset, size, index, chunk_size, total_groups, + *partition_entries, data_entries); + }; + + for (const Partition& partition : partitions) + { + // If a partition is odd in some way that prevents us from encoding it as a partition, + // we encode it as raw data instead by skipping the current loop iteration. + // Partitions can always be encoded as raw data, but it is less space efficient. + + if (partition.offset < last_partition_end_offset) + { + WARN_LOG(DISCIO, "Overlapping partitions at %" PRIx64, partition.offset); + continue; + } + + if (volume->ReadSwapped<u32>(partition.offset, PARTITION_NONE) != u32(0x10001)) + { + // This looks more like garbage data than an actual partition. + // The values of data_offset and data_size will very likely also be garbage. + // Some WBFS writing programs scrub the SSBB Masterpiece partitions without + // removing them from the partition table, causing this problem. + WARN_LOG(DISCIO, "Invalid partition at %" PRIx64, partition.offset); + continue; + } + + std::optional<u64> data_offset = + volume->ReadSwappedAndShifted(partition.offset + 0x2b8, PARTITION_NONE); + std::optional<u64> data_size = + volume->ReadSwappedAndShifted(partition.offset + 0x2bc, PARTITION_NONE); + + if (!data_offset || !data_size) + return ConversionResultCode::ReadFailed; + + const u64 data_start = partition.offset + *data_offset; + const u64 data_end = data_start + *data_size; + + if (data_start % VolumeWii::BLOCK_TOTAL_SIZE != 0) + { + WARN_LOG(DISCIO, "Misaligned partition at %" PRIx64, partition.offset); + continue; + } + + if (*data_size < VolumeWii::BLOCK_TOTAL_SIZE) + { + WARN_LOG(DISCIO, "Very small partition at %" PRIx64, partition.offset); + continue; + } + + if (data_end > iso_size) + { + WARN_LOG(DISCIO, "Too large partition at %" PRIx64, partition.offset); + *data_size = iso_size - *data_offset - partition.offset; + } + + const std::optional<u64> fst_offset = GetFSTOffset(*volume, partition); + const std::optional<u64> fst_size = GetFSTSize(*volume, partition); + + if (!fst_offset || !fst_size) + return ConversionResultCode::ReadFailed; + + const IOS::ES::TicketReader& ticket = volume->GetTicket(partition); + if (!ticket.IsValid()) + return ConversionResultCode::ReadFailed; + + add_raw_data_entry(last_partition_end_offset, partition.offset - last_partition_end_offset); + + add_raw_data_entry(partition.offset, *data_offset); + + const u64 fst_end = volume->PartitionOffsetToRawOffset(*fst_offset + *fst_size, partition); + const u64 split_point = std::min( + data_end, Common::AlignUp(fst_end - data_start, VolumeWii::GROUP_TOTAL_SIZE) + data_start); + + PartitionEntry partition_entry; + partition_entry.partition_key = ticket.GetTitleKey(); + partition_entry.data_entries[0] = + create_partition_data_entry(data_start, split_point - data_start, 0); + partition_entry.data_entries[1] = + create_partition_data_entry(split_point, data_end - split_point, 1); + + // Note: We can't simply set last_partition_end_offset to data_end, + // because construct_partition_data_entry may have rounded it + last_partition_end_offset = + (Common::swap32(partition_entry.data_entries[1].first_sector) + + Common::swap32(partition_entry.data_entries[1].number_of_sectors)) * + VolumeWii::BLOCK_TOTAL_SIZE; + + partition_entries->emplace_back(std::move(partition_entry)); + partition_file_systems->emplace_back(volume->GetFileSystem(partition)); + } + + add_raw_data_entry(last_partition_end_offset, iso_size - last_partition_end_offset); + + return ConversionResultCode::Success; +} + +template <bool RVZ> +std::optional<std::vector<u8>> WIARVZFileReader<RVZ>::Compress(Compressor* compressor, + const u8* data, size_t size) +{ + if (compressor) + { + if (!compressor->Start() || !compressor->Compress(data, size) || !compressor->End()) + return std::nullopt; + + data = compressor->GetData(); + size = compressor->GetSize(); + } + + return std::vector<u8>(data, data + size); +} + +template <bool RVZ> +void WIARVZFileReader<RVZ>::SetUpCompressor(std::unique_ptr<Compressor>* compressor, + WIARVZCompressionType compression_type, + int compression_level, WIAHeader2* header_2) +{ + switch (compression_type) + { + case WIARVZCompressionType::None: + *compressor = nullptr; + break; + case WIARVZCompressionType::Purge: + *compressor = std::make_unique<PurgeCompressor>(); + break; + case WIARVZCompressionType::Bzip2: + *compressor = std::make_unique<Bzip2Compressor>(compression_level); + break; + case WIARVZCompressionType::LZMA: + case WIARVZCompressionType::LZMA2: + { + u8* compressor_data = nullptr; + u8* compressor_data_size = nullptr; + + if (header_2) + { + compressor_data = header_2->compressor_data; + compressor_data_size = &header_2->compressor_data_size; + } + + const bool lzma2 = compression_type == WIARVZCompressionType::LZMA2; + *compressor = std::make_unique<LZMACompressor>(lzma2, compression_level, compressor_data, + compressor_data_size); + break; + } + case WIARVZCompressionType::Zstd: + *compressor = std::make_unique<ZstdCompressor>(compression_level); + break; + } +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::TryReuse(std::map<ReuseID, GroupEntry>* reusable_groups, + std::mutex* reusable_groups_mutex, + OutputParametersEntry* entry) +{ + if (entry->reused_group) + return true; + + if (!entry->reuse_id) + return false; + + std::lock_guard guard(*reusable_groups_mutex); + const auto it = reusable_groups->find(*entry->reuse_id); + if (it == reusable_groups->end()) + return false; + + entry->reused_group = it->second; + return true; +} + +static bool AllAre(const std::vector<u8>& data, u8 x) +{ + return std::all_of(data.begin(), data.end(), [x](u8 y) { return x == y; }); +}; + +static bool AllAre(const u8* begin, const u8* end, u8 x) +{ + return std::all_of(begin, end, [x](u8 y) { return x == y; }); +}; + +static bool AllZero(const std::vector<u8>& data) +{ + return AllAre(data, 0); +}; + +static bool AllSame(const std::vector<u8>& data) +{ + return AllAre(data, data.front()); +}; + +static bool AllSame(const u8* begin, const u8* end) +{ + return AllAre(begin, end, *begin); +}; + +template <typename OutputParametersEntry> +static void RVZPack(const u8* in, OutputParametersEntry* out, u64 bytes_per_chunk, size_t chunks, + u64 total_size, u64 data_offset, u64 in_offset, bool multipart, + bool allow_junk_reuse, bool compression, const FileSystem* file_system) +{ + using Seed = std::array<u32, LaggedFibonacciGenerator::SEED_SIZE>; + struct JunkInfo + { + size_t start_offset; + Seed seed; + }; + + constexpr size_t SEED_SIZE = LaggedFibonacciGenerator::SEED_SIZE * sizeof(u32); + + // Maps end_offset -> (start_offset, seed) + std::map<size_t, JunkInfo> junk_info; + + size_t position = 0; + while (position < total_size) + { + // Skip the 0 to 32 zero bytes that typically come after a file + size_t zeroes = 0; + while (position + zeroes < total_size && in[in_offset + position + zeroes] == 0) + ++zeroes; + + // If there are very many zero bytes (perhaps the PRNG junk data has been scrubbed?) + // and we aren't using compression, it makes sense to encode the zero bytes as junk. + // If we are using compression, the compressor will likely encode zeroes better than we can + if (!compression && zeroes > SEED_SIZE) + junk_info.emplace(position + zeroes, JunkInfo{position, {}}); + + position += zeroes; + data_offset += zeroes; + + const size_t bytes_to_read = + std::min(Common::AlignUp(data_offset + 1, VolumeWii::BLOCK_TOTAL_SIZE) - data_offset, + total_size - position); + + const size_t data_offset_mod = static_cast<size_t>(data_offset % VolumeWii::BLOCK_TOTAL_SIZE); + + Seed seed; + const size_t bytes_reconstructed = LaggedFibonacciGenerator::GetSeed( + in + in_offset + position, bytes_to_read, data_offset_mod, seed.data()); + + if (bytes_reconstructed > 0) + junk_info.emplace(position + bytes_reconstructed, JunkInfo{position, seed}); + + if (file_system) + { + const std::unique_ptr<DiscIO::FileInfo> file_info = + file_system->FindFileInfo(data_offset + bytes_reconstructed); + + // If we're at a file and there's more space in this block after the file, + // continue after the file instead of skipping to the next block + if (file_info) + { + const u64 file_end_offset = file_info->GetOffset() + file_info->GetSize(); + if (file_end_offset < data_offset + bytes_to_read) + { + position += file_end_offset - data_offset; + data_offset = file_end_offset; + continue; + } + } + } + + position += bytes_to_read; + data_offset += bytes_to_read; + } + + for (size_t i = 0; i < chunks; ++i) + { + OutputParametersEntry& entry = out[i]; + if (entry.reused_group) + continue; + + u64 current_offset = i * bytes_per_chunk; + const u64 end_offset = std::min(current_offset + bytes_per_chunk, total_size); + + const bool store_junk_efficiently = allow_junk_reuse || !entry.reuse_id; + + // TODO: It would be possible to support skipping RVZ packing even when the chunk size is larger + // than 2 MiB (multipart == true), but it would be more effort than it's worth since Dolphin's + // converter doesn't expose chunk sizes larger than 2 MiB to the user anyway + bool first_loop_iteration = !multipart; + + while (current_offset < end_offset) + { + u64 next_junk_start = end_offset; + u64 next_junk_end = end_offset; + Seed* seed = nullptr; + if (store_junk_efficiently && end_offset - current_offset > SEED_SIZE) + { + const auto next_junk_it = junk_info.upper_bound(current_offset + SEED_SIZE); + if (next_junk_it != junk_info.end() && + next_junk_it->second.start_offset + SEED_SIZE < end_offset) + { + next_junk_start = std::max<u64>(current_offset, next_junk_it->second.start_offset); + next_junk_end = std::min<u64>(end_offset, next_junk_it->first); + seed = &next_junk_it->second.seed; + } + } + + if (first_loop_iteration) + { + if (next_junk_start == end_offset) + { + // Storing this chunk without RVZ packing would be inefficient, so store it without + PushBack(&entry.main_data, in + in_offset + current_offset, in + in_offset + end_offset); + break; + } + + first_loop_iteration = false; + } + + const u64 non_junk_bytes = next_junk_start - current_offset; + if (non_junk_bytes > 0) + { + const u8* ptr = in + in_offset + current_offset; + + PushBack(&entry.main_data, Common::swap32(static_cast<u32>(non_junk_bytes))); + PushBack(&entry.main_data, ptr, ptr + non_junk_bytes); + + current_offset += non_junk_bytes; + entry.rvz_packed_size += sizeof(u32) + non_junk_bytes; + } + + const u64 junk_bytes = next_junk_end - current_offset; + if (junk_bytes > 0) + { + PushBack(&entry.main_data, Common::swap32(static_cast<u32>(junk_bytes) | 0x80000000)); + PushBack(&entry.main_data, *seed); + + current_offset += junk_bytes; + entry.rvz_packed_size += sizeof(u32) + SEED_SIZE; + } + } + } +} + +template <typename OutputParametersEntry> +static void RVZPack(const u8* in, OutputParametersEntry* out, u64 size, u64 data_offset, + bool allow_junk_reuse, bool compression, const FileSystem* file_system) +{ + RVZPack(in, out, size, 1, size, data_offset, 0, false, allow_junk_reuse, compression, + file_system); +} + +template <bool RVZ> +ConversionResult<typename WIARVZFileReader<RVZ>::OutputParameters> +WIARVZFileReader<RVZ>::ProcessAndCompress(CompressThreadState* state, CompressParameters parameters, + const std::vector<PartitionEntry>& partition_entries, + const std::vector<DataEntry>& data_entries, + const FileSystem* file_system, + std::map<ReuseID, GroupEntry>* reusable_groups, + std::mutex* reusable_groups_mutex, + u64 chunks_per_wii_group, u64 exception_lists_per_chunk, + bool compressed_exception_lists, bool compression) +{ + std::vector<OutputParametersEntry> output_entries; + + if (!parameters.data_entry->is_partition) + { + OutputParametersEntry& entry = output_entries.emplace_back(); + std::vector<u8>& data = parameters.data; + + if (AllSame(data)) + entry.reuse_id = ReuseID{nullptr, data.size(), false, data.front()}; + + if constexpr (RVZ) + { + RVZPack(data.data(), output_entries.data(), data.size(), parameters.data_offset, true, + compression, file_system); + } + else + { + entry.main_data = std::move(data); + } + } + else + { + const PartitionEntry& partition_entry = partition_entries[parameters.data_entry->index]; + + mbedtls_aes_context aes_context; + mbedtls_aes_setkey_dec(&aes_context, partition_entry.partition_key.data(), 128); + + const u64 groups = Common::AlignUp(parameters.data.size(), VolumeWii::GROUP_TOTAL_SIZE) / + VolumeWii::GROUP_TOTAL_SIZE; + + ASSERT(parameters.data.size() % VolumeWii::BLOCK_TOTAL_SIZE == 0); + const u64 blocks = parameters.data.size() / VolumeWii::BLOCK_TOTAL_SIZE; + + const u64 blocks_per_chunk = chunks_per_wii_group == 1 ? + exception_lists_per_chunk * VolumeWii::BLOCKS_PER_GROUP : + VolumeWii::BLOCKS_PER_GROUP / chunks_per_wii_group; + + const u64 chunks = Common::AlignUp(blocks, blocks_per_chunk) / blocks_per_chunk; + + const u64 in_data_per_chunk = blocks_per_chunk * VolumeWii::BLOCK_TOTAL_SIZE; + const u64 out_data_per_chunk = blocks_per_chunk * VolumeWii::BLOCK_DATA_SIZE; + + const size_t first_chunk = output_entries.size(); + + const auto create_reuse_id = [&partition_entry, blocks, + blocks_per_chunk](u8 value, bool encrypted, u64 block) { + const u64 size = std::min(blocks - block, blocks_per_chunk) * VolumeWii::BLOCK_DATA_SIZE; + return ReuseID{&partition_entry.partition_key, size, encrypted, value}; + }; + + const u8* parameters_data_end = parameters.data.data() + parameters.data.size(); + for (u64 i = 0; i < chunks; ++i) + { + const u64 block_index = i * blocks_per_chunk; + + OutputParametersEntry& entry = output_entries.emplace_back(); + std::optional<ReuseID>& reuse_id = entry.reuse_id; + + // Set this chunk as reusable if the encrypted data is AllSame + const u8* data = parameters.data.data() + block_index * VolumeWii::BLOCK_TOTAL_SIZE; + if (AllSame(data, std::min(parameters_data_end, data + in_data_per_chunk))) + reuse_id = create_reuse_id(parameters.data.front(), true, i * blocks_per_chunk); + + TryReuse(reusable_groups, reusable_groups_mutex, &entry); + if (!entry.reused_group && reuse_id) + { + const auto it = std::find_if(output_entries.begin(), output_entries.begin() + i, + [reuse_id](const auto& e) { return e.reuse_id == reuse_id; }); + if (it != output_entries.begin() + i) + entry.reused_group = it->reused_group; + } + } + + if (!std::all_of(output_entries.begin(), output_entries.end(), + [](const OutputParametersEntry& entry) { return entry.reused_group; })) + { + const u64 number_of_exception_lists = + chunks_per_wii_group == 1 ? exception_lists_per_chunk : chunks; + std::vector<std::vector<HashExceptionEntry>> exception_lists(number_of_exception_lists); + + for (u64 i = 0; i < groups; ++i) + { + const u64 offset_of_group = i * VolumeWii::GROUP_TOTAL_SIZE; + const u64 write_offset_of_group = i * VolumeWii::GROUP_DATA_SIZE; + + const u64 blocks_in_this_group = + std::min<u64>(VolumeWii::BLOCKS_PER_GROUP, blocks - i * VolumeWii::BLOCKS_PER_GROUP); + + for (u32 j = 0; j < VolumeWii::BLOCKS_PER_GROUP; ++j) + { + if (j < blocks_in_this_group) + { + const u64 offset_of_block = offset_of_group + j * VolumeWii::BLOCK_TOTAL_SIZE; + VolumeWii::DecryptBlockData(parameters.data.data() + offset_of_block, + state->decryption_buffer[j].data(), &aes_context); + } + else + { + state->decryption_buffer[j].fill(0); + } + } + + VolumeWii::HashGroup(state->decryption_buffer.data(), state->hash_buffer.data()); + + for (u64 j = 0; j < blocks_in_this_group; ++j) + { + const u64 chunk_index = j / blocks_per_chunk; + const u64 block_index_in_chunk = j % blocks_per_chunk; + + if (output_entries[chunk_index].reused_group) + continue; + + const u64 exception_list_index = chunks_per_wii_group == 1 ? i : chunk_index; + + const u64 offset_of_block = offset_of_group + j * VolumeWii::BLOCK_TOTAL_SIZE; + const u64 hash_offset_of_block = block_index_in_chunk * VolumeWii::BLOCK_HEADER_SIZE; + + VolumeWii::HashBlock hashes; + VolumeWii::DecryptBlockHashes(parameters.data.data() + offset_of_block, &hashes, + &aes_context); + + const auto compare_hash = [&](size_t offset_in_block) { + ASSERT(offset_in_block + sizeof(SHA1) <= VolumeWii::BLOCK_HEADER_SIZE); + + const u8* desired_hash = reinterpret_cast<u8*>(&hashes) + offset_in_block; + const u8* computed_hash = + reinterpret_cast<u8*>(&state->hash_buffer[j]) + offset_in_block; + + // We want to store a hash exception either if there is a hash mismatch, or if this + // chunk might get reused in a context where it is paired up (within a 2 MiB Wii group) + // with chunks that are different from the chunks it currently is paired up with, since + // that affects the recalculated hashes. Chunks which have been marked as reusable at + // this point normally have zero matching hashes anyway, so this shouldn't waste space. + if ((chunks_per_wii_group != 1 && output_entries[chunk_index].reuse_id) || + !std::equal(desired_hash, desired_hash + sizeof(SHA1), computed_hash)) + { + const u64 hash_offset = hash_offset_of_block + offset_in_block; + ASSERT(hash_offset <= std::numeric_limits<u16>::max()); + + HashExceptionEntry& exception = exception_lists[exception_list_index].emplace_back(); + exception.offset = static_cast<u16>(Common::swap16(hash_offset)); + std::memcpy(exception.hash.data(), desired_hash, sizeof(SHA1)); + } + }; + + const auto compare_hashes = [&compare_hash](size_t offset, size_t size) { + for (size_t l = 0; l < size; l += sizeof(SHA1)) + // The std::min is to ensure that we don't go beyond the end of HashBlock with + // padding_2, which is 32 bytes long (not divisible by sizeof(SHA1), which is 20). + compare_hash(offset + std::min(l, size - sizeof(SHA1))); + }; + + using HashBlock = VolumeWii::HashBlock; + compare_hashes(offsetof(HashBlock, h0), sizeof(HashBlock::h0)); + compare_hashes(offsetof(HashBlock, padding_0), sizeof(HashBlock::padding_0)); + compare_hashes(offsetof(HashBlock, h1), sizeof(HashBlock::h1)); + compare_hashes(offsetof(HashBlock, padding_1), sizeof(HashBlock::padding_1)); + compare_hashes(offsetof(HashBlock, h2), sizeof(HashBlock::h2)); + compare_hashes(offsetof(HashBlock, padding_2), sizeof(HashBlock::padding_2)); + } + + static_assert(std::is_trivially_copyable_v<typename decltype( + CompressThreadState::decryption_buffer)::value_type>); + if constexpr (RVZ) + { + // We must not store junk efficiently for chunks that may get reused at a position + // which has a different value of data_offset % VolumeWii::BLOCK_TOTAL_SIZE + const bool allow_junk_reuse = chunks_per_wii_group == 1; + + const u64 bytes_per_chunk = std::min(out_data_per_chunk, VolumeWii::GROUP_DATA_SIZE); + const u64 total_size = blocks_in_this_group * VolumeWii::BLOCK_DATA_SIZE; + const u64 data_offset = parameters.data_offset + write_offset_of_group; + + RVZPack(state->decryption_buffer[0].data(), output_entries.data() + first_chunk, + bytes_per_chunk, chunks, total_size, data_offset, write_offset_of_group, + groups > 1, allow_junk_reuse, compression, file_system); + } + else + { + const u8* in_ptr = state->decryption_buffer[0].data(); + for (u64 j = 0; j < chunks; ++j) + { + OutputParametersEntry& entry = output_entries[first_chunk + j]; + + if (!entry.reused_group) + { + const u64 bytes_left = (blocks - j * blocks_per_chunk) * VolumeWii::BLOCK_DATA_SIZE; + const u64 bytes_to_write_total = std::min(out_data_per_chunk, bytes_left); + + if (i == 0) + entry.main_data.resize(bytes_to_write_total); + + const u64 bytes_to_write = std::min(bytes_to_write_total, VolumeWii::GROUP_DATA_SIZE); + + std::memcpy(entry.main_data.data() + write_offset_of_group, in_ptr, bytes_to_write); + + // Set this chunk as reusable if the decrypted data is AllSame. + // There is also a requirement that it lacks exceptions, but this is checked later + if (i == 0 && !entry.reuse_id) + { + if (AllSame(in_ptr, in_ptr + bytes_to_write)) + entry.reuse_id = create_reuse_id(*in_ptr, false, j * blocks_per_chunk); + } + else + { + if (entry.reuse_id && !entry.reuse_id->encrypted && + (!AllSame(in_ptr, in_ptr + bytes_to_write) || entry.reuse_id->value != *in_ptr)) + { + entry.reuse_id.reset(); + } + } + } + + in_ptr += out_data_per_chunk; + } + } + } + + for (size_t i = 0; i < exception_lists.size(); ++i) + { + OutputParametersEntry& entry = output_entries[chunks_per_wii_group == 1 ? 0 : i]; + if (entry.reused_group) + continue; + + const std::vector<HashExceptionEntry>& in = exception_lists[i]; + std::vector<u8>& out = entry.exception_lists; + + const u16 exceptions = Common::swap16(static_cast<u16>(in.size())); + PushBack(&out, exceptions); + for (const HashExceptionEntry& exception : in) + PushBack(&out, exception); + } + + for (u64 i = 0; i < output_entries.size(); ++i) + { + OutputParametersEntry& entry = output_entries[i]; + + // If this chunk was set as reusable because the decrypted data is AllSame, + // but it has exceptions, unmark it as reusable + if (entry.reuse_id && !entry.reuse_id->encrypted && !AllZero(entry.exception_lists)) + entry.reuse_id.reset(); + } + } + } + + for (OutputParametersEntry& entry : output_entries) + { + TryReuse(reusable_groups, reusable_groups_mutex, &entry); + if (entry.reused_group) + continue; + + // Special case - a compressed size of zero is treated by WIA as meaning the data is all zeroes + if (entry.reuse_id && !entry.reuse_id->encrypted && entry.reuse_id->value == 0) + { + entry.exception_lists.clear(); + entry.main_data.clear(); + if constexpr (RVZ) + { + entry.rvz_packed_size = 0; + entry.compressed = false; + } + continue; + } + + const auto pad_exception_lists = [&entry]() { + while (entry.exception_lists.size() % 4 != 0) + entry.exception_lists.push_back(0); + }; + + if (state->compressor) + { + if (!state->compressor->Start()) + return ConversionResultCode::InternalError; + } + + if (!entry.exception_lists.empty()) + { + if (compressed_exception_lists && state->compressor) + { + if (!state->compressor->Compress(entry.exception_lists.data(), + entry.exception_lists.size())) + { + return ConversionResultCode::InternalError; + } + } + else + { + if (!compressed_exception_lists) + pad_exception_lists(); + + if (state->compressor) + { + if (!state->compressor->AddPrecedingDataOnlyForPurgeHashing(entry.exception_lists.data(), + entry.exception_lists.size())) + { + return ConversionResultCode::InternalError; + } + } + } + } + + if (state->compressor) + { + if (!state->compressor->Compress(entry.main_data.data(), entry.main_data.size())) + return ConversionResultCode::InternalError; + if (!state->compressor->End()) + return ConversionResultCode::InternalError; + } + + bool compressed = !!state->compressor; + if constexpr (RVZ) + { + size_t uncompressed_size = entry.main_data.size(); + if (compressed_exception_lists) + uncompressed_size += Common::AlignUp(entry.exception_lists.size(), 4); + + compressed = state->compressor && state->compressor->GetSize() < uncompressed_size; + entry.compressed = compressed; + + if (!compressed) + pad_exception_lists(); + } + + if (compressed) + { + const u8* data = state->compressor->GetData(); + const size_t size = state->compressor->GetSize(); + + entry.main_data.resize(size); + std::copy(data, data + size, entry.main_data.data()); + + if (compressed_exception_lists) + entry.exception_lists.clear(); + } + } + + return OutputParameters{std::move(output_entries), parameters.bytes_read, parameters.group_index}; +} + +template <bool RVZ> +ConversionResultCode WIARVZFileReader<RVZ>::Output(std::vector<OutputParametersEntry>* entries, + File::IOFile* outfile, + std::map<ReuseID, GroupEntry>* reusable_groups, + std::mutex* reusable_groups_mutex, + GroupEntry* group_entry, u64* bytes_written) +{ + for (OutputParametersEntry& entry : *entries) + { + TryReuse(reusable_groups, reusable_groups_mutex, &entry); + if (entry.reused_group) + { + *group_entry = *entry.reused_group; + ++group_entry; + continue; + } + + if (*bytes_written >> 2 > std::numeric_limits<u32>::max()) + return ConversionResultCode::InternalError; + + ASSERT((*bytes_written & 3) == 0); + group_entry->data_offset = Common::swap32(static_cast<u32>(*bytes_written >> 2)); + + u32 data_size = static_cast<u32>(entry.exception_lists.size() + entry.main_data.size()); + if constexpr (RVZ) + { + data_size = (data_size & 0x7FFFFFFF) | (static_cast<u32>(entry.compressed) << 31); + group_entry->rvz_packed_size = Common::swap32(static_cast<u32>(entry.rvz_packed_size)); + } + group_entry->data_size = Common::swap32(data_size); + + if (!outfile->WriteArray(entry.exception_lists.data(), entry.exception_lists.size())) + return ConversionResultCode::WriteFailed; + if (!outfile->WriteArray(entry.main_data.data(), entry.main_data.size())) + return ConversionResultCode::WriteFailed; + + *bytes_written += entry.exception_lists.size() + entry.main_data.size(); + + if (entry.reuse_id) + { + std::lock_guard guard(*reusable_groups_mutex); + reusable_groups->emplace(*entry.reuse_id, *group_entry); + } + + if (!PadTo4(outfile, bytes_written)) + return ConversionResultCode::WriteFailed; + + ++group_entry; + } + + return ConversionResultCode::Success; +} + +template <bool RVZ> +ConversionResultCode +WIARVZFileReader<RVZ>::RunCallback(size_t groups_written, u64 bytes_read, u64 bytes_written, + u32 total_groups, u64 iso_size, CompressCB callback, void* arg) +{ + int ratio = 0; + if (bytes_read != 0) + ratio = static_cast<int>(100 * bytes_written / bytes_read); + + const std::string text = + StringFromFormat(Common::GetStringT("%i of %i blocks. Compression ratio %i%%").c_str(), + groups_written, total_groups, ratio); + + const float completion = static_cast<float>(bytes_read) / iso_size; + + return callback(text, completion, arg) ? ConversionResultCode::Success : + ConversionResultCode::Canceled; +} + +template <bool RVZ> +bool WIARVZFileReader<RVZ>::WriteHeader(File::IOFile* file, const u8* data, size_t size, + u64 upper_bound, u64* bytes_written, u64* offset_out) +{ + // The first part of the check is to prevent this from running more than once. If *bytes_written + // is past the upper bound, we are already at the end of the file, so we don't need to do anything + if (*bytes_written <= upper_bound && *bytes_written + size > upper_bound) + { + WARN_LOG(DISCIO, "Headers did not fit in the allocated space. Writing to end of file instead"); + if (!file->Seek(0, SEEK_END)) + return false; + *bytes_written = file->Tell(); + } + + *offset_out = *bytes_written; + if (!file->WriteArray(data, size)) + return false; + *bytes_written += size; + return PadTo4(file, bytes_written); +} + +template <bool RVZ> +ConversionResultCode +WIARVZFileReader<RVZ>::Convert(BlobReader* infile, const VolumeDisc* infile_volume, + File::IOFile* outfile, WIARVZCompressionType compression_type, + int compression_level, int chunk_size, CompressCB callback, + void* arg) +{ + ASSERT(infile->IsDataSizeAccurate()); + ASSERT(chunk_size > 0); + + const u64 iso_size = infile->GetDataSize(); + const u64 chunks_per_wii_group = std::max<u64>(1, VolumeWii::GROUP_TOTAL_SIZE / chunk_size); + const u64 exception_lists_per_chunk = std::max<u64>(1, chunk_size / VolumeWii::GROUP_TOTAL_SIZE); + const bool compressed_exception_lists = compression_type > WIARVZCompressionType::Purge; + + u64 bytes_read = 0; + u64 bytes_written = 0; + size_t groups_processed = 0; + + WIAHeader1 header_1{}; + WIAHeader2 header_2{}; + + std::vector<PartitionEntry> partition_entries; + std::vector<RawDataEntry> raw_data_entries; + std::vector<GroupEntry> group_entries; + + u32 total_groups; + std::vector<DataEntry> data_entries; + + const FileSystem* non_partition_file_system = + infile_volume ? infile_volume->GetFileSystem(PARTITION_NONE) : nullptr; + std::vector<const FileSystem*> partition_file_systems; + + const ConversionResultCode set_up_data_entries_result = SetUpDataEntriesForWriting( + infile_volume, chunk_size, iso_size, &total_groups, &partition_entries, &raw_data_entries, + &data_entries, &partition_file_systems); + if (set_up_data_entries_result != ConversionResultCode::Success) + return set_up_data_entries_result; + + group_entries.resize(total_groups); + + const size_t partition_entries_size = partition_entries.size() * sizeof(PartitionEntry); + const size_t raw_data_entries_size = raw_data_entries.size() * sizeof(RawDataEntry); + const size_t group_entries_size = group_entries.size() * sizeof(GroupEntry); + + // Conservative estimate for how much space will be taken up by headers. + // The compression methods None and Purge have very predictable overhead, + // and the other methods are able to compress group entries well + const u64 headers_size_upper_bound = [&] { + u64 upper_bound = sizeof(WIAHeader1) + sizeof(WIAHeader2) + partition_entries_size + + raw_data_entries_size + 0x100; + + // RVZ's added data in GroupEntry usually compresses well + if (RVZ && compression_type > WIARVZCompressionType::Purge) + upper_bound += group_entries_size / 2; + else + upper_bound += group_entries_size; + + return Common::AlignUp(upper_bound, VolumeWii::BLOCK_TOTAL_SIZE); + }(); + + std::vector<u8> buffer; + + buffer.resize(headers_size_upper_bound); + outfile->WriteBytes(buffer.data(), buffer.size()); + bytes_written = headers_size_upper_bound; + + if (!infile->Read(0, header_2.disc_header.size(), header_2.disc_header.data())) + return ConversionResultCode::ReadFailed; + // We intentially do not increment bytes_read here, since these bytes will be read again + + std::map<ReuseID, GroupEntry> reusable_groups; + std::mutex reusable_groups_mutex; + + const auto set_up_compress_thread_state = [&](CompressThreadState* state) { + SetUpCompressor(&state->compressor, compression_type, compression_level, nullptr); + return ConversionResultCode::Success; + }; + + const auto process_and_compress = [&](CompressThreadState* state, CompressParameters parameters) { + const DataEntry& data_entry = *parameters.data_entry; + const FileSystem* file_system = data_entry.is_partition ? + partition_file_systems[data_entry.index] : + non_partition_file_system; + + const bool compression = compression_type != WIARVZCompressionType::None; + + return ProcessAndCompress(state, std::move(parameters), partition_entries, data_entries, + file_system, &reusable_groups, &reusable_groups_mutex, + chunks_per_wii_group, exception_lists_per_chunk, + compressed_exception_lists, compression); + }; + + const auto output = [&](OutputParameters parameters) { + const ConversionResultCode result = + Output(¶meters.entries, outfile, &reusable_groups, &reusable_groups_mutex, + &group_entries[parameters.group_index], &bytes_written); + + if (result != ConversionResultCode::Success) + return result; + + return RunCallback(parameters.group_index + parameters.entries.size(), parameters.bytes_read, + bytes_written, total_groups, iso_size, callback, arg); + }; + + MultithreadedCompressor<CompressThreadState, CompressParameters, OutputParameters> mt_compressor( + set_up_compress_thread_state, process_and_compress, output); + + for (const DataEntry& data_entry : data_entries) + { + u32 first_group; + u32 last_group; + + u64 data_offset; + u64 data_size; + + u64 data_offset_in_partition; + + if (data_entry.is_partition) + { + const PartitionEntry& partition_entry = partition_entries[data_entry.index]; + const PartitionDataEntry& partition_data_entry = + partition_entry.data_entries[data_entry.partition_data_index]; + + first_group = Common::swap32(partition_data_entry.group_index); + last_group = first_group + Common::swap32(partition_data_entry.number_of_groups); + + const u32 first_sector = Common::swap32(partition_data_entry.first_sector); + data_offset = first_sector * VolumeWii::BLOCK_TOTAL_SIZE; + data_size = + Common::swap32(partition_data_entry.number_of_sectors) * VolumeWii::BLOCK_TOTAL_SIZE; + + const u32 block_in_partition = + first_sector - Common::swap32(partition_entry.data_entries[0].first_sector); + data_offset_in_partition = block_in_partition * VolumeWii::BLOCK_DATA_SIZE; + } + else + { + const RawDataEntry& raw_data_entry = raw_data_entries[data_entry.index]; + + first_group = Common::swap32(raw_data_entry.group_index); + last_group = first_group + Common::swap32(raw_data_entry.number_of_groups); + + data_offset = Common::swap64(raw_data_entry.data_offset); + data_size = Common::swap64(raw_data_entry.data_size); + + const u64 skipped_data = data_offset % VolumeWii::BLOCK_TOTAL_SIZE; + data_offset -= skipped_data; + data_size += skipped_data; + + data_offset_in_partition = data_offset; + } + + ASSERT(groups_processed == first_group); + ASSERT(bytes_read == data_offset); + + while (groups_processed < last_group) + { + const ConversionResultCode status = mt_compressor.GetStatus(); + if (status != ConversionResultCode::Success) + return status; + + u64 bytes_to_read = chunk_size; + if (data_entry.is_partition) + bytes_to_read = std::max<u64>(bytes_to_read, VolumeWii::GROUP_TOTAL_SIZE); + bytes_to_read = std::min<u64>(bytes_to_read, data_offset + data_size - bytes_read); + + buffer.resize(bytes_to_read); + if (!infile->Read(bytes_read, bytes_to_read, buffer.data())) + return ConversionResultCode::ReadFailed; + bytes_read += bytes_to_read; + + mt_compressor.CompressAndWrite(CompressParameters{ + buffer, &data_entry, data_offset_in_partition, bytes_read, groups_processed}); + + data_offset += bytes_to_read; + data_size -= bytes_to_read; + + if (data_entry.is_partition) + { + data_offset_in_partition += + bytes_to_read / VolumeWii::BLOCK_TOTAL_SIZE * VolumeWii::BLOCK_DATA_SIZE; + } + else + { + data_offset_in_partition += bytes_to_read; + } + + groups_processed += Common::AlignUp(bytes_to_read, chunk_size) / chunk_size; + } + + ASSERT(data_size == 0); + } + + ASSERT(groups_processed == total_groups); + ASSERT(bytes_read == iso_size); + + mt_compressor.Shutdown(); + + const ConversionResultCode status = mt_compressor.GetStatus(); + if (status != ConversionResultCode::Success) + return status; + + std::unique_ptr<Compressor> compressor; + SetUpCompressor(&compressor, compression_type, compression_level, &header_2); + + const std::optional<std::vector<u8>> compressed_raw_data_entries = Compress( + compressor.get(), reinterpret_cast<u8*>(raw_data_entries.data()), raw_data_entries_size); + if (!compressed_raw_data_entries) + return ConversionResultCode::InternalError; + + const std::optional<std::vector<u8>> compressed_group_entries = + Compress(compressor.get(), reinterpret_cast<u8*>(group_entries.data()), group_entries_size); + if (!compressed_group_entries) + return ConversionResultCode::InternalError; + + bytes_written = sizeof(WIAHeader1) + sizeof(WIAHeader2); + if (!outfile->Seek(sizeof(WIAHeader1) + sizeof(WIAHeader2), SEEK_SET)) + return ConversionResultCode::WriteFailed; + + u64 partition_entries_offset; + if (!WriteHeader(outfile, reinterpret_cast<u8*>(partition_entries.data()), partition_entries_size, + headers_size_upper_bound, &bytes_written, &partition_entries_offset)) + { + return ConversionResultCode::WriteFailed; + } + + u64 raw_data_entries_offset; + if (!WriteHeader(outfile, compressed_raw_data_entries->data(), + compressed_raw_data_entries->size(), headers_size_upper_bound, &bytes_written, + &raw_data_entries_offset)) + { + return ConversionResultCode::WriteFailed; + } + + u64 group_entries_offset; + if (!WriteHeader(outfile, compressed_group_entries->data(), compressed_group_entries->size(), + headers_size_upper_bound, &bytes_written, &group_entries_offset)) + { + return ConversionResultCode::WriteFailed; + } + + u32 disc_type = 0; + if (infile_volume) + { + if (infile_volume->GetVolumeType() == Platform::GameCubeDisc) + disc_type = 1; + else if (infile_volume->GetVolumeType() == Platform::WiiDisc) + disc_type = 2; + } + + header_2.disc_type = Common::swap32(disc_type); + header_2.compression_type = Common::swap32(static_cast<u32>(compression_type)); + header_2.compression_level = Common::swap32(static_cast<u32>(compression_level)); + header_2.chunk_size = Common::swap32(static_cast<u32>(chunk_size)); + + header_2.number_of_partition_entries = Common::swap32(static_cast<u32>(partition_entries.size())); + header_2.partition_entry_size = Common::swap32(sizeof(PartitionEntry)); + header_2.partition_entries_offset = Common::swap64(partition_entries_offset); + + if (partition_entries.data() == nullptr) + partition_entries.reserve(1); // Avoid a crash in mbedtls_sha1_ret + mbedtls_sha1_ret(reinterpret_cast<const u8*>(partition_entries.data()), partition_entries_size, + header_2.partition_entries_hash.data()); + + header_2.number_of_raw_data_entries = Common::swap32(static_cast<u32>(raw_data_entries.size())); + header_2.raw_data_entries_offset = Common::swap64(raw_data_entries_offset); + header_2.raw_data_entries_size = + Common::swap32(static_cast<u32>(compressed_raw_data_entries->size())); + + header_2.number_of_group_entries = Common::swap32(static_cast<u32>(group_entries.size())); + header_2.group_entries_offset = Common::swap64(group_entries_offset); + header_2.group_entries_size = Common::swap32(static_cast<u32>(compressed_group_entries->size())); + + header_1.magic = RVZ ? RVZ_MAGIC : WIA_MAGIC; + header_1.version = Common::swap32(RVZ ? RVZ_VERSION : WIA_VERSION); + header_1.version_compatible = + Common::swap32(RVZ ? RVZ_VERSION_WRITE_COMPATIBLE : WIA_VERSION_WRITE_COMPATIBLE); + header_1.header_2_size = Common::swap32(sizeof(WIAHeader2)); + mbedtls_sha1_ret(reinterpret_cast<const u8*>(&header_2), sizeof(header_2), + header_1.header_2_hash.data()); + header_1.iso_file_size = Common::swap64(infile->GetDataSize()); + header_1.wia_file_size = Common::swap64(outfile->GetSize()); + mbedtls_sha1_ret(reinterpret_cast<const u8*>(&header_1), offsetof(WIAHeader1, header_1_hash), + header_1.header_1_hash.data()); + + if (!outfile->Seek(0, SEEK_SET)) + return ConversionResultCode::WriteFailed; + + if (!outfile->WriteArray(&header_1, 1)) + return ConversionResultCode::WriteFailed; + if (!outfile->WriteArray(&header_2, 1)) + return ConversionResultCode::WriteFailed; + + return ConversionResultCode::Success; +} + +bool ConvertToWIAOrRVZ(BlobReader* infile, const std::string& infile_path, + const std::string& outfile_path, bool rvz, + WIARVZCompressionType compression_type, int compression_level, + int chunk_size, CompressCB callback, void* arg) +{ + File::IOFile outfile(outfile_path, "wb"); + if (!outfile) + { + PanicAlertT("Failed to open the output file \"%s\".\n" + "Check that you have permissions to write the target folder and that the media can " + "be written.", + outfile_path.c_str()); + return false; + } + + std::unique_ptr<VolumeDisc> infile_volume = CreateDisc(infile_path); + + const auto convert = rvz ? RVZFileReader::Convert : WIAFileReader::Convert; + const ConversionResultCode result = + convert(infile, infile_volume.get(), &outfile, compression_type, compression_level, + chunk_size, callback, arg); + + if (result == ConversionResultCode::ReadFailed) + PanicAlertT("Failed to read from the input file \"%s\".", infile_path.c_str()); + + if (result == ConversionResultCode::WriteFailed) + { + PanicAlertT("Failed to write the output file \"%s\".\n" + "Check that you have enough space available on the target drive.", + outfile_path.c_str()); + } + + if (result != ConversionResultCode::Success) + { + // Remove the incomplete output file + outfile.Close(); + File::Delete(outfile_path); + } + + return result == ConversionResultCode::Success; +} + +template class WIARVZFileReader<false>; +template class WIARVZFileReader<true>; + +} // namespace DiscIO diff --git a/Source/Core/DiscIO/WIABlob.h b/Source/Core/DiscIO/WIABlob.h new file mode 100644 index 0000000000..789a3e7dd5 --- /dev/null +++ b/Source/Core/DiscIO/WIABlob.h @@ -0,0 +1,390 @@ +// Copyright 2018 Dolphin Emulator Project +// Licensed under GPLv2+ +// Refer to the license.txt file included. + +#pragma once + +#include <array> +#include <limits> +#include <map> +#include <memory> +#include <mutex> +#include <type_traits> +#include <utility> + +#include "Common/CommonTypes.h" +#include "Common/File.h" +#include "Common/Swap.h" +#include "DiscIO/Blob.h" +#include "DiscIO/MultithreadedCompressor.h" +#include "DiscIO/WIACompression.h" +#include "DiscIO/WiiEncryptionCache.h" + +namespace DiscIO +{ +class FileSystem; +class VolumeDisc; + +enum class WIARVZCompressionType : u32 +{ + None = 0, + Purge = 1, + Bzip2 = 2, + LZMA = 3, + LZMA2 = 4, + Zstd = 5, +}; + +std::pair<int, int> GetAllowedCompressionLevels(WIARVZCompressionType compression_type); + +constexpr u32 WIA_MAGIC = 0x01414957; // "WIA\x1" (byteswapped to little endian) +constexpr u32 RVZ_MAGIC = 0x015A5652; // "RVZ\x1" (byteswapped to little endian) + +template <bool RVZ> +class WIARVZFileReader : public BlobReader +{ +public: + ~WIARVZFileReader(); + + static std::unique_ptr<WIARVZFileReader> Create(File::IOFile file, const std::string& path); + + BlobType GetBlobType() const override; + + u64 GetRawSize() const override { return Common::swap64(m_header_1.wia_file_size); } + u64 GetDataSize() const override { return Common::swap64(m_header_1.iso_file_size); } + bool IsDataSizeAccurate() const override { return true; } + + u64 GetBlockSize() const override { return Common::swap32(m_header_2.chunk_size); } + bool HasFastRandomAccessInBlock() const override { return false; } + + bool Read(u64 offset, u64 size, u8* out_ptr) override; + bool SupportsReadWiiDecrypted() const override; + bool ReadWiiDecrypted(u64 offset, u64 size, u8* out_ptr, u64 partition_data_offset) override; + + static ConversionResultCode Convert(BlobReader* infile, const VolumeDisc* infile_volume, + File::IOFile* outfile, WIARVZCompressionType compression_type, + int compression_level, int chunk_size, CompressCB callback, + void* arg); + +private: + using SHA1 = std::array<u8, 20>; + using WiiKey = std::array<u8, 16>; + + // See docs/WIA.md for details about the format + +#pragma pack(push, 1) + struct WIAHeader1 + { + u32 magic; + u32 version; + u32 version_compatible; + u32 header_2_size; + SHA1 header_2_hash; + u64 iso_file_size; + u64 wia_file_size; + SHA1 header_1_hash; + }; + static_assert(sizeof(WIAHeader1) == 0x48, "Wrong size for WIA header 1"); + + struct WIAHeader2 + { + u32 disc_type; + u32 compression_type; + u32 compression_level; // Informative only + u32 chunk_size; + + std::array<u8, 0x80> disc_header; + + u32 number_of_partition_entries; + u32 partition_entry_size; + u64 partition_entries_offset; + SHA1 partition_entries_hash; + + u32 number_of_raw_data_entries; + u64 raw_data_entries_offset; + u32 raw_data_entries_size; + + u32 number_of_group_entries; + u64 group_entries_offset; + u32 group_entries_size; + + u8 compressor_data_size; + u8 compressor_data[7]; + }; + static_assert(sizeof(WIAHeader2) == 0xdc, "Wrong size for WIA header 2"); + + struct PartitionDataEntry + { + u32 first_sector; + u32 number_of_sectors; + u32 group_index; + u32 number_of_groups; + }; + static_assert(sizeof(PartitionDataEntry) == 0x10, "Wrong size for WIA partition data entry"); + + struct PartitionEntry + { + WiiKey partition_key; + std::array<PartitionDataEntry, 2> data_entries; + }; + static_assert(sizeof(PartitionEntry) == 0x30, "Wrong size for WIA partition entry"); + + struct RawDataEntry + { + u64 data_offset; + u64 data_size; + u32 group_index; + u32 number_of_groups; + }; + static_assert(sizeof(RawDataEntry) == 0x18, "Wrong size for WIA raw data entry"); + + struct WIAGroupEntry + { + u32 data_offset; // >> 2 + u32 data_size; + }; + static_assert(sizeof(WIAGroupEntry) == 0x08, "Wrong size for WIA group entry"); + + struct RVZGroupEntry + { + u32 data_offset; // >> 2 + u32 data_size; + u32 rvz_packed_size; + }; + static_assert(sizeof(RVZGroupEntry) == 0x0c, "Wrong size for RVZ group entry"); + + using GroupEntry = std::conditional_t<RVZ, RVZGroupEntry, WIAGroupEntry>; + + struct HashExceptionEntry + { + u16 offset; + SHA1 hash; + }; + static_assert(sizeof(HashExceptionEntry) == 0x16, "Wrong size for WIA hash exception entry"); +#pragma pack(pop) + + struct DataEntry + { + u32 index; + bool is_partition; + u8 partition_data_index; + + DataEntry(size_t index_) : index(static_cast<u32>(index_)), is_partition(false) {} + DataEntry(size_t index_, size_t partition_data_index_) + : index(static_cast<u32>(index_)), is_partition(true), + partition_data_index(static_cast<u8>(partition_data_index_)) + { + } + }; + + class Chunk + { + public: + Chunk(); + Chunk(File::IOFile* file, u64 offset_in_file, u64 compressed_size, u64 decompressed_size, + u32 exception_lists, bool compressed_exception_lists, u32 rvz_packed_size, + u64 data_offset, std::unique_ptr<Decompressor> decompressor); + + bool Read(u64 offset, u64 size, u8* out_ptr); + + // This can only be called once at least one byte of data has been read + void GetHashExceptions(std::vector<HashExceptionEntry>* exception_list, + u64 exception_list_index, u16 additional_offset) const; + + template <typename T> + bool ReadAll(std::vector<T>* vector) + { + return Read(0, vector->size() * sizeof(T), reinterpret_cast<u8*>(vector->data())); + } + + private: + bool Decompress(); + bool HandleExceptions(const u8* data, size_t bytes_allocated, size_t bytes_written, + size_t* bytes_used, bool align); + + DecompressionBuffer m_in; + DecompressionBuffer m_out; + size_t m_in_bytes_read = 0; + + std::unique_ptr<Decompressor> m_decompressor = nullptr; + File::IOFile* m_file = nullptr; + u64 m_offset_in_file = 0; + + size_t m_out_bytes_allocated_for_exceptions = 0; + size_t m_out_bytes_used_for_exceptions = 0; + size_t m_in_bytes_used_for_exceptions = 0; + u32 m_exception_lists = 0; + bool m_compressed_exception_lists = false; + u32 m_rvz_packed_size = 0; + u64 m_data_offset = 0; + }; + + explicit WIARVZFileReader(File::IOFile file, const std::string& path); + bool Initialize(const std::string& path); + bool HasDataOverlap() const; + + bool ReadFromGroups(u64* offset, u64* size, u8** out_ptr, u64 chunk_size, u32 sector_size, + u64 data_offset, u64 data_size, u32 group_index, u32 number_of_groups, + u32 exception_lists); + Chunk& ReadCompressedData(u64 offset_in_file, u64 compressed_size, u64 decompressed_size, + WIARVZCompressionType compression_type, u32 exception_lists = 0, + u32 rvz_packed_size = 0, u64 data_offset = 0); + + static bool ApplyHashExceptions(const std::vector<HashExceptionEntry>& exception_list, + VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP]); + + static std::string VersionToString(u32 version); + + struct ReuseID + { + bool operator==(const ReuseID& other) const + { + return std::tie(partition_key, data_size, encrypted, value) == + std::tie(other.partition_key, other.data_size, other.encrypted, other.value); + } + bool operator<(const ReuseID& other) const + { + return std::tie(partition_key, data_size, encrypted, value) < + std::tie(other.partition_key, other.data_size, other.encrypted, other.value); + } + bool operator>(const ReuseID& other) const + { + return std::tie(partition_key, data_size, encrypted, value) > + std::tie(other.partition_key, other.data_size, other.encrypted, other.value); + } + bool operator!=(const ReuseID& other) const { return !operator==(other); } + bool operator>=(const ReuseID& other) const { return !operator<(other); } + bool operator<=(const ReuseID& other) const { return !operator>(other); } + + const WiiKey* partition_key; + u64 data_size; + bool encrypted; + u8 value; + }; + + struct CompressThreadState + { + using WiiBlockData = std::array<u8, VolumeWii::BLOCK_DATA_SIZE>; + + std::unique_ptr<Compressor> compressor; + + std::vector<WiiBlockData> decryption_buffer = + std::vector<WiiBlockData>(VolumeWii::BLOCKS_PER_GROUP); + + std::vector<VolumeWii::HashBlock> hash_buffer = + std::vector<VolumeWii::HashBlock>(VolumeWii::BLOCKS_PER_GROUP); + }; + + struct CompressParameters + { + std::vector<u8> data; + const DataEntry* data_entry; + u64 data_offset; + u64 bytes_read; + size_t group_index; + }; + + struct WIAOutputParametersEntry + { + std::vector<u8> exception_lists; + std::vector<u8> main_data; + std::optional<ReuseID> reuse_id; + std::optional<GroupEntry> reused_group; + }; + + struct RVZOutputParametersEntry + { + std::vector<u8> exception_lists; + std::vector<u8> main_data; + std::optional<ReuseID> reuse_id; + std::optional<GroupEntry> reused_group; + size_t rvz_packed_size = 0; + bool compressed = false; + }; + + using OutputParametersEntry = + std::conditional_t<RVZ, RVZOutputParametersEntry, WIAOutputParametersEntry>; + + struct OutputParameters + { + std::vector<OutputParametersEntry> entries; + u64 bytes_read; + size_t group_index; + }; + + static bool PadTo4(File::IOFile* file, u64* bytes_written); + static void AddRawDataEntry(u64 offset, u64 size, int chunk_size, u32* total_groups, + std::vector<RawDataEntry>* raw_data_entries, + std::vector<DataEntry>* data_entries); + static PartitionDataEntry + CreatePartitionDataEntry(u64 offset, u64 size, u32 index, int chunk_size, u32* total_groups, + const std::vector<PartitionEntry>& partition_entries, + std::vector<DataEntry>* data_entries); + static ConversionResultCode SetUpDataEntriesForWriting( + const VolumeDisc* volume, int chunk_size, u64 iso_size, u32* total_groups, + std::vector<PartitionEntry>* partition_entries, std::vector<RawDataEntry>* raw_data_entries, + std::vector<DataEntry>* data_entries, std::vector<const FileSystem*>* partition_file_systems); + static std::optional<std::vector<u8>> Compress(Compressor* compressor, const u8* data, + size_t size); + static bool WriteHeader(File::IOFile* file, const u8* data, size_t size, u64 upper_bound, + u64* bytes_written, u64* offset_out); + + static void SetUpCompressor(std::unique_ptr<Compressor>* compressor, + WIARVZCompressionType compression_type, int compression_level, + WIAHeader2* header_2); + static bool TryReuse(std::map<ReuseID, GroupEntry>* reusable_groups, + std::mutex* reusable_groups_mutex, OutputParametersEntry* entry); + static ConversionResult<OutputParameters> + ProcessAndCompress(CompressThreadState* state, CompressParameters parameters, + const std::vector<PartitionEntry>& partition_entries, + const std::vector<DataEntry>& data_entries, const FileSystem* file_system, + std::map<ReuseID, GroupEntry>* reusable_groups, + std::mutex* reusable_groups_mutex, u64 chunks_per_wii_group, + u64 exception_lists_per_chunk, bool compressed_exception_lists, + bool compression); + static ConversionResultCode Output(std::vector<OutputParametersEntry>* entries, + File::IOFile* outfile, + std::map<ReuseID, GroupEntry>* reusable_groups, + std::mutex* reusable_groups_mutex, GroupEntry* group_entry, + u64* bytes_written); + static ConversionResultCode RunCallback(size_t groups_written, u64 bytes_read, u64 bytes_written, + u32 total_groups, u64 iso_size, CompressCB callback, + void* arg); + + bool m_valid; + WIARVZCompressionType m_compression_type; + + File::IOFile m_file; + Chunk m_cached_chunk; + u64 m_cached_chunk_offset = std::numeric_limits<u64>::max(); + WiiEncryptionCache m_encryption_cache; + + std::vector<HashExceptionEntry> m_exception_list; + bool m_write_to_exception_list = false; + u64 m_exception_list_last_group_index; + + WIAHeader1 m_header_1; + WIAHeader2 m_header_2; + std::vector<PartitionEntry> m_partition_entries; + std::vector<RawDataEntry> m_raw_data_entries; + std::vector<GroupEntry> m_group_entries; + + std::map<u64, DataEntry> m_data_entries; + + // Perhaps we could set WIA_VERSION_WRITE_COMPATIBLE to 0.9, but WIA version 0.9 was never in + // any official release of wit, and interim versions (either source or binaries) are hard to find. + // Since we've been unable to check if we're write compatible with 0.9, we set it 1.0 to be safe. + + static constexpr u32 WIA_VERSION = 0x01000000; + static constexpr u32 WIA_VERSION_WRITE_COMPATIBLE = 0x01000000; + static constexpr u32 WIA_VERSION_READ_COMPATIBLE = 0x00080000; + + static constexpr u32 RVZ_VERSION = 0x01000000; + static constexpr u32 RVZ_VERSION_WRITE_COMPATIBLE = 0x00030000; + static constexpr u32 RVZ_VERSION_READ_COMPATIBLE = 0x00030000; +}; + +using WIAFileReader = WIARVZFileReader<false>; +using RVZFileReader = WIARVZFileReader<true>; + +} // namespace DiscIO diff --git a/Source/Core/DiscIO/WIACompression.cpp b/Source/Core/DiscIO/WIACompression.cpp new file mode 100644 index 0000000000..d6d3acae2b --- /dev/null +++ b/Source/Core/DiscIO/WIACompression.cpp @@ -0,0 +1,810 @@ +// Copyright 2020 Dolphin Emulator Project +// Licensed under GPLv2+ +// Refer to the license.txt file included. + +#include "DiscIO/WIACompression.h" + +#include <algorithm> +#include <cstddef> +#include <cstring> +#include <memory> +#include <optional> +#include <vector> + +#include <bzlib.h> +#include <lzma.h> +#include <mbedtls/sha1.h> +#include <zstd.h> + +#include "Common/Assert.h" +#include "Common/CommonTypes.h" +#include "Common/Swap.h" +#include "DiscIO/LaggedFibonacciGenerator.h" + +namespace DiscIO +{ +static u32 LZMA2DictionarySize(u8 p) +{ + return (static_cast<u32>(2) | (p & 1)) << (p / 2 + 11); +} + +Decompressor::~Decompressor() = default; + +bool NoneDecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) +{ + const size_t length = + std::min(in.bytes_written - *in_bytes_read, out->data.size() - out->bytes_written); + + std::memcpy(out->data.data() + out->bytes_written, in.data.data() + *in_bytes_read, length); + + *in_bytes_read += length; + out->bytes_written += length; + + m_done = in.data.size() == *in_bytes_read; + return true; +} + +PurgeDecompressor::PurgeDecompressor(u64 decompressed_size) : m_decompressed_size(decompressed_size) +{ + mbedtls_sha1_init(&m_sha1_context); +} + +bool PurgeDecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) +{ + if (!m_started) + { + mbedtls_sha1_starts_ret(&m_sha1_context); + + // Include the exception lists in the SHA-1 calculation (but not in the compression...) + mbedtls_sha1_update_ret(&m_sha1_context, in.data.data(), *in_bytes_read); + + m_started = true; + } + + while (!m_done && in.bytes_written != *in_bytes_read && + (m_segment_bytes_written < sizeof(m_segment) || out->data.size() != out->bytes_written)) + { + if (m_segment_bytes_written == 0 && *in_bytes_read == in.data.size() - sizeof(SHA1)) + { + const size_t zeroes_to_write = std::min<size_t>(m_decompressed_size - m_out_bytes_written, + out->data.size() - out->bytes_written); + + std::memset(out->data.data() + out->bytes_written, 0, zeroes_to_write); + + out->bytes_written += zeroes_to_write; + m_out_bytes_written += zeroes_to_write; + + if (m_out_bytes_written == m_decompressed_size && in.bytes_written == in.data.size()) + { + SHA1 actual_hash; + mbedtls_sha1_finish_ret(&m_sha1_context, actual_hash.data()); + + SHA1 expected_hash; + std::memcpy(expected_hash.data(), in.data.data() + *in_bytes_read, expected_hash.size()); + + *in_bytes_read += expected_hash.size(); + m_done = true; + + if (actual_hash != expected_hash) + return false; + } + + return true; + } + + if (m_segment_bytes_written < sizeof(m_segment)) + { + const size_t bytes_to_copy = + std::min(in.bytes_written - *in_bytes_read, sizeof(m_segment) - m_segment_bytes_written); + + std::memcpy(reinterpret_cast<u8*>(&m_segment) + m_segment_bytes_written, + in.data.data() + *in_bytes_read, bytes_to_copy); + mbedtls_sha1_update_ret(&m_sha1_context, in.data.data() + *in_bytes_read, bytes_to_copy); + + *in_bytes_read += bytes_to_copy; + m_bytes_read += bytes_to_copy; + m_segment_bytes_written += bytes_to_copy; + } + + if (m_segment_bytes_written < sizeof(m_segment)) + return true; + + const size_t offset = Common::swap32(m_segment.offset); + const size_t size = Common::swap32(m_segment.size); + + if (m_out_bytes_written < offset) + { + const size_t zeroes_to_write = + std::min(offset - m_out_bytes_written, out->data.size() - out->bytes_written); + + std::memset(out->data.data() + out->bytes_written, 0, zeroes_to_write); + + out->bytes_written += zeroes_to_write; + m_out_bytes_written += zeroes_to_write; + } + + if (m_out_bytes_written >= offset && m_out_bytes_written < offset + size) + { + const size_t bytes_to_copy = std::min( + std::min(offset + size - m_out_bytes_written, out->data.size() - out->bytes_written), + in.bytes_written - *in_bytes_read); + + std::memcpy(out->data.data() + out->bytes_written, in.data.data() + *in_bytes_read, + bytes_to_copy); + mbedtls_sha1_update_ret(&m_sha1_context, in.data.data() + *in_bytes_read, bytes_to_copy); + + *in_bytes_read += bytes_to_copy; + m_bytes_read += bytes_to_copy; + out->bytes_written += bytes_to_copy; + m_out_bytes_written += bytes_to_copy; + } + + if (m_out_bytes_written >= offset + size) + m_segment_bytes_written = 0; + } + + return true; +} + +Bzip2Decompressor::~Bzip2Decompressor() +{ + if (m_started) + BZ2_bzDecompressEnd(&m_stream); +} + +bool Bzip2Decompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) +{ + if (!m_started) + { + if (BZ2_bzDecompressInit(&m_stream, 0, 0) != BZ_OK) + return false; + + m_started = true; + } + + constexpr auto clamped_cast = [](size_t x) { + return static_cast<unsigned int>( + std::min<size_t>(std::numeric_limits<unsigned int>().max(), x)); + }; + + char* const in_ptr = reinterpret_cast<char*>(const_cast<u8*>(in.data.data() + *in_bytes_read)); + m_stream.next_in = in_ptr; + m_stream.avail_in = clamped_cast(in.bytes_written - *in_bytes_read); + + char* const out_ptr = reinterpret_cast<char*>(out->data.data() + out->bytes_written); + m_stream.next_out = out_ptr; + m_stream.avail_out = clamped_cast(out->data.size() - out->bytes_written); + + const int result = BZ2_bzDecompress(&m_stream); + + *in_bytes_read += m_stream.next_in - in_ptr; + out->bytes_written += m_stream.next_out - out_ptr; + + m_done = result == BZ_STREAM_END; + return result == BZ_OK || result == BZ_STREAM_END; +} + +LZMADecompressor::LZMADecompressor(bool lzma2, const u8* filter_options, size_t filter_options_size) +{ + m_options.preset_dict = nullptr; + + if (!lzma2 && filter_options_size == 5) + { + // The dictionary size is stored as a 32-bit little endian unsigned integer + static_assert(sizeof(m_options.dict_size) == sizeof(u32)); + std::memcpy(&m_options.dict_size, filter_options + 1, sizeof(u32)); + + const u8 d = filter_options[0]; + if (d >= (9 * 5 * 5)) + { + m_error_occurred = true; + } + else + { + m_options.lc = d % 9; + const u8 e = d / 9; + m_options.pb = e / 5; + m_options.lp = e % 5; + } + } + else if (lzma2 && filter_options_size == 1) + { + const u8 d = filter_options[0]; + if (d > 40) + m_error_occurred = true; + else + m_options.dict_size = d == 40 ? 0xFFFFFFFF : LZMA2DictionarySize(d); + } + else + { + m_error_occurred = true; + } + + m_filters[0].id = lzma2 ? LZMA_FILTER_LZMA2 : LZMA_FILTER_LZMA1; + m_filters[0].options = &m_options; + m_filters[1].id = LZMA_VLI_UNKNOWN; + m_filters[1].options = nullptr; +} + +LZMADecompressor::~LZMADecompressor() +{ + if (m_started) + lzma_end(&m_stream); +} + +bool LZMADecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) +{ + if (!m_started) + { + if (m_error_occurred || lzma_raw_decoder(&m_stream, m_filters) != LZMA_OK) + return false; + + m_started = true; + } + + const u8* const in_ptr = in.data.data() + *in_bytes_read; + m_stream.next_in = in_ptr; + m_stream.avail_in = in.bytes_written - *in_bytes_read; + + u8* const out_ptr = out->data.data() + out->bytes_written; + m_stream.next_out = out_ptr; + m_stream.avail_out = out->data.size() - out->bytes_written; + + const lzma_ret result = lzma_code(&m_stream, LZMA_RUN); + + *in_bytes_read += m_stream.next_in - in_ptr; + out->bytes_written += m_stream.next_out - out_ptr; + + m_done = result == LZMA_STREAM_END; + return result == LZMA_OK || result == LZMA_STREAM_END; +} + +ZstdDecompressor::ZstdDecompressor() +{ + m_stream = ZSTD_createDStream(); +} + +ZstdDecompressor::~ZstdDecompressor() +{ + ZSTD_freeDStream(m_stream); +} + +bool ZstdDecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) +{ + if (!m_stream) + return false; + + ZSTD_inBuffer in_buffer{in.data.data(), in.bytes_written, *in_bytes_read}; + ZSTD_outBuffer out_buffer{out->data.data(), out->data.size(), out->bytes_written}; + + const size_t result = ZSTD_decompressStream(m_stream, &out_buffer, &in_buffer); + + *in_bytes_read = in_buffer.pos; + out->bytes_written = out_buffer.pos; + + m_done = result == 0; + return !ZSTD_isError(result); +} + +RVZPackDecompressor::RVZPackDecompressor(std::unique_ptr<Decompressor> decompressor, + DecompressionBuffer decompressed, u64 data_offset, + u32 rvz_packed_size) + : m_decompressor(std::move(decompressor)), m_decompressed(std::move(decompressed)), + m_data_offset(data_offset), m_rvz_packed_size(rvz_packed_size) +{ + m_bytes_read = m_decompressed.bytes_written; +} + +bool RVZPackDecompressor::IncrementBytesRead(size_t x) +{ + m_bytes_read += x; + return m_bytes_read <= m_rvz_packed_size; +} + +std::optional<bool> RVZPackDecompressor::ReadToDecompressed(const DecompressionBuffer& in, + size_t* in_bytes_read, + size_t decompressed_bytes_read, + size_t bytes_to_read) +{ + if (m_decompressed.data.size() < decompressed_bytes_read + bytes_to_read) + m_decompressed.data.resize(decompressed_bytes_read + bytes_to_read); + + if (m_decompressed.bytes_written < decompressed_bytes_read + bytes_to_read) + { + const size_t prev_bytes_written = m_decompressed.bytes_written; + + if (!m_decompressor->Decompress(in, &m_decompressed, in_bytes_read)) + return false; + + if (!IncrementBytesRead(m_decompressed.bytes_written - prev_bytes_written)) + return false; + + if (m_decompressed.bytes_written < decompressed_bytes_read + bytes_to_read) + return true; + } + + return std::nullopt; +} + +bool RVZPackDecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) +{ + while (out->data.size() != out->bytes_written && !Done()) + { + if (m_size == 0) + { + if (m_decompressed.bytes_written == m_decompressed_bytes_read) + { + m_decompressed.data.resize(sizeof(u32)); + m_decompressed.bytes_written = 0; + m_decompressed_bytes_read = 0; + } + + std::optional<bool> result = + ReadToDecompressed(in, in_bytes_read, m_decompressed_bytes_read, sizeof(u32)); + if (result) + return *result; + + m_size = Common::swap32(m_decompressed.data.data() + m_decompressed_bytes_read); + + m_junk = m_size & 0x80000000; + if (m_junk) + { + m_size &= 0x7FFFFFFF; + + constexpr size_t SEED_SIZE = LaggedFibonacciGenerator::SEED_SIZE * sizeof(u32); + constexpr size_t BLOCK_SIZE = 0x8000; + + result = ReadToDecompressed(in, in_bytes_read, m_decompressed_bytes_read + sizeof(u32), + SEED_SIZE); + if (result) + return *result; + + m_lfg.SetSeed(m_decompressed.data.data() + m_decompressed_bytes_read + sizeof(u32)); + m_lfg.Forward(m_data_offset % BLOCK_SIZE); + + m_decompressed_bytes_read += SEED_SIZE; + } + + m_decompressed_bytes_read += sizeof(u32); + } + + size_t bytes_to_write = std::min<size_t>(m_size, out->data.size() - out->bytes_written); + if (m_junk) + { + m_lfg.GetBytes(bytes_to_write, out->data.data() + out->bytes_written); + out->bytes_written += bytes_to_write; + } + else + { + if (m_decompressed.bytes_written != m_decompressed_bytes_read) + { + bytes_to_write = + std::min(bytes_to_write, m_decompressed.bytes_written - m_decompressed_bytes_read); + + std::memcpy(out->data.data() + out->bytes_written, + m_decompressed.data.data() + m_decompressed_bytes_read, bytes_to_write); + + m_decompressed_bytes_read += bytes_to_write; + out->bytes_written += bytes_to_write; + } + else + { + const size_t prev_out_bytes_written = out->bytes_written; + const size_t old_out_size = out->data.size(); + const size_t new_out_size = out->bytes_written + bytes_to_write; + + if (new_out_size < old_out_size) + out->data.resize(new_out_size); + + if (!m_decompressor->Decompress(in, out, in_bytes_read)) + return false; + + out->data.resize(old_out_size); + + bytes_to_write = out->bytes_written - prev_out_bytes_written; + + if (!IncrementBytesRead(bytes_to_write)) + return false; + + if (bytes_to_write == 0) + return true; + } + } + + m_data_offset += bytes_to_write; + m_size -= static_cast<u32>(bytes_to_write); + } + + // If out is full but not all data has been read from in, give the decompressor a chance to read + // from in anyway. This is needed for the case where zstd has read everything except the checksum. + if (out->data.size() == out->bytes_written && in.bytes_written != *in_bytes_read) + { + if (!m_decompressor->Decompress(in, out, in_bytes_read)) + return false; + } + + return true; +} + +bool RVZPackDecompressor::Done() const +{ + return m_size == 0 && m_rvz_packed_size == m_bytes_read && + m_decompressed.bytes_written == m_decompressed_bytes_read && m_decompressor->Done(); +} + +Compressor::~Compressor() = default; + +PurgeCompressor::PurgeCompressor() +{ + mbedtls_sha1_init(&m_sha1_context); +} + +PurgeCompressor::~PurgeCompressor() = default; + +bool PurgeCompressor::Start() +{ + m_buffer.clear(); + m_bytes_written = 0; + + mbedtls_sha1_starts_ret(&m_sha1_context); + + return true; +} + +bool PurgeCompressor::AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) +{ + mbedtls_sha1_update_ret(&m_sha1_context, data, size); + return true; +} + +bool PurgeCompressor::Compress(const u8* data, size_t size) +{ + // We could add support for calling this twice if we're fine with + // making the code more complicated, but there's no need to support it + ASSERT_MSG(DISCIO, m_bytes_written == 0, + "Calling PurgeCompressor::Compress() twice is not supported"); + + m_buffer.resize(size + sizeof(PurgeSegment) + sizeof(SHA1)); + + size_t bytes_read = 0; + + while (true) + { + const auto first_non_zero = + std::find_if(data + bytes_read, data + size, [](u8 x) { return x != 0; }); + + const u32 non_zero_data_start = static_cast<u32>(first_non_zero - data); + if (non_zero_data_start == size) + break; + + size_t non_zero_data_end = non_zero_data_start; + size_t sequence_length = 0; + for (size_t i = non_zero_data_start; i < size; ++i) + { + if (data[i] == 0) + { + ++sequence_length; + } + else + { + sequence_length = 0; + non_zero_data_end = i + 1; + } + + // To avoid wasting space, only count runs of zeroes that are of a certain length + // (unless there is nothing after the run of zeroes, then we might as well always count it) + if (sequence_length > sizeof(PurgeSegment)) + break; + } + + const u32 non_zero_data_length = static_cast<u32>(non_zero_data_end - non_zero_data_start); + + const PurgeSegment segment{Common::swap32(non_zero_data_start), + Common::swap32(non_zero_data_length)}; + std::memcpy(m_buffer.data() + m_bytes_written, &segment, sizeof(segment)); + m_bytes_written += sizeof(segment); + + std::memcpy(m_buffer.data() + m_bytes_written, data + non_zero_data_start, + non_zero_data_length); + m_bytes_written += non_zero_data_length; + + bytes_read = non_zero_data_end; + } + + return true; +} + +bool PurgeCompressor::End() +{ + mbedtls_sha1_update_ret(&m_sha1_context, m_buffer.data(), m_bytes_written); + + mbedtls_sha1_finish_ret(&m_sha1_context, m_buffer.data() + m_bytes_written); + m_bytes_written += sizeof(SHA1); + + ASSERT(m_bytes_written <= m_buffer.size()); + + return true; +} + +const u8* PurgeCompressor::GetData() const +{ + return m_buffer.data(); +} + +size_t PurgeCompressor::GetSize() const +{ + return m_bytes_written; +} + +Bzip2Compressor::Bzip2Compressor(int compression_level) : m_compression_level(compression_level) +{ +} + +Bzip2Compressor::~Bzip2Compressor() +{ + BZ2_bzCompressEnd(&m_stream); +} + +bool Bzip2Compressor::Start() +{ + ASSERT_MSG(DISCIO, m_stream.state == nullptr, + "Called Bzip2Compressor::Start() twice without calling Bzip2Compressor::End()"); + + m_buffer.clear(); + m_stream.next_out = reinterpret_cast<char*>(m_buffer.data()); + + return BZ2_bzCompressInit(&m_stream, m_compression_level, 0, 0) == BZ_OK; +} + +bool Bzip2Compressor::Compress(const u8* data, size_t size) +{ + m_stream.next_in = reinterpret_cast<char*>(const_cast<u8*>(data)); + m_stream.avail_in = static_cast<unsigned int>(size); + + ExpandBuffer(size); + + while (m_stream.avail_in != 0) + { + if (m_stream.avail_out == 0) + ExpandBuffer(0x100); + + if (BZ2_bzCompress(&m_stream, BZ_RUN) != BZ_RUN_OK) + return false; + } + + return true; +} + +bool Bzip2Compressor::End() +{ + bool success = true; + + while (true) + { + if (m_stream.avail_out == 0) + ExpandBuffer(0x100); + + const int result = BZ2_bzCompress(&m_stream, BZ_FINISH); + if (result != BZ_FINISH_OK && result != BZ_STREAM_END) + success = false; + if (result != BZ_FINISH_OK) + break; + } + + if (BZ2_bzCompressEnd(&m_stream) != BZ_OK) + success = false; + + return success; +} + +void Bzip2Compressor::ExpandBuffer(size_t bytes_to_add) +{ + const size_t bytes_written = GetSize(); + m_buffer.resize(m_buffer.size() + bytes_to_add); + m_stream.next_out = reinterpret_cast<char*>(m_buffer.data()) + bytes_written; + m_stream.avail_out = static_cast<unsigned int>(m_buffer.size() - bytes_written); +} + +const u8* Bzip2Compressor::GetData() const +{ + return m_buffer.data(); +} + +size_t Bzip2Compressor::GetSize() const +{ + return static_cast<size_t>(reinterpret_cast<u8*>(m_stream.next_out) - m_buffer.data()); +} + +LZMACompressor::LZMACompressor(bool lzma2, int compression_level, u8 compressor_data_out[7], + u8* compressor_data_size_out) +{ + // lzma_lzma_preset returns false on success for some reason + if (lzma_lzma_preset(&m_options, static_cast<uint32_t>(compression_level))) + { + m_initialization_failed = true; + return; + } + + if (!lzma2) + { + if (compressor_data_size_out) + *compressor_data_size_out = 5; + + if (compressor_data_out) + { + ASSERT(m_options.lc < 9); + ASSERT(m_options.lp < 5); + ASSERT(m_options.pb < 5); + compressor_data_out[0] = + static_cast<u8>((m_options.pb * 5 + m_options.lp) * 9 + m_options.lc); + + // The dictionary size is stored as a 32-bit little endian unsigned integer + static_assert(sizeof(m_options.dict_size) == sizeof(u32)); + std::memcpy(compressor_data_out + 1, &m_options.dict_size, sizeof(u32)); + } + } + else + { + if (compressor_data_size_out) + *compressor_data_size_out = 1; + + if (compressor_data_out) + { + u8 encoded_dict_size = 0; + while (encoded_dict_size < 40 && m_options.dict_size > LZMA2DictionarySize(encoded_dict_size)) + ++encoded_dict_size; + + compressor_data_out[0] = encoded_dict_size; + } + } + + m_filters[0].id = lzma2 ? LZMA_FILTER_LZMA2 : LZMA_FILTER_LZMA1; + m_filters[0].options = &m_options; + m_filters[1].id = LZMA_VLI_UNKNOWN; + m_filters[1].options = nullptr; +} + +LZMACompressor::~LZMACompressor() +{ + lzma_end(&m_stream); +} + +bool LZMACompressor::Start() +{ + if (m_initialization_failed) + return false; + + m_buffer.clear(); + m_stream.next_out = m_buffer.data(); + + return lzma_raw_encoder(&m_stream, m_filters) == LZMA_OK; +} + +bool LZMACompressor::Compress(const u8* data, size_t size) +{ + m_stream.next_in = data; + m_stream.avail_in = size; + + ExpandBuffer(size); + + while (m_stream.avail_in != 0) + { + if (m_stream.avail_out == 0) + ExpandBuffer(0x100); + + if (lzma_code(&m_stream, LZMA_RUN) != LZMA_OK) + return false; + } + + return true; +} + +bool LZMACompressor::End() +{ + while (true) + { + if (m_stream.avail_out == 0) + ExpandBuffer(0x100); + + switch (lzma_code(&m_stream, LZMA_FINISH)) + { + case LZMA_OK: + break; + case LZMA_STREAM_END: + return true; + default: + return false; + } + } +} + +void LZMACompressor::ExpandBuffer(size_t bytes_to_add) +{ + const size_t bytes_written = GetSize(); + m_buffer.resize(m_buffer.size() + bytes_to_add); + m_stream.next_out = m_buffer.data() + bytes_written; + m_stream.avail_out = m_buffer.size() - bytes_written; +} + +const u8* LZMACompressor::GetData() const +{ + return m_buffer.data(); +} + +size_t LZMACompressor::GetSize() const +{ + return static_cast<size_t>(m_stream.next_out - m_buffer.data()); +} + +ZstdCompressor::ZstdCompressor(int compression_level) +{ + m_stream = ZSTD_createCStream(); + + if (ZSTD_isError(ZSTD_CCtx_setParameter(m_stream, ZSTD_c_compressionLevel, compression_level))) + m_stream = nullptr; +} + +ZstdCompressor::~ZstdCompressor() +{ + ZSTD_freeCStream(m_stream); +} + +bool ZstdCompressor::Start() +{ + if (!m_stream) + return false; + + m_buffer.clear(); + m_out_buffer = {}; + + return !ZSTD_isError(ZSTD_CCtx_reset(m_stream, ZSTD_reset_session_only)); +} + +bool ZstdCompressor::Compress(const u8* data, size_t size) +{ + ZSTD_inBuffer in_buffer{data, size, 0}; + + ExpandBuffer(size); + + while (in_buffer.size != in_buffer.pos) + { + if (m_out_buffer.size == m_out_buffer.pos) + ExpandBuffer(0x100); + + if (ZSTD_isError(ZSTD_compressStream(m_stream, &m_out_buffer, &in_buffer))) + return false; + } + + return true; +} + +bool ZstdCompressor::End() +{ + while (true) + { + if (m_out_buffer.size == m_out_buffer.pos) + ExpandBuffer(0x100); + + const size_t result = ZSTD_endStream(m_stream, &m_out_buffer); + if (ZSTD_isError(result)) + return false; + if (result == 0) + return true; + } +} + +void ZstdCompressor::ExpandBuffer(size_t bytes_to_add) +{ + m_buffer.resize(m_buffer.size() + bytes_to_add); + + m_out_buffer.dst = m_buffer.data(); + m_out_buffer.size = m_buffer.size(); +} + +} // namespace DiscIO diff --git a/Source/Core/DiscIO/WIACompression.h b/Source/Core/DiscIO/WIACompression.h new file mode 100644 index 0000000000..37e8cf3dc3 --- /dev/null +++ b/Source/Core/DiscIO/WIACompression.h @@ -0,0 +1,252 @@ +// Copyright 2020 Dolphin Emulator Project +// Licensed under GPLv2+ +// Refer to the license.txt file included. + +#pragma once + +#include <cstddef> +#include <memory> +#include <optional> +#include <vector> + +#include <bzlib.h> +#include <lzma.h> +#include <mbedtls/sha1.h> +#include <zstd.h> + +#include "Common/CommonTypes.h" +#include "DiscIO/LaggedFibonacciGenerator.h" + +namespace DiscIO +{ +struct DecompressionBuffer +{ + std::vector<u8> data; + size_t bytes_written = 0; +}; + +using SHA1 = std::array<u8, 20>; + +struct PurgeSegment +{ + u32 offset; + u32 size; +}; +static_assert(sizeof(PurgeSegment) == 0x08, "Wrong size for WIA purge segment"); + +class Decompressor +{ +public: + virtual ~Decompressor(); + + virtual bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) = 0; + virtual bool Done() const { return m_done; }; + +protected: + bool m_done = false; +}; + +class NoneDecompressor final : public Decompressor +{ +public: + bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) override; +}; + +// This class assumes that more bytes won't be added to in once in.bytes_written == in.data.size() +// and that *in_bytes_read initially will be equal to the size of the exception lists +class PurgeDecompressor final : public Decompressor +{ +public: + PurgeDecompressor(u64 decompressed_size); + bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) override; + +private: + const u64 m_decompressed_size; + + PurgeSegment m_segment = {}; + size_t m_bytes_read = 0; + size_t m_segment_bytes_written = 0; + size_t m_out_bytes_written = 0; + bool m_started = false; + + mbedtls_sha1_context m_sha1_context; +}; + +class Bzip2Decompressor final : public Decompressor +{ +public: + ~Bzip2Decompressor(); + + bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) override; + +private: + bz_stream m_stream = {}; + bool m_started = false; +}; + +class LZMADecompressor final : public Decompressor +{ +public: + LZMADecompressor(bool lzma2, const u8* filter_options, size_t filter_options_size); + ~LZMADecompressor(); + + bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) override; + +private: + lzma_stream m_stream = LZMA_STREAM_INIT; + lzma_options_lzma m_options = {}; + lzma_filter m_filters[2]; + bool m_started = false; + bool m_error_occurred = false; +}; + +class ZstdDecompressor final : public Decompressor +{ +public: + ZstdDecompressor(); + ~ZstdDecompressor(); + + bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) override; + +private: + ZSTD_DStream* m_stream; +}; + +class RVZPackDecompressor final : public Decompressor +{ +public: + RVZPackDecompressor(std::unique_ptr<Decompressor> decompressor, DecompressionBuffer decompressed, + u64 data_offset, u32 rvz_packed_size); + + bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out, + size_t* in_bytes_read) override; + + bool Done() const override; + +private: + bool IncrementBytesRead(size_t x); + std::optional<bool> ReadToDecompressed(const DecompressionBuffer& in, size_t* in_bytes_read, + size_t decompressed_bytes_read, size_t bytes_to_read); + + std::unique_ptr<Decompressor> m_decompressor; + DecompressionBuffer m_decompressed; + size_t m_decompressed_bytes_read = 0; + size_t m_bytes_read; + u64 m_data_offset; + u32 m_rvz_packed_size; + + u32 m_size = 0; + bool m_junk; + LaggedFibonacciGenerator m_lfg; +}; + +class Compressor +{ +public: + virtual ~Compressor(); + + // First call Start, then AddDataOnlyForPurgeHashing/Compress any number of times, + // then End, then GetData/GetSize any number of times. + + virtual bool Start() = 0; + virtual bool AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) { return true; } + virtual bool Compress(const u8* data, size_t size) = 0; + virtual bool End() = 0; + + virtual const u8* GetData() const = 0; + virtual size_t GetSize() const = 0; +}; + +class PurgeCompressor final : public Compressor +{ +public: + PurgeCompressor(); + ~PurgeCompressor(); + + bool Start() override; + bool AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) override; + bool Compress(const u8* data, size_t size) override; + bool End() override; + + const u8* GetData() const override; + size_t GetSize() const override; + +private: + std::vector<u8> m_buffer; + size_t m_bytes_written; + mbedtls_sha1_context m_sha1_context; +}; + +class Bzip2Compressor final : public Compressor +{ +public: + Bzip2Compressor(int compression_level); + ~Bzip2Compressor(); + + bool Start() override; + bool Compress(const u8* data, size_t size) override; + bool End() override; + + const u8* GetData() const override; + size_t GetSize() const override; + +private: + void ExpandBuffer(size_t bytes_to_add); + + bz_stream m_stream = {}; + std::vector<u8> m_buffer; + int m_compression_level; +}; + +class LZMACompressor final : public Compressor +{ +public: + LZMACompressor(bool lzma2, int compression_level, u8 compressor_data_out[7], + u8* compressor_data_size_out); + ~LZMACompressor(); + + bool Start() override; + bool Compress(const u8* data, size_t size) override; + bool End() override; + + const u8* GetData() const override; + size_t GetSize() const override; + +private: + void ExpandBuffer(size_t bytes_to_add); + + lzma_stream m_stream = LZMA_STREAM_INIT; + lzma_options_lzma m_options = {}; + lzma_filter m_filters[2]; + std::vector<u8> m_buffer; + bool m_initialization_failed = false; +}; + +class ZstdCompressor final : public Compressor +{ +public: + ZstdCompressor(int compression_level); + ~ZstdCompressor(); + + bool Start() override; + bool Compress(const u8* data, size_t size) override; + bool End() override; + + const u8* GetData() const override { return m_buffer.data(); } + size_t GetSize() const override { return m_out_buffer.pos; } + +private: + void ExpandBuffer(size_t bytes_to_add); + + ZSTD_CStream* m_stream; + ZSTD_outBuffer m_out_buffer; + std::vector<u8> m_buffer; +}; + +} // namespace DiscIO diff --git a/Source/Core/DiscIO/WiiEncryptionCache.cpp b/Source/Core/DiscIO/WiiEncryptionCache.cpp index c5a2111daa..0bb892da71 100644 --- a/Source/Core/DiscIO/WiiEncryptionCache.cpp +++ b/Source/Core/DiscIO/WiiEncryptionCache.cpp @@ -24,7 +24,8 @@ WiiEncryptionCache::~WiiEncryptionCache() = default; const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>* WiiEncryptionCache::EncryptGroup(u64 offset, u64 partition_data_offset, - u64 partition_data_decrypted_size, const Key& key) + u64 partition_data_decrypted_size, const Key& key, + const HashExceptionCallback& hash_exception_callback) { // Only allocate memory if this function actually ends up getting called if (!m_cache) @@ -40,8 +41,20 @@ WiiEncryptionCache::EncryptGroup(u64 offset, u64 partition_data_offset, if (m_cached_offset != group_offset_on_disc) { + std::function<void(VolumeWii::HashBlock * hash_blocks)> hash_exception_callback_2; + + if (hash_exception_callback) + { + hash_exception_callback_2 = + [offset, &hash_exception_callback]( + VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP]) { + return hash_exception_callback(hash_blocks, offset); + }; + } + if (!VolumeWii::EncryptGroup(group_offset_in_partition, partition_data_offset, - partition_data_decrypted_size, key, m_blob, m_cache.get())) + partition_data_decrypted_size, key, m_blob, m_cache.get(), + hash_exception_callback_2)) { m_cached_offset = std::numeric_limits<u64>::max(); // Invalidate the cache return nullptr; @@ -54,13 +67,14 @@ WiiEncryptionCache::EncryptGroup(u64 offset, u64 partition_data_offset, } bool WiiEncryptionCache::EncryptGroups(u64 offset, u64 size, u8* out_ptr, u64 partition_data_offset, - u64 partition_data_decrypted_size, const Key& key) + u64 partition_data_decrypted_size, const Key& key, + const HashExceptionCallback& hash_exception_callback) { while (size > 0) { const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>* group = EncryptGroup(Common::AlignDown(offset, VolumeWii::GROUP_TOTAL_SIZE), partition_data_offset, - partition_data_decrypted_size, key); + partition_data_decrypted_size, key, hash_exception_callback); if (!group) return false; diff --git a/Source/Core/DiscIO/WiiEncryptionCache.h b/Source/Core/DiscIO/WiiEncryptionCache.h index 0c8f4b489a..d0a48b054d 100644 --- a/Source/Core/DiscIO/WiiEncryptionCache.h +++ b/Source/Core/DiscIO/WiiEncryptionCache.h @@ -19,6 +19,8 @@ class WiiEncryptionCache { public: using Key = std::array<u8, VolumeWii::AES_KEY_SIZE>; + using HashExceptionCallback = std::function<void( + VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP], u64 offset)>; // The blob pointer is kept around for the lifetime of this object. explicit WiiEncryptionCache(BlobReader* blob); @@ -28,15 +30,15 @@ public: // If the returned pointer is nullptr, reading from the blob failed. // If the returned pointer is not nullptr, it is guaranteed to be valid until // the next call of this function or the destruction of this object. - const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>* EncryptGroup(u64 offset, - u64 partition_data_offset, - u64 partition_data_decrypted_size, - const Key& key); + const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>* + EncryptGroup(u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size, + const Key& key, const HashExceptionCallback& hash_exception_callback = {}); // Encrypts a variable number of groups, as determined by the offset and size parameters. // Supports reading groups partially. bool EncryptGroups(u64 offset, u64 size, u8* out_ptr, u64 partition_data_offset, - u64 partition_data_decrypted_size, const Key& key); + u64 partition_data_decrypted_size, const Key& key, + const HashExceptionCallback& hash_exception_callback = {}); private: BlobReader* m_blob; diff --git a/Source/Core/DolphinQt/ConvertDialog.cpp b/Source/Core/DolphinQt/ConvertDialog.cpp index 0d01785dc4..d7676cc81a 100644 --- a/Source/Core/DolphinQt/ConvertDialog.cpp +++ b/Source/Core/DolphinQt/ConvertDialog.cpp @@ -26,6 +26,7 @@ #include "Common/Logging/Log.h" #include "DiscIO/Blob.h" #include "DiscIO/ScrubbedBlob.h" +#include "DiscIO/WIABlob.h" #include "DolphinQt/QtUtils/ModalMessageBox.h" #include "DolphinQt/QtUtils/ParallelProgressDialog.h" #include "UICommon/GameFile.h" @@ -57,6 +58,8 @@ ConvertDialog::ConvertDialog(QList<std::shared_ptr<const UICommon::GameFile>> fi m_format = new QComboBox; m_format->addItem(QStringLiteral("ISO"), static_cast<int>(DiscIO::BlobType::PLAIN)); m_format->addItem(QStringLiteral("GCZ"), static_cast<int>(DiscIO::BlobType::GCZ)); + m_format->addItem(QStringLiteral("WIA"), static_cast<int>(DiscIO::BlobType::WIA)); + m_format->addItem(QStringLiteral("RVZ"), static_cast<int>(DiscIO::BlobType::RVZ)); if (std::all_of(m_files.begin(), m_files.end(), [](const auto& file) { return file->GetBlobType() == DiscIO::BlobType::PLAIN; })) { @@ -69,9 +72,17 @@ ConvertDialog::ConvertDialog(QList<std::shared_ptr<const UICommon::GameFile>> fi grid_layout->addWidget(new QLabel(tr("Block Size:")), 1, 0); grid_layout->addWidget(m_block_size, 1, 1); + m_compression = new QComboBox; + grid_layout->addWidget(new QLabel(tr("Compression:")), 2, 0); + grid_layout->addWidget(m_compression, 2, 1); + + m_compression_level = new QComboBox; + grid_layout->addWidget(new QLabel(tr("Compression Level:")), 3, 0); + grid_layout->addWidget(m_compression_level, 3, 1); + m_scrub = new QCheckBox; - grid_layout->addWidget(new QLabel(tr("Remove Junk Data (Irreversible):")), 2, 0); - grid_layout->addWidget(m_scrub, 2, 1); + grid_layout->addWidget(new QLabel(tr("Remove Junk Data (Irreversible):")), 4, 0); + grid_layout->addWidget(m_scrub, 4, 1); m_scrub->setEnabled( std::none_of(m_files.begin(), m_files.end(), std::mem_fn(&UICommon::GameFile::IsDatelDisc))); @@ -83,12 +94,17 @@ ConvertDialog::ConvertDialog(QList<std::shared_ptr<const UICommon::GameFile>> fi QGroupBox* options_group = new QGroupBox(tr("Options")); options_group->setLayout(options_layout); - QLabel* info_text = - new QLabel(tr("ISO: A simple and robust format which is supported by many programs. " - "It takes up more space than any other format.\n\n" - "GCZ: A basic compressed format which is compatible with most versions of " - "Dolphin and some other programs. It can't efficiently compress junk data " - "(unless removed) or encrypted Wii data.")); + QLabel* info_text = new QLabel( + tr("ISO: A simple and robust format which is supported by many programs. It takes up more " + "space than any other format.\n\n" + "GCZ: A basic compressed format which is compatible with most versions of Dolphin and " + "some other programs. It can't efficiently compress junk data (unless removed) or " + "encrypted Wii data.\n\n" + "WIA: An advanced compressed format which is compatible with recent versions of Dolphin " + "and a few other programs. It can efficiently compress encrypted Wii data, but not junk " + "data (unless removed).\n\n" + "RVZ: An advanced compressed format which is compatible with recent versions of Dolphin. " + "It can efficiently compress both junk data and encrypted Wii data.")); info_text->setWordWrap(true); QVBoxLayout* info_layout = new QVBoxLayout; @@ -104,14 +120,34 @@ ConvertDialog::ConvertDialog(QList<std::shared_ptr<const UICommon::GameFile>> fi connect(m_format, QOverload<int>::of(&QComboBox::currentIndexChanged), this, &ConvertDialog::OnFormatChanged); + connect(m_compression, QOverload<int>::of(&QComboBox::currentIndexChanged), this, + &ConvertDialog::OnCompressionChanged); connect(convert_button, &QPushButton::clicked, this, &ConvertDialog::Convert); OnFormatChanged(); + OnCompressionChanged(); } void ConvertDialog::AddToBlockSizeComboBox(int size) { m_block_size->addItem(QString::fromStdString(UICommon::FormatSize(size, 0)), size); + + // Select 128 KiB by default, or if it is not available, the size closest to it. + // This code assumes that sizes get added to the combo box in increasing order. + constexpr int DEFAULT_SIZE = 0x20000; + if (size <= DEFAULT_SIZE) + m_block_size->setCurrentIndex(m_block_size->count() - 1); +} + +void ConvertDialog::AddToCompressionComboBox(const QString& name, + DiscIO::WIARVZCompressionType type) +{ + m_compression->addItem(name, static_cast<int>(type)); +} + +void ConvertDialog::AddToCompressionLevelComboBox(int level) +{ + m_compression_level->addItem(QString::number(level), level); } void ConvertDialog::OnFormatChanged() @@ -127,6 +163,9 @@ void ConvertDialog::OnFormatChanged() const DiscIO::BlobType format = static_cast<DiscIO::BlobType>(m_format->currentData().toInt()); m_block_size->clear(); + m_compression->clear(); + + // Populate m_block_size switch (format) { case DiscIO::BlobType::GCZ: @@ -166,11 +205,90 @@ void ConvertDialog::OnFormatChanged() break; } + case DiscIO::BlobType::WIA: + m_block_size->setEnabled(true); + + // This is the smallest block size supported by WIA. For performance, larger sizes are avoided. + AddToBlockSizeComboBox(0x200000); + + break; + case DiscIO::BlobType::RVZ: + m_block_size->setEnabled(true); + + for (int block_size = MIN_BLOCK_SIZE; block_size <= MAX_BLOCK_SIZE; block_size *= 2) + AddToBlockSizeComboBox(block_size); + + break; + default: + break; + } + + // Populate m_compression + switch (format) + { + case DiscIO::BlobType::GCZ: + m_compression->setEnabled(true); + AddToCompressionComboBox(QStringLiteral("Deflate"), DiscIO::WIARVZCompressionType::None); + break; + case DiscIO::BlobType::WIA: + case DiscIO::BlobType::RVZ: + { + m_compression->setEnabled(true); + + // i18n: %1 is the name of a compression method (e.g. LZMA) + const QString slow = tr("%1 (slow)"); + + AddToCompressionComboBox(tr("No Compression"), DiscIO::WIARVZCompressionType::None); + + if (format == DiscIO::BlobType::WIA) + AddToCompressionComboBox(QStringLiteral("Purge"), DiscIO::WIARVZCompressionType::Purge); + + AddToCompressionComboBox(slow.arg(QStringLiteral("bzip2")), + DiscIO::WIARVZCompressionType::Bzip2); + + AddToCompressionComboBox(slow.arg(QStringLiteral("LZMA")), DiscIO::WIARVZCompressionType::LZMA); + + AddToCompressionComboBox(slow.arg(QStringLiteral("LZMA2")), + DiscIO::WIARVZCompressionType::LZMA2); + + if (format == DiscIO::BlobType::RVZ) + { + AddToCompressionComboBox(QStringLiteral("Zstandard"), DiscIO::WIARVZCompressionType::Zstd); + m_compression->setCurrentIndex(m_compression->count() - 1); + } + + break; + } default: + m_compression->setEnabled(false); break; } m_block_size->setEnabled(m_block_size->count() > 1); + m_compression->setEnabled(m_compression->count() > 1); + + m_scrub->setEnabled(format != DiscIO::BlobType::RVZ); + if (format == DiscIO::BlobType::RVZ) + m_scrub->setChecked(false); +} + +void ConvertDialog::OnCompressionChanged() +{ + m_compression_level->clear(); + + const auto compression_type = + static_cast<DiscIO::WIARVZCompressionType>(m_compression->currentData().toInt()); + + const std::pair<int, int> range = DiscIO::GetAllowedCompressionLevels(compression_type); + + for (int i = range.first; i <= range.second; ++i) + { + AddToCompressionLevelComboBox(i); + if (i == 5) + m_compression_level->setCurrentIndex(m_compression_level->count() - 1); + } + + m_compression_level->setEnabled(m_compression_level->count() > 1); } bool ConvertDialog::ShowAreYouSureDialog(const QString& text) @@ -189,6 +307,9 @@ void ConvertDialog::Convert() { const DiscIO::BlobType format = static_cast<DiscIO::BlobType>(m_format->currentData().toInt()); const int block_size = m_block_size->currentData().toInt(); + const DiscIO::WIARVZCompressionType compression = + static_cast<DiscIO::WIARVZCompressionType>(m_compression->currentData().toInt()); + const int compression_level = m_compression_level->currentData().toInt(); const bool scrub = m_scrub->isChecked(); if (scrub && format == DiscIO::BlobType::PLAIN) @@ -224,7 +345,15 @@ void ConvertDialog::Convert() break; case DiscIO::BlobType::GCZ: extension = QStringLiteral(".gcz"); - filter = tr("Compressed GC/Wii images (*.gcz)"); + filter = tr("GCZ GC/Wii images (*.gcz)"); + break; + case DiscIO::BlobType::WIA: + extension = QStringLiteral(".wia"); + filter = tr("WIA GC/Wii images (*.wia)"); + break; + case DiscIO::BlobType::RVZ: + extension = QStringLiteral(".rvz"); + filter = tr("RVZ GC/Wii images (*.rvz)"); break; default: ASSERT(false); @@ -330,8 +459,9 @@ void ConvertDialog::Convert() { std::future<bool> good; - if (format == DiscIO::BlobType::PLAIN) + switch (format) { + case DiscIO::BlobType::PLAIN: good = std::async(std::launch::async, [&] { const bool good = DiscIO::ConvertToPlain(blob_reader.get(), original_path, dst_path.toStdString(), @@ -339,9 +469,9 @@ void ConvertDialog::Convert() progress_dialog.Reset(); return good; }); - } - else if (format == DiscIO::BlobType::GCZ) - { + break; + + case DiscIO::BlobType::GCZ: good = std::async(std::launch::async, [&] { const bool good = DiscIO::ConvertToGCZ(blob_reader.get(), original_path, dst_path.toStdString(), @@ -350,6 +480,19 @@ void ConvertDialog::Convert() progress_dialog.Reset(); return good; }); + break; + + case DiscIO::BlobType::WIA: + case DiscIO::BlobType::RVZ: + good = std::async(std::launch::async, [&] { + const bool good = DiscIO::ConvertToWIAOrRVZ( + blob_reader.get(), original_path, dst_path.toStdString(), + format == DiscIO::BlobType::RVZ, compression, compression_level, block_size, + &CompressCB, &progress_dialog); + progress_dialog.Reset(); + return good; + }); + break; } progress_dialog.GetRaw()->exec(); diff --git a/Source/Core/DolphinQt/ConvertDialog.h b/Source/Core/DolphinQt/ConvertDialog.h index 1a3a9bbf4a..9a53265f86 100644 --- a/Source/Core/DolphinQt/ConvertDialog.h +++ b/Source/Core/DolphinQt/ConvertDialog.h @@ -14,6 +14,11 @@ class QCheckBox; class QComboBox; +namespace DiscIO +{ +enum class WIARVZCompressionType : u32; +} + namespace UICommon { class GameFile; @@ -29,15 +34,20 @@ public: private slots: void OnFormatChanged(); + void OnCompressionChanged(); void Convert(); private: void AddToBlockSizeComboBox(int size); + void AddToCompressionComboBox(const QString& name, DiscIO::WIARVZCompressionType type); + void AddToCompressionLevelComboBox(int level); bool ShowAreYouSureDialog(const QString& text); QComboBox* m_format; QComboBox* m_block_size; + QComboBox* m_compression; + QComboBox* m_compression_level; QCheckBox* m_scrub; QList<std::shared_ptr<const UICommon::GameFile>> m_files; }; diff --git a/Source/Core/DolphinQt/GameList/GameTracker.cpp b/Source/Core/DolphinQt/GameList/GameTracker.cpp index 4caf643dd5..1f2b7383dd 100644 --- a/Source/Core/DolphinQt/GameList/GameTracker.cpp +++ b/Source/Core/DolphinQt/GameList/GameTracker.cpp @@ -24,6 +24,7 @@ static const QStringList game_filters{ QStringLiteral("*.[gG][cC][mM]"), QStringLiteral("*.[iI][sS][oO]"), QStringLiteral("*.[tT][gG][cC]"), QStringLiteral("*.[cC][iI][sS][oO]"), QStringLiteral("*.[gG][cC][zZ]"), QStringLiteral("*.[wW][bB][fF][sS]"), + QStringLiteral("*.[wW][iI][aA]"), QStringLiteral("*.[rR][vV][zZ]"), QStringLiteral("*.[wW][aA][dD]"), QStringLiteral("*.[eE][lL][fF]"), QStringLiteral("*.[dD][oO][lL]")}; diff --git a/Source/Core/DolphinQt/Info.plist.in b/Source/Core/DolphinQt/Info.plist.in index 7fd60fa8e0..b5f3a3f44b 100644 --- a/Source/Core/DolphinQt/Info.plist.in +++ b/Source/Core/DolphinQt/Info.plist.in @@ -14,8 +14,10 @@ <string>gcz</string> <string>iso</string> <string>m3u</string> + <string>rvz</string> <string>tgc</string> <string>wad</string> + <string>wia</string> <string>wbfs</string> </array> <key>CFBundleTypeIconFile</key> diff --git a/Source/Core/DolphinQt/MainWindow.cpp b/Source/Core/DolphinQt/MainWindow.cpp index 8b8eb37f40..1ecb875986 100644 --- a/Source/Core/DolphinQt/MainWindow.cpp +++ b/Source/Core/DolphinQt/MainWindow.cpp @@ -686,8 +686,8 @@ QStringList MainWindow::PromptFileNames() QStringList paths = QFileDialog::getOpenFileNames( this, tr("Select a File"), settings.value(QStringLiteral("mainwindow/lastdir"), QString{}).toString(), - tr("All GC/Wii files (*.elf *.dol *.gcm *.iso *.tgc *.wbfs *.ciso *.gcz *.wad *.dff *.m3u);;" - "All Files (*)")); + tr("All GC/Wii files (*.elf *.dol *.gcm *.iso *.tgc *.wbfs *.ciso *.gcz *.wia *.rvz *.wad " + "*.dff *.m3u);;All Files (*)")); if (!paths.isEmpty()) { diff --git a/Source/Core/DolphinQt/Settings/PathPane.cpp b/Source/Core/DolphinQt/Settings/PathPane.cpp index ab59f185cc..fd85d42a22 100644 --- a/Source/Core/DolphinQt/Settings/PathPane.cpp +++ b/Source/Core/DolphinQt/Settings/PathPane.cpp @@ -42,10 +42,10 @@ void PathPane::Browse() void PathPane::BrowseDefaultGame() { - QString file = QDir::toNativeSeparators(QFileDialog::getOpenFileName( - this, tr("Select a Game"), Settings::Instance().GetDefaultGame(), - tr("All GC/Wii files (*.elf *.dol *.gcm *.iso *.tgc *.wbfs *.ciso *.gcz *.wad *.m3u);;" - "All Files (*)"))); + QString file = QDir::toNativeSeparators( + QFileDialog::getOpenFileName(this, tr("Select a Game"), Settings::Instance().GetDefaultGame(), + tr("All GC/Wii files (*.elf *.dol *.gcm *.iso *.tgc *.wbfs " + "*.ciso *.gcz *.wia *.rvz *.wad *.m3u);;All Files (*)"))); if (!file.isEmpty()) Settings::Instance().SetDefaultGame(file); diff --git a/Source/Core/UICommon/GameFileCache.cpp b/Source/Core/UICommon/GameFileCache.cpp index 9b2729407e..1af33ae44b 100644 --- a/Source/Core/UICommon/GameFileCache.cpp +++ b/Source/Core/UICommon/GameFileCache.cpp @@ -33,7 +33,7 @@ std::vector<std::string> FindAllGamePaths(const std::vector<std::string>& direct bool recursive_scan) { static const std::vector<std::string> search_extensions = { - ".gcm", ".tgc", ".iso", ".ciso", ".gcz", ".wbfs", ".wad", ".dol", ".elf"}; + ".gcm", ".tgc", ".iso", ".ciso", ".gcz", ".wbfs", ".wia", ".rvz", ".wad", ".dol", ".elf"}; // TODO: We could process paths iteratively as they are found return Common::DoFileSearch(directories_to_scan, search_extensions, recursive_scan); |
