From e936c4acd821a12f4d45d163df66cf062f64f902 Mon Sep 17 00:00:00 2001 From: JosJuice Date: Wed, 15 Apr 2020 21:15:08 +0200 Subject: [PATCH] WIA: Write hash exceptions --- Source/Core/DiscIO/VolumeWii.cpp | 111 ++++++++++++++-------------- Source/Core/DiscIO/VolumeWii.h | 9 +++ Source/Core/DiscIO/WIABlob.cpp | 119 +++++++++++++++++++++++++------ 3 files changed, 167 insertions(+), 72 deletions(-) diff --git a/Source/Core/DiscIO/VolumeWii.cpp b/Source/Core/DiscIO/VolumeWii.cpp index 227cedb5cf..6e2b40041f 100644 --- a/Source/Core/DiscIO/VolumeWii.cpp +++ b/Source/Core/DiscIO/VolumeWii.cpp @@ -473,9 +473,7 @@ bool VolumeWii::CheckBlockIntegrity(u64 block_index, const std::vector& 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(&hashes)); + DecryptBlockHashes(encrypted_data.data(), &hashes, aes_context); u8 cluster_data[BLOCK_DATA_SIZE]; DecryptBlockData(encrypted_data.data(), cluster_data, aes_context); @@ -521,55 +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& key, BlobReader* blob, - std::array* out, - const std::function& hash_exception_callback) +bool VolumeWii::HashGroup(const std::array in[BLOCKS_PER_GROUP], + HashBlock out[BLOCKS_PER_GROUP], + const std::function& read_function) { - std::vector> unencrypted_data(BLOCKS_PER_GROUP); - std::vector unencrypted_hashes(BLOCKS_PER_GROUP); - std::array, 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(unencrypted_hashes[i].h0), sizeof(HashBlock::h0), - unencrypted_hashes[h1_base].h1[i - h1_base]); + mbedtls_sha1_ret(reinterpret_cast(out[i].h0), sizeof(HashBlock::h0), + out[h1_base].h1[i - h1_base]); } if (i % 8 == 7) @@ -577,21 +553,18 @@ bool VolumeWii::EncryptGroup( 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(unencrypted_hashes[i].h1), sizeof(HashBlock::h1), - unencrypted_hashes[0].h2[h1_base / 8]); + mbedtls_sha1_ret(reinterpret_cast(out[i].h1), sizeof(HashBlock::h1), + out[0].h2[h1_base / 8]); } if (i == BLOCKS_PER_GROUP - 1) @@ -599,17 +572,14 @@ bool VolumeWii::EncryptGroup( 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)); } } } @@ -619,7 +589,36 @@ bool VolumeWii::EncryptGroup( // 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& key, BlobReader* blob, + std::array* out, + const std::function& hash_exception_callback) +{ + std::vector> unencrypted_data(BLOCKS_PER_GROUP); + std::vector 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) @@ -660,6 +659,14 @@ bool VolumeWii::EncryptGroup( return true; } +void VolumeWii::DecryptBlockHashes(const u8* in, HashBlock* out, mbedtls_aes_context* aes_context) +{ + std::array iv; + iv.fill(0); + mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, sizeof(HashBlock), iv.data(), in, + reinterpret_cast(out)); +} + void VolumeWii::DecryptBlockData(const u8* in, u8* out, mbedtls_aes_context* aes_context) { std::array iv; diff --git a/Source/Core/DiscIO/VolumeWii.h b/Source/Core/DiscIO/VolumeWii.h index b0fdf713cb..0f2bb43d01 100644 --- a/Source/Core/DiscIO/VolumeWii.h +++ b/Source/Core/DiscIO/VolumeWii.h @@ -98,12 +98,21 @@ 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 in[BLOCKS_PER_GROUP], + HashBlock out[BLOCKS_PER_GROUP], + const std::function& read_function = {}); + static bool EncryptGroup(u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size, const std::array& key, BlobReader* blob, std::array* out, const std::function& 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: diff --git a/Source/Core/DiscIO/WIABlob.cpp b/Source/Core/DiscIO/WIABlob.cpp index af11451c61..157a59f01b 100644 --- a/Source/Core/DiscIO/WIABlob.cpp +++ b/Source/Core/DiscIO/WIABlob.cpp @@ -1068,6 +1068,7 @@ WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile, ASSERT(chunk_size > 0); const u64 iso_size = infile->GetDataSize(); + const u64 exception_lists_per_chunk = chunk_size / VolumeWii::GROUP_TOTAL_SIZE; u64 bytes_read = 0; u64 bytes_written = 0; @@ -1120,8 +1121,12 @@ WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile, return ConversionResult::ReadFailed; // We intentially do not increment bytes_read here, since these bytes will be read again + using WiiBlockData = std::array; + std::vector buffer(chunk_size); - std::vector decryption_buffer(VolumeWii::BLOCK_DATA_SIZE); + std::vector decryption_buffer(VolumeWii::BLOCKS_PER_GROUP); + std::vector hash_buffer(VolumeWii::BLOCKS_PER_GROUP); + for (const DataEntry& data_entry : data_entries) { if (data_entry.is_partition) @@ -1148,39 +1153,109 @@ WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile, { const u64 bytes_to_read = std::min(chunk_size, data_offset + data_size - bytes_read); + const u64 groups = Common::AlignUp(bytes_to_read, VolumeWii::GROUP_TOTAL_SIZE) / + VolumeWii::GROUP_TOTAL_SIZE; + ASSERT(bytes_to_read % VolumeWii::BLOCK_TOTAL_SIZE == 0); - const u64 bytes_to_write = - bytes_to_read / VolumeWii::BLOCK_TOTAL_SIZE * VolumeWii::BLOCK_DATA_SIZE; + const u64 blocks = bytes_to_read / VolumeWii::BLOCK_TOTAL_SIZE; + const u64 bytes_to_write = blocks * VolumeWii::BLOCK_DATA_SIZE; if (!infile->Read(bytes_read, bytes_to_read, buffer.data())) return ConversionResult::ReadFailed; - const u64 exception_lists = Common::AlignUp(bytes_to_read, VolumeWii::GROUP_TOTAL_SIZE) / - VolumeWii::GROUP_TOTAL_SIZE; + std::vector> exception_lists(exception_lists_per_chunk); - const u64 exceptions_size = Common::AlignUp(exception_lists * sizeof(u16), 4); - const u64 total_size = exceptions_size + bytes_to_write; - ASSERT((bytes_written & 3) == 0); - group_entries[i].data_offset = Common::swap32(static_cast(bytes_written >> 2)); - group_entries[i].data_size = Common::swap32(static_cast(total_size)); - - for (u64 j = 0; j < exception_lists; ++j) + for (u64 j = 0; j < groups; ++j) { - const u16 exceptions = 0; + const u64 offset_of_group = j * VolumeWii::GROUP_TOTAL_SIZE; + const u64 write_offset_of_group = j * VolumeWii::GROUP_DATA_SIZE; + + const u64 blocks_in_this_group = + std::min(VolumeWii::BLOCKS_PER_GROUP, blocks - j * VolumeWii::BLOCKS_PER_GROUP); + + for (u32 k = 0; k < VolumeWii::BLOCKS_PER_GROUP; ++k) + { + if (k < blocks_in_this_group) + { + const u64 offset_of_block = offset_of_group + k * VolumeWii::BLOCK_TOTAL_SIZE; + VolumeWii::DecryptBlockData(buffer.data() + offset_of_block, + decryption_buffer[k].data(), &aes_context); + } + else + { + decryption_buffer[k].fill(0); + } + } + + VolumeWii::HashGroup(decryption_buffer.data(), hash_buffer.data()); + + for (u64 k = 0; k < blocks_in_this_group; ++k) + { + const u64 offset_of_block = offset_of_group + k * VolumeWii::BLOCK_TOTAL_SIZE; + const u64 hash_offset_of_block = k * VolumeWii::BLOCK_HEADER_SIZE; + + VolumeWii::HashBlock hashes; + VolumeWii::DecryptBlockHashes(buffer.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(&hashes) + offset_in_block; + const u8* computed_hash = reinterpret_cast(&hash_buffer[k]) + offset_in_block; + + if (!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::max()); + + HashExceptionEntry& exception = exception_lists[j].emplace_back(); + exception.offset = static_cast(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)); + } + + for (u64 k = 0; k < blocks_in_this_group; ++k) + { + std::memcpy(buffer.data() + write_offset_of_group + k * VolumeWii::BLOCK_DATA_SIZE, + decryption_buffer[k].data(), VolumeWii::BLOCK_DATA_SIZE); + } + } + + const u64 write_offset = bytes_written; + + for (const std::vector& exception_list : exception_lists) + { + const u16 exceptions = Common::swap16(static_cast(exception_list.size())); if (!outfile->WriteArray(&exceptions, 1)) return ConversionResult::WriteFailed; - bytes_written += sizeof(u16); + + if (!outfile->WriteArray(exception_list.data(), exception_list.size())) + return ConversionResult::WriteFailed; + + bytes_written += sizeof(u16) + exception_list.size() * sizeof(HashExceptionEntry); } if (!PadTo4(outfile, &bytes_written)) return ConversionResult::WriteFailed; - for (u64 j = 0; j < bytes_to_read; j += VolumeWii::BLOCK_TOTAL_SIZE) - { - VolumeWii::DecryptBlockData(buffer.data() + j, decryption_buffer.data(), &aes_context); - if (!outfile->WriteArray(decryption_buffer.data(), VolumeWii::BLOCK_DATA_SIZE)) - return ConversionResult::WriteFailed; - } + if (!outfile->WriteArray(buffer.data(), bytes_to_write)) + return ConversionResult::WriteFailed; bytes_read += bytes_to_read; bytes_written += bytes_to_write; @@ -1188,6 +1263,10 @@ WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile, if (!PadTo4(outfile, &bytes_written)) return ConversionResult::WriteFailed; + ASSERT((write_offset & 3) == 0); + group_entries[i].data_offset = Common::swap32(static_cast(write_offset >> 2)); + group_entries[i].data_size = Common::swap32(static_cast(bytes_written - write_offset)); + if (!run_callback()) return ConversionResult::Canceled; }