WIA: Make use of the exception lists

This commit is contained in:
JosJuice 2020-01-27 16:12:56 +01:00
parent 47067f661a
commit 791e363c9a
6 changed files with 93 additions and 15 deletions

View File

@ -532,10 +532,11 @@ 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<u8, AES_KEY_SIZE>& key, BlobReader* blob,
std::array<u8, GROUP_TOTAL_SIZE>* out)
bool VolumeWii::EncryptGroup(
u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size,
const std::array<u8, AES_KEY_SIZE>& key, BlobReader* blob,
std::array<u8, GROUP_TOTAL_SIZE>* out,
const std::function<void(HashBlock hash_blocks[BLOCKS_PER_GROUP])>& hash_exception_callback)
{
std::vector<std::array<u8, BLOCK_DATA_SIZE>> unencrypted_data(BLOCKS_PER_GROUP);
std::vector<HashBlock> unencrypted_hashes(BLOCKS_PER_GROUP);
@ -632,6 +633,9 @@ bool VolumeWii::EncryptGroup(u64 offset, u64 partition_data_offset,
if (error_occurred)
return false;
if (hash_exception_callback)
hash_exception_callback(unencrypted_hashes.data());
const unsigned int threads =
std::min(BLOCKS_PER_GROUP, std::max<unsigned int>(1, std::thread::hardware_concurrency()));

View File

@ -5,6 +5,7 @@
#pragma once
#include <array>
#include <functional>
#include <map>
#include <memory>
#include <optional>
@ -99,7 +100,9 @@ public:
static bool EncryptGroup(u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size,
const std::array<u8, AES_KEY_SIZE>& key, BlobReader* blob,
std::array<u8, GROUP_TOTAL_SIZE>* out);
std::array<u8, GROUP_TOTAL_SIZE>* out,
const std::function<void(HashBlock hash_blocks[BLOCKS_PER_GROUP])>&
hash_exception_callback = {});
protected:
u32 GetOffsetShift() const override { return 2; }

View File

@ -277,12 +277,26 @@ bool WIAFileReader::Read(u64 offset, u64 size, u8* out_ptr)
const u64 bytes_to_read = std::min(data_size - (offset - data_offset), size);
bool hash_exception_error = false;
if (!m_encryption_cache.EncryptGroups(
offset - partition_data_offset, bytes_to_read, out_ptr, partition_data_offset,
partition_total_sectors * VolumeWii::BLOCK_DATA_SIZE, partition.partition_key))
partition_total_sectors * VolumeWii::BLOCK_DATA_SIZE, partition.partition_key,
[this, chunk_size, first_sector, partition_first_sector, &hash_exception_error](
VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP], u64 offset) {
const u64 partition_part_offset =
(first_sector - partition_first_sector) * VolumeWii::BLOCK_TOTAL_SIZE;
const u64 index =
(offset - partition_part_offset) % chunk_size / VolumeWii::GROUP_TOTAL_SIZE;
// EncryptGroups calls ReadWiiDecrypted, which populates m_cached_chunk
if (!m_cached_chunk.ApplyHashExceptions(hash_blocks, index))
hash_exception_error = true;
}))
{
return false;
}
if (hash_exception_error)
return false;
offset += bytes_to_read;
size -= bytes_to_read;
@ -836,4 +850,41 @@ bool WIAFileReader::Chunk::HandleExceptions(const u8* data, size_t bytes_allocat
return true;
}
bool WIAFileReader::Chunk::ApplyHashExceptions(
VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP], u64 exception_list_index) const
{
if (m_exception_lists > 0)
return false; // We still have exception lists left to read
const u8* data = m_compressed_exception_lists ? m_out.data.data() : m_in.data.data();
for (u64 i = exception_list_index; i > 0; --i)
data += Common::swap16(data) * sizeof(HashExceptionEntry) + sizeof(u16);
const u16 exceptions = Common::swap16(data);
data += sizeof(u16);
for (size_t i = 0; i < exceptions; ++i)
{
HashExceptionEntry exception;
std::memcpy(&exception, data, sizeof(HashExceptionEntry));
data += sizeof(HashExceptionEntry);
const u16 offset = Common::swap16(exception.offset);
const size_t block_index = offset / VolumeWii::BLOCK_HEADER_SIZE;
if (block_index > VolumeWii::BLOCKS_PER_GROUP)
return false;
const size_t offset_in_block = offset % VolumeWii::BLOCK_HEADER_SIZE;
if (offset_in_block + sizeof(SHA1) > VolumeWii::BLOCK_HEADER_SIZE)
return false;
std::memcpy(reinterpret_cast<u8*>(&hash_blocks[block_index]) + offset_in_block, &exception.hash,
sizeof(SHA1));
}
return true;
}
} // namespace DiscIO

