From f754a1a548f2e599e4cc2d1d799f2d21b520076f Mon Sep 17 00:00:00 2001 From: JosJuice Date: Mon, 5 Aug 2019 13:58:27 +0200 Subject: [PATCH] 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;