WIA: Make use of the exception lists

Source/Core/DiscIO/VolumeWii.cpp

@@ -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,
+bool VolumeWii::EncryptGroup(
-                             u64 partition_data_decrypted_size,
+    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)
 {
   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()));
 

Source/Core/DiscIO/VolumeWii.h

@@ -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; }

Source/Core/DiscIO/WIABlob.cpp

@@ -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

Source/Core/DiscIO/WIABlob.h

@@ -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)
     {

Source/Core/DiscIO/WiiEncryptionCache.cpp

@@ -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;

Source/Core/DiscIO/WiiEncryptionCache.h

@@ -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,
+  const std::array<u8, VolumeWii::GROUP_TOTAL_SIZE>*
-                                                                  u64 partition_data_offset,
+  EncryptGroup(u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size,
-                                                                  u64 partition_data_decrypted_size,
+               const Key& key, const HashExceptionCallback& hash_exception_callback = {});
-                                                                  const Key& key);
 
   // 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;