View File

@ -234,6 +234,10 @@ private:
bool Read(u64 offset, u64 size, u8* out_ptr);
// This can only be called once at least one byte of data has been read
bool ApplyHashExceptions(VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP],
u64 exception_list_index) const;
template <typename T>
bool ReadAll(std::vector<T>* vector)
{

View File

@ -24,7 +24,8 @@ WiiEncryptionCache::~WiiEncryptionCache() = default;
const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>*
WiiEncryptionCache::EncryptGroup(u64 offset, u64 partition_data_offset,
u64 partition_data_decrypted_size, const Key& key)
u64 partition_data_decrypted_size, const Key& key,
const HashExceptionCallback& hash_exception_callback)
{
// Only allocate memory if this function actually ends up getting called
if (!m_cache)
@ -40,8 +41,20 @@ WiiEncryptionCache::EncryptGroup(u64 offset, u64 partition_data_offset,
if (m_cached_offset != group_offset_on_disc)
{
std::function<void(VolumeWii::HashBlock * hash_blocks)> hash_exception_callback_2;
if (hash_exception_callback)
{
hash_exception_callback_2 =
[offset, &hash_exception_callback](
VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP]) {
return hash_exception_callback(hash_blocks, offset);
};
}
if (!VolumeWii::EncryptGroup(group_offset_in_partition, partition_data_offset,
partition_data_decrypted_size, key, m_blob, m_cache.get()))
partition_data_decrypted_size, key, m_blob, m_cache.get(),
hash_exception_callback_2))
{
m_cached_offset = std::numeric_limits<u64>::max(); // Invalidate the cache
return nullptr;
@ -54,13 +67,14 @@ WiiEncryptionCache::EncryptGroup(u64 offset, u64 partition_data_offset,
}
bool WiiEncryptionCache::EncryptGroups(u64 offset, u64 size, u8* out_ptr, u64 partition_data_offset,
u64 partition_data_decrypted_size, const Key& key)
u64 partition_data_decrypted_size, const Key& key,
const HashExceptionCallback& hash_exception_callback)
{
while (size > 0)
{
const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>* group =
EncryptGroup(Common::AlignDown(offset, VolumeWii::GROUP_TOTAL_SIZE), partition_data_offset,
partition_data_decrypted_size, key);
partition_data_decrypted_size, key, hash_exception_callback);
if (!group)
return false;

View File

@ -19,6 +19,8 @@ class WiiEncryptionCache
{
public:
using Key = std::array<u8, VolumeWii::AES_KEY_SIZE>;
using HashExceptionCallback = std::function<void(
VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP], u64 offset)>;
// The blob pointer is kept around for the lifetime of this object.
explicit WiiEncryptionCache(BlobReader* blob);
@ -28,15 +30,15 @@ public:
// If the returned pointer is nullptr, reading from the blob failed.
// If the returned pointer is not nullptr, it is guaranteed to be valid until
// the next call of this function or the destruction of this object.
const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>* EncryptGroup(u64 offset,
u64 partition_data_offset,
u64 partition_data_decrypted_size,
const Key& key);
const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>*
EncryptGroup(u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size,
const Key& key, const HashExceptionCallback& hash_exception_callback = {});
// Encrypts a variable number of groups, as determined by the offset and size parameters.
// Supports reading groups partially.
bool EncryptGroups(u64 offset, u64 size, u8* out_ptr, u64 partition_data_offset,
u64 partition_data_decrypted_size, const Key& key);
u64 partition_data_decrypted_size, const Key& key,
const HashExceptionCallback& hash_exception_callback = {});
private:
BlobReader* m_blob;