Skip to content
Permalink
Browse files

Merge pull request #8299 from JosJuice/volumeverifier-performance

VolumeVerifier: Performance improvements
  • Loading branch information...
stenzek committed Aug 9, 2019
2 parents 22ed2c0 + 34fb608 commit 2df522d4eb27961dc10c3423f52277f5d3a84234
@@ -76,6 +76,12 @@ class Volume
}
virtual std::vector<u8> GetContent(u16 index) const { return {}; }
virtual std::vector<u64> GetContentOffsets() const { return {}; }
virtual bool CheckContentIntegrity(const IOS::ES::Content& content,
const std::vector<u8>& 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<u8>& encrypted_data,
const Partition& partition) const
{
return false;
}
virtual bool CheckBlockIntegrity(u64 block_index, const Partition& partition) const
{
return false;
@@ -6,7 +6,9 @@

#include <algorithm>
#include <cinttypes>
#include <future>
#include <limits>
#include <mutex>
#include <optional>
#include <string>
#include <unordered_set>
@@ -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<u8> 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,70 +740,114 @@ 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)
{
bytes_to_read = std::min(bytes_to_read, m_blocks[m_block_index].offset - m_progress);
}
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<u8> 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;
}
else
{
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<unsigned int>(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<unsigned int>(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)
@@ -4,7 +4,9 @@

#pragma once

#include <future>
#include <map>
#include <mutex>
#include <optional>
#include <string>
#include <vector>
@@ -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<u8> m_data;
std::mutex m_volume_mutex;
std::future<void> m_crc32_future;
std::future<void> m_md5_future;
std::future<void> m_sha1_future;
std::future<void> m_content_future;
std::future<void> m_block_future;

DiscScrubber m_scrubber;
IOS::ES::TicketReader m_ticket;
std::vector<u64> m_content_offsets;
@@ -156,6 +156,9 @@ bool VolumeWAD::CheckContentIntegrity(const IOS::ES::Content& content,
const std::vector<u8>& 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<u8, 16> key = ticket.GetTitleKey();
mbedtls_aes_setkey_dec(&context, key.data(), 128);
@@ -39,6 +39,8 @@ class VolumeWAD : public Volume
GetCertificateChain(const Partition& partition = PARTITION_NONE) const override;
std::vector<u8> GetContent(u16 index) const override;
std::vector<u64> GetContentOffsets() const override;
bool CheckContentIntegrity(const IOS::ES::Content& content, const std::vector<u8>& 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<u8>& encrypted_data,
const IOS::ES::TicketReader& ticket) const;

std::unique_ptr<BlobReader> m_reader;
IOS::ES::TicketReader m_ticket;
IOS::ES::TMDReader m_tmd;
@@ -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<u8>& 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<u8> cluster(BLOCK_TOTAL_SIZE);
if (!m_reader->Read(cluster_offset, cluster.size(), cluster.data()))
return false;
return CheckBlockIntegrity(block_index, cluster, partition);
}

} // namespace DiscIO
@@ -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<u8>& encrypted_data,
const Partition& partition) const override;
bool CheckBlockIntegrity(u64 block_index, const Partition& partition) const override;

Region GetRegion() const override;

0 comments on commit 2df522d

Please sign in to comment.
You can’t perform that action at this time.