From f754a1a548f2e599e4cc2d1d799f2d21b520076f Mon Sep 17 00:00:00 2001 From: JosJuice Date: Mon, 5 Aug 2019 13:58:27 +0200 Subject: [PATCH 1/2] VolumeVerifier: Don't read data multiple times --- Source/Core/DiscIO/Volume.h | 11 ++++++++ Source/Core/DiscIO/VolumeVerifier.cpp | 36 ++++++++++++++++++++------- Source/Core/DiscIO/VolumeWad.cpp | 3 +++ Source/Core/DiscIO/VolumeWad.h | 5 ++-- Source/Core/DiscIO/VolumeWii.cpp | 35 ++++++++++++++++++-------- Source/Core/DiscIO/VolumeWii.h | 2 ++ 6 files changed, 69 insertions(+), 23 deletions(-) diff --git a/Source/Core/DiscIO/Volume.h b/Source/Core/DiscIO/Volume.h index 05325aced2..81d986f547 100644 --- a/Source/Core/DiscIO/Volume.h +++ b/Source/Core/DiscIO/Volume.h @@ -76,6 +76,12 @@ public: } virtual std::vector GetContent(u16 index) const { return {}; } virtual std::vector GetContentOffsets() const { return {}; } + virtual bool CheckContentIntegrity(const IOS::ES::Content& content, + const std::vector& encrypted_data, + const IOS::ES::TicketReader& ticket) const + { + return false; + } virtual bool CheckContentIntegrity(const IOS::ES::Content& content, u64 content_offset, const IOS::ES::TicketReader& ticket) const { @@ -109,6 +115,11 @@ public: virtual Platform GetVolumeType() const = 0; virtual bool SupportsIntegrityCheck() const { return false; } virtual bool CheckH3TableIntegrity(const Partition& partition) const { return false; } + virtual bool CheckBlockIntegrity(u64 block_index, const std::vector& encrypted_data, + const Partition& partition) const + { + return false; + } virtual bool CheckBlockIntegrity(u64 block_index, const Partition& partition) const { return false; diff --git a/Source/Core/DiscIO/VolumeVerifier.cpp b/Source/Core/DiscIO/VolumeVerifier.cpp index 3850adad22..33a3a03e98 100644 --- a/Source/Core/DiscIO/VolumeVerifier.cpp +++ b/Source/Core/DiscIO/VolumeVerifier.cpp @@ -691,8 +691,9 @@ void VolumeVerifier::Process() if (m_progress == m_max_progress) return; - IOS::ES::Content content; + IOS::ES::Content content{}; bool content_read = false; + bool block_read = false; u64 bytes_to_read = BLOCK_SIZE; if (m_content_index < m_content_offsets.size() && m_content_offsets[m_content_index] == m_progress) @@ -709,6 +710,7 @@ void VolumeVerifier::Process() else if (m_block_index < m_blocks.size() && m_blocks[m_block_index].offset == m_progress) { bytes_to_read = VolumeWii::BLOCK_TOTAL_SIZE; + block_read = true; } else if (m_block_index < m_blocks.size() && m_blocks[m_block_index].offset > m_progress) { @@ -716,10 +718,17 @@ void VolumeVerifier::Process() } bytes_to_read = std::min(bytes_to_read, m_max_progress - m_progress); + bool read_succeeded = false; + std::vector data(bytes_to_read); + if (m_calculating_any_hash || content_read || block_read) + { + if (m_volume.Read(m_progress, bytes_to_read, data.data(), PARTITION_NONE)) + read_succeeded = true; + } + if (m_calculating_any_hash) { - std::vector data(bytes_to_read); - if (!m_volume.Read(m_progress, bytes_to_read, data.data(), PARTITION_NONE)) + if (!read_succeeded) { m_calculating_any_hash = false; } @@ -740,11 +749,9 @@ void VolumeVerifier::Process() } } - m_progress += bytes_to_read; - if (content_read) { - if (!m_volume.CheckContentIntegrity(content, m_content_offsets[m_content_index], m_ticket)) + if (!read_succeeded || !m_volume.CheckContentIntegrity(content, data, m_ticket)) { AddProblem( Severity::High, @@ -754,10 +761,16 @@ void VolumeVerifier::Process() m_content_index++; } - while (m_block_index < m_blocks.size() && m_blocks[m_block_index].offset < m_progress) + while (m_block_index < m_blocks.size() && + m_blocks[m_block_index].offset < m_progress + bytes_to_read) { - if (!m_volume.CheckBlockIntegrity(m_blocks[m_block_index].block_index, - m_blocks[m_block_index].partition)) + const bool success = m_blocks[m_block_index].offset == m_progress ? + read_succeeded && m_volume.CheckBlockIntegrity( + m_blocks[m_block_index].block_index, data, + m_blocks[m_block_index].partition) : + m_volume.CheckBlockIntegrity(m_blocks[m_block_index].block_index, + m_blocks[m_block_index].partition); + if (!success) { const u64 offset = m_blocks[m_block_index].offset; if (m_scrubber.CanBlockBeScrubbed(offset)) @@ -773,6 +786,8 @@ void VolumeVerifier::Process() } m_block_index++; } + + m_progress += bytes_to_read; } u64 VolumeVerifier::GetBytesProcessed() const @@ -791,6 +806,9 @@ void VolumeVerifier::Finish() return; m_done = true; + ASSERT(m_content_index == m_content_offsets.size()); + ASSERT(m_block_index == m_blocks.size()); + if (m_calculating_any_hash) { if (m_hashes_to_calculate.crc32) diff --git a/Source/Core/DiscIO/VolumeWad.cpp b/Source/Core/DiscIO/VolumeWad.cpp index cc83f59081..e471edcd1d 100644 --- a/Source/Core/DiscIO/VolumeWad.cpp +++ b/Source/Core/DiscIO/VolumeWad.cpp @@ -156,6 +156,9 @@ bool VolumeWAD::CheckContentIntegrity(const IOS::ES::Content& content, const std::vector& encrypted_data, const IOS::ES::TicketReader& ticket) const { + if (encrypted_data.size() != Common::AlignUp(content.size, 0x40)) + return false; + mbedtls_aes_context context; const std::array key = ticket.GetTitleKey(); mbedtls_aes_setkey_dec(&context, key.data(), 128); diff --git a/Source/Core/DiscIO/VolumeWad.h b/Source/Core/DiscIO/VolumeWad.h index 02bac94e62..85df789228 100644 --- a/Source/Core/DiscIO/VolumeWad.h +++ b/Source/Core/DiscIO/VolumeWad.h @@ -39,6 +39,8 @@ public: GetCertificateChain(const Partition& partition = PARTITION_NONE) const override; std::vector GetContent(u16 index) const override; std::vector GetContentOffsets() const override; + bool CheckContentIntegrity(const IOS::ES::Content& content, const std::vector& encrypted_data, + const IOS::ES::TicketReader& ticket) const override; bool CheckContentIntegrity(const IOS::ES::Content& content, u64 content_offset, const IOS::ES::TicketReader& ticket) const override; IOS::ES::TicketReader GetTicketWithFixedCommonKey() const override; @@ -66,9 +68,6 @@ public: u64 GetRawSize() const override; private: - bool CheckContentIntegrity(const IOS::ES::Content& content, const std::vector& encrypted_data, - const IOS::ES::TicketReader& ticket) const; - std::unique_ptr m_reader; IOS::ES::TicketReader m_ticket; IOS::ES::TMDReader m_tmd; diff --git a/Source/Core/DiscIO/VolumeWii.cpp b/Source/Core/DiscIO/VolumeWii.cpp index 8ab3722873..5b6524f807 100644 --- a/Source/Core/DiscIO/VolumeWii.cpp +++ b/Source/Core/DiscIO/VolumeWii.cpp @@ -447,8 +447,12 @@ bool VolumeWii::CheckH3TableIntegrity(const Partition& partition) const return h3_table_sha1 == contents[0].sha1; } -bool VolumeWii::CheckBlockIntegrity(u64 block_index, const Partition& partition) const +bool VolumeWii::CheckBlockIntegrity(u64 block_index, const std::vector& encrypted_data, + const Partition& partition) const { + if (encrypted_data.size() != BLOCK_TOTAL_SIZE) + return false; + auto it = m_partitions.find(partition); if (it == m_partitions.end()) return false; @@ -462,21 +466,15 @@ bool VolumeWii::CheckBlockIntegrity(u64 block_index, const Partition& partition) if (!aes_context) return false; - const u64 cluster_offset = - partition.offset + *partition_details.data_offset + block_index * BLOCK_TOTAL_SIZE; - - // Read and decrypt the cluster metadata - u8 cluster_metadata_crypted[BLOCK_HEADER_SIZE]; u8 cluster_metadata[BLOCK_HEADER_SIZE]; u8 iv[16] = {0}; - if (!m_reader->Read(cluster_offset, BLOCK_HEADER_SIZE, cluster_metadata_crypted)) - return false; mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, BLOCK_HEADER_SIZE, iv, - cluster_metadata_crypted, cluster_metadata); + encrypted_data.data(), cluster_metadata); u8 cluster_data[BLOCK_DATA_SIZE]; - if (!Read(block_index * BLOCK_DATA_SIZE, BLOCK_DATA_SIZE, cluster_data, partition)) - return false; + std::memcpy(iv, encrypted_data.data() + 0x3D0, 16); + mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, BLOCK_DATA_SIZE, iv, + encrypted_data.data() + BLOCK_HEADER_SIZE, cluster_data); for (u32 hash_index = 0; hash_index < 31; ++hash_index) { @@ -504,4 +502,19 @@ bool VolumeWii::CheckBlockIntegrity(u64 block_index, const Partition& partition) return true; } +bool VolumeWii::CheckBlockIntegrity(u64 block_index, const Partition& partition) const +{ + auto it = m_partitions.find(partition); + if (it == m_partitions.end()) + return false; + const PartitionDetails& partition_details = it->second; + const u64 cluster_offset = + partition.offset + *partition_details.data_offset + block_index * BLOCK_TOTAL_SIZE; + + std::vector cluster(BLOCK_TOTAL_SIZE); + if (!m_reader->Read(cluster_offset, cluster.size(), cluster.data())) + return false; + return CheckBlockIntegrity(block_index, cluster, partition); +} + } // namespace DiscIO diff --git a/Source/Core/DiscIO/VolumeWii.h b/Source/Core/DiscIO/VolumeWii.h index 38dfc1290b..fb55169ad0 100644 --- a/Source/Core/DiscIO/VolumeWii.h +++ b/Source/Core/DiscIO/VolumeWii.h @@ -58,6 +58,8 @@ public: Platform GetVolumeType() const override; bool SupportsIntegrityCheck() const override { return m_encrypted; } bool CheckH3TableIntegrity(const Partition& partition) const override; + bool CheckBlockIntegrity(u64 block_index, const std::vector& encrypted_data, + const Partition& partition) const override; bool CheckBlockIntegrity(u64 block_index, const Partition& partition) const override; Region GetRegion() const override; From 34fb608dd6b4137ecc98e1c03228165a27eb800a Mon Sep 17 00:00:00 2001 From: JosJuice Date: Tue, 6 Aug 2019 11:39:09 +0200 Subject: [PATCH 2/2] VolumeVerifier: Multithreading --- Source/Core/DiscIO/VolumeVerifier.cpp | 140 +++++++++++++++++++------- Source/Core/DiscIO/VolumeVerifier.h | 12 +++ 2 files changed, 113 insertions(+), 39 deletions(-) diff --git a/Source/Core/DiscIO/VolumeVerifier.cpp b/Source/Core/DiscIO/VolumeVerifier.cpp index 33a3a03e98..4d1bdc9a40 100644 --- a/Source/Core/DiscIO/VolumeVerifier.cpp +++ b/Source/Core/DiscIO/VolumeVerifier.cpp @@ -6,7 +6,9 @@ #include #include +#include #include +#include #include #include #include @@ -683,6 +685,34 @@ void VolumeVerifier::SetUpHashing() } } +void VolumeVerifier::WaitForAsyncOperations() const +{ + if (m_crc32_future.valid()) + m_crc32_future.wait(); + if (m_md5_future.valid()) + m_md5_future.wait(); + if (m_sha1_future.valid()) + m_sha1_future.wait(); + if (m_content_future.valid()) + m_content_future.wait(); + if (m_block_future.valid()) + m_block_future.wait(); +} + +bool VolumeVerifier::ReadChunkAndWaitForAsyncOperations(u64 bytes_to_read) +{ + std::vector data(bytes_to_read); + { + std::lock_guard lk(m_volume_mutex); + if (!m_volume.Read(m_progress, bytes_to_read, data.data(), PARTITION_NONE)) + return false; + } + + WaitForAsyncOperations(); + m_data = std::move(data); + return true; +} + void VolumeVerifier::Process() { ASSERT(m_started); @@ -718,13 +748,8 @@ void VolumeVerifier::Process() } bytes_to_read = std::min(bytes_to_read, m_max_progress - m_progress); - bool read_succeeded = false; - std::vector data(bytes_to_read); - if (m_calculating_any_hash || content_read || block_read) - { - if (m_volume.Read(m_progress, bytes_to_read, data.data(), PARTITION_NONE)) - read_succeeded = true; - } + const bool is_data_needed = m_calculating_any_hash || content_read || block_read; + const bool read_succeeded = is_data_needed && ReadChunkAndWaitForAsyncOperations(bytes_to_read); if (m_calculating_any_hash) { @@ -736,55 +761,90 @@ void VolumeVerifier::Process() { if (m_hashes_to_calculate.crc32) { - // It would be nice to use crc32_z here instead of crc32, but it isn't available on Android - m_crc32_context = - crc32(m_crc32_context, data.data(), static_cast(bytes_to_read)); + m_crc32_future = std::async(std::launch::async, [this] { + // Would be nice to use crc32_z here instead of crc32, but it isn't available on Android + m_crc32_context = + crc32(m_crc32_context, m_data.data(), static_cast(m_data.size())); + }); } if (m_hashes_to_calculate.md5) - mbedtls_md5_update_ret(&m_md5_context, data.data(), bytes_to_read); + { + m_md5_future = std::async(std::launch::async, [this] { + mbedtls_md5_update_ret(&m_md5_context, m_data.data(), m_data.size()); + }); + } if (m_hashes_to_calculate.sha1) - mbedtls_sha1_update_ret(&m_sha1_context, data.data(), bytes_to_read); + { + m_sha1_future = std::async(std::launch::async, [this] { + mbedtls_sha1_update_ret(&m_sha1_context, m_data.data(), m_data.size()); + }); + } } } if (content_read) { - if (!read_succeeded || !m_volume.CheckContentIntegrity(content, data, m_ticket)) - { - AddProblem( - Severity::High, - StringFromFormat(Common::GetStringT("Content %08x is corrupt.").c_str(), content.id)); - } + m_content_future = std::async(std::launch::async, [this, read_succeeded, content] { + if (!read_succeeded || !m_volume.CheckContentIntegrity(content, m_data, m_ticket)) + { + AddProblem( + Severity::High, + StringFromFormat(Common::GetStringT("Content %08x is corrupt.").c_str(), content.id)); + } + }); m_content_index++; } - while (m_block_index < m_blocks.size() && - m_blocks[m_block_index].offset < m_progress + bytes_to_read) + if (m_block_index < m_blocks.size() && + m_blocks[m_block_index].offset < m_progress + bytes_to_read) { - const bool success = m_blocks[m_block_index].offset == m_progress ? - read_succeeded && m_volume.CheckBlockIntegrity( - m_blocks[m_block_index].block_index, data, - m_blocks[m_block_index].partition) : - m_volume.CheckBlockIntegrity(m_blocks[m_block_index].block_index, - m_blocks[m_block_index].partition); - if (!success) + m_md5_future = std::async( + std::launch::async, + [this, read_succeeded, bytes_to_read](size_t block_index, u64 progress) { + while (block_index < m_blocks.size() && + m_blocks[block_index].offset < progress + bytes_to_read) + { + bool success; + if (m_blocks[block_index].offset == progress) + { + success = read_succeeded && + m_volume.CheckBlockIntegrity(m_blocks[block_index].block_index, m_data, + m_blocks[block_index].partition); + } + else + { + std::lock_guard lk(m_volume_mutex); + success = m_volume.CheckBlockIntegrity(m_blocks[block_index].block_index, + m_blocks[block_index].partition); + } + + if (!success) + { + const u64 offset = m_blocks[block_index].offset; + if (m_scrubber.CanBlockBeScrubbed(offset)) + { + WARN_LOG(DISCIO, "Integrity check failed for unused block at 0x%" PRIx64, offset); + m_unused_block_errors[m_blocks[block_index].partition]++; + } + else + { + WARN_LOG(DISCIO, "Integrity check failed for block at 0x%" PRIx64, offset); + m_block_errors[m_blocks[block_index].partition]++; + } + } + block_index++; + } + }, + m_block_index, m_progress); + + while (m_block_index < m_blocks.size() && + m_blocks[m_block_index].offset < m_progress + bytes_to_read) { - const u64 offset = m_blocks[m_block_index].offset; - if (m_scrubber.CanBlockBeScrubbed(offset)) - { - WARN_LOG(DISCIO, "Integrity check failed for unused block at 0x%" PRIx64, offset); - m_unused_block_errors[m_blocks[m_block_index].partition]++; - } - else - { - WARN_LOG(DISCIO, "Integrity check failed for block at 0x%" PRIx64, offset); - m_block_errors[m_blocks[m_block_index].partition]++; - } + m_block_index++; } - m_block_index++; } m_progress += bytes_to_read; @@ -806,6 +866,8 @@ void VolumeVerifier::Finish() return; m_done = true; + WaitForAsyncOperations(); + ASSERT(m_content_index == m_content_offsets.size()); ASSERT(m_block_index == m_blocks.size()); diff --git a/Source/Core/DiscIO/VolumeVerifier.h b/Source/Core/DiscIO/VolumeVerifier.h index 9bcf2fa8c2..713836de5b 100644 --- a/Source/Core/DiscIO/VolumeVerifier.h +++ b/Source/Core/DiscIO/VolumeVerifier.h @@ -4,7 +4,9 @@ #pragma once +#include #include +#include #include #include #include @@ -98,6 +100,8 @@ private: u64 GetBiggestUsedOffset(const FileInfo& file_info) const; void CheckMisc(); void SetUpHashing(); + void WaitForAsyncOperations() const; + bool ReadChunkAndWaitForAsyncOperations(u64 bytes_to_read); void AddProblem(Severity severity, std::string text); @@ -113,6 +117,14 @@ private: mbedtls_md5_context m_md5_context; mbedtls_sha1_context m_sha1_context; + std::vector m_data; + std::mutex m_volume_mutex; + std::future m_crc32_future; + std::future m_md5_future; + std::future m_sha1_future; + std::future m_content_future; + std::future m_block_future; + DiscScrubber m_scrubber; IOS::ES::TicketReader m_ticket; std::vector m_content_offsets;