diff --git a/Source/Core/DiscIO/Volume.h b/Source/Core/DiscIO/Volume.h index 05325aced28a..81d986f547c8 100644 --- a/Source/Core/DiscIO/Volume.h +++ b/Source/Core/DiscIO/Volume.h @@ -76,6 +76,12 @@ class Volume } 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 @@ class Volume 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 3850adad226a..4d1bdc9a40b5 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); @@ -691,8 +721,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 +740,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 +748,12 @@ void VolumeVerifier::Process() } bytes_to_read = std::min(bytes_to_read, m_max_progress - m_progress); + 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) { - 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; } @@ -727,52 +761,93 @@ 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()); + }); + } } } - m_progress += bytes_to_read; - if (content_read) { - if (!m_volume.CheckContentIntegrity(content, m_content_offsets[m_content_index], 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) + if (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)) + 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; } u64 VolumeVerifier::GetBytesProcessed() const @@ -791,6 +866,11 @@ void VolumeVerifier::Finish() return; m_done = true; + WaitForAsyncOperations(); + + 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/VolumeVerifier.h b/Source/Core/DiscIO/VolumeVerifier.h index 9bcf2fa8c21f..713836de5b8b 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 @@ class VolumeVerifier final 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 @@ class VolumeVerifier final 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; diff --git a/Source/Core/DiscIO/VolumeWad.cpp b/Source/Core/DiscIO/VolumeWad.cpp index cc83f59081ea..e471edcd1d10 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 02bac94e624b..85df789228a5 100644 --- a/Source/Core/DiscIO/VolumeWad.h +++ b/Source/Core/DiscIO/VolumeWad.h @@ -39,6 +39,8 @@ class VolumeWAD : public Volume 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 @@ class VolumeWAD : public Volume 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 8ab37228732b..5b6524f807c2 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 38dfc1290bdc..fb55169ad0a5 100644 --- a/Source/Core/DiscIO/VolumeWii.h +++ b/Source/Core/DiscIO/VolumeWii.h @@ -58,6 +58,8 @@ class VolumeWii : public VolumeDisc 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;