WIA: Write hash exceptions

Source/Core/DiscIO/VolumeWii.cpp

@@ -473,9 +473,7 @@ bool VolumeWii::CheckBlockIntegrity(u64 block_index, const std::vector<u8>& encr
     return false;
 
   HashBlock hashes;
-  u8 iv[16] = {0};
+  DecryptBlockHashes(encrypted_data.data(), &hashes, aes_context);
-  mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, sizeof(HashBlock), iv,
-                        encrypted_data.data(), reinterpret_cast<u8*>(&hashes));
 
   u8 cluster_data[BLOCK_DATA_SIZE];
   DecryptBlockData(encrypted_data.data(), cluster_data, aes_context);
@@ -521,55 +519,33 @@ bool VolumeWii::CheckBlockIntegrity(u64 block_index, const Partition& partition)
   return CheckBlockIntegrity(block_index, cluster, partition);
 }
 
-bool VolumeWii::EncryptGroup(
+bool VolumeWii::HashGroup(const std::array<u8, BLOCK_DATA_SIZE> in[BLOCKS_PER_GROUP],
-    u64 offset, u64 partition_data_offset, u64 partition_data_decrypted_size,
+                          HashBlock out[BLOCKS_PER_GROUP],
-    const std::array<u8, AES_KEY_SIZE>& key, BlobReader* blob,
+                          const std::function<bool(size_t block)>& read_function)
-    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);
-
   std::array<std::future<void>, BLOCKS_PER_GROUP> hash_futures;
-  bool error_occurred = false;
+  bool success = true;
 
   for (size_t i = 0; i < BLOCKS_PER_GROUP; ++i)
   {
-    if (!error_occurred)
+    if (read_function && success)
-    {
+      success = read_function(i);
-      if (offset + (i + 1) * BLOCK_DATA_SIZE <= partition_data_decrypted_size)
-      {
-        if (!blob->ReadWiiDecrypted(offset + i * BLOCK_DATA_SIZE, BLOCK_DATA_SIZE,
-                                    unencrypted_data[i].data(), partition_data_offset))
-        {
-          error_occurred = true;
-        }
-      }
-      else
-      {
-        unencrypted_data[i].fill(0);
-      }
-    }
 
-    hash_futures[i] = std::async(std::launch::async, [&unencrypted_data, &unencrypted_hashes,
+    hash_futures[i] = std::async(std::launch::async, [&in, &out, &hash_futures, success, i]() {
-                                                      &hash_futures, error_occurred, i]() {
       const size_t h1_base = Common::AlignDown(i, 8);
 
-      if (!error_occurred)
+      if (success)
       {
         // H0 hashes
         for (size_t j = 0; j < 31; ++j)
-        {
+          mbedtls_sha1_ret(in[i].data() + j * 0x400, 0x400, out[i].h0[j]);
-          mbedtls_sha1_ret(unencrypted_data[i].data() + j * 0x400, 0x400,
-                           unencrypted_hashes[i].h0[j]);
-        }
 
         // H0 padding
-        std::memset(unencrypted_hashes[i].padding_0, 0, sizeof(HashBlock::padding_0));
+        std::memset(out[i].padding_0, 0, sizeof(HashBlock::padding_0));
 
         // H1 hash
-        mbedtls_sha1_ret(reinterpret_cast<u8*>(unencrypted_hashes[i].h0), sizeof(HashBlock::h0),
+        mbedtls_sha1_ret(reinterpret_cast<u8*>(out[i].h0), sizeof(HashBlock::h0),
-                         unencrypted_hashes[h1_base].h1[i - h1_base]);
+                         out[h1_base].h1[i - h1_base]);
       }
 
       if (i % 8 == 7)
@@ -577,21 +553,18 @@ bool VolumeWii::EncryptGroup(
         for (size_t j = 0; j < 7; ++j)
           hash_futures[h1_base + j].get();
 
-        if (!error_occurred)
+        if (success)
         {
           // H1 padding
-          std::memset(unencrypted_hashes[h1_base].padding_1, 0, sizeof(HashBlock::padding_1));
+          std::memset(out[h1_base].padding_1, 0, sizeof(HashBlock::padding_1));
 
           // H1 copies
           for (size_t j = 1; j < 8; ++j)
-          {
+            std::memcpy(out[h1_base + j].h1, out[h1_base].h1, sizeof(HashBlock::h1));
-            std::memcpy(unencrypted_hashes[h1_base + j].h1, unencrypted_hashes[h1_base].h1,
-                        sizeof(HashBlock::h1));
-          }
 
           // H2 hash
-          mbedtls_sha1_ret(reinterpret_cast<u8*>(unencrypted_hashes[i].h1), sizeof(HashBlock::h1),
+          mbedtls_sha1_ret(reinterpret_cast<u8*>(out[i].h1), sizeof(HashBlock::h1),
-                           unencrypted_hashes[0].h2[h1_base / 8]);
+                           out[0].h2[h1_base / 8]);
         }
 
         if (i == BLOCKS_PER_GROUP - 1)
@@ -599,17 +572,14 @@ bool VolumeWii::EncryptGroup(
           for (size_t j = 0; j < 7; ++j)
             hash_futures[j * 8 + 7].get();
 
-          if (!error_occurred)
+          if (success)
           {
             // H2 padding
-            std::memset(unencrypted_hashes[0].padding_2, 0, sizeof(HashBlock::padding_2));
+            std::memset(out[0].padding_2, 0, sizeof(HashBlock::padding_2));
 
             // H2 copies
             for (size_t j = 1; j < BLOCKS_PER_GROUP; ++j)
-            {
+              std::memcpy(out[j].h2, out[0].h2, sizeof(HashBlock::h2));
-              std::memcpy(unencrypted_hashes[j].h2, unencrypted_hashes[0].h2,
-                          sizeof(HashBlock::h2));
-            }
           }
         }
       }
@@ -619,7 +589,36 @@ bool VolumeWii::EncryptGroup(
   // Wait for all the async tasks to finish
   hash_futures.back().get();
 
-  if (error_occurred)
+  return success;
+}
+
+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);
+
+  const bool success =
+      HashGroup(unencrypted_data.data(), unencrypted_hashes.data(), [&](size_t block) {
+        if (offset + (block + 1) * BLOCK_DATA_SIZE <= partition_data_decrypted_size)
+        {
+          if (!blob->ReadWiiDecrypted(offset + block * BLOCK_DATA_SIZE, BLOCK_DATA_SIZE,
+                                      unencrypted_data[block].data(), partition_data_offset))
+          {
+            return false;
+          }
+        }
+        else
+        {
+          unencrypted_data[block].fill(0);
+        }
+        return true;
+      });
+
+  if (!success)
     return false;
 
   if (hash_exception_callback)
@@ -660,6 +659,14 @@ bool VolumeWii::EncryptGroup(
   return true;
 }
 
+void VolumeWii::DecryptBlockHashes(const u8* in, HashBlock* out, mbedtls_aes_context* aes_context)
+{
+  std::array<u8, 16> iv;
+  iv.fill(0);
+  mbedtls_aes_crypt_cbc(aes_context, MBEDTLS_AES_DECRYPT, sizeof(HashBlock), iv.data(), in,
+                        reinterpret_cast<u8*>(out));
+}
+
 void VolumeWii::DecryptBlockData(const u8* in, u8* out, mbedtls_aes_context* aes_context)
 {
   std::array<u8, 16> iv;

Source/Core/DiscIO/VolumeWii.h

@@ -98,12 +98,21 @@ public:
   u64 GetRawSize() const override;
   const BlobReader& GetBlobReader() const;
 
+  // The in parameter can either contain all the data to begin with,
+  // or read_function can write data into the in parameter when called.
+  // The latter lets reading run in parallel with hashing.
+  // This function returns false iff read_function returns false.
+  static bool HashGroup(const std::array<u8, BLOCK_DATA_SIZE> in[BLOCKS_PER_GROUP],
+                        HashBlock out[BLOCKS_PER_GROUP],
+                        const std::function<bool(size_t block)>& read_function = {});
+
   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,
                            const std::function<void(HashBlock hash_blocks[BLOCKS_PER_GROUP])>&
                                hash_exception_callback = {});
 
+  static void DecryptBlockHashes(const u8* in, HashBlock* out, mbedtls_aes_context* aes_context);
   static void DecryptBlockData(const u8* in, u8* out, mbedtls_aes_context* aes_context);
 
 protected:

Source/Core/DiscIO/WIABlob.cpp

@@ -1068,6 +1068,7 @@ WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile,
   ASSERT(chunk_size > 0);
 
   const u64 iso_size = infile->GetDataSize();
+  const u64 exception_lists_per_chunk = chunk_size / VolumeWii::GROUP_TOTAL_SIZE;
 
   u64 bytes_read = 0;
   u64 bytes_written = 0;
@@ -1120,8 +1121,12 @@ WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile,
     return ConversionResult::ReadFailed;
   // We intentially do not increment bytes_read here, since these bytes will be read again
 
+  using WiiBlockData = std::array<u8, VolumeWii::BLOCK_DATA_SIZE>;
+
   std::vector<u8> buffer(chunk_size);
-  std::vector<u8> decryption_buffer(VolumeWii::BLOCK_DATA_SIZE);
+  std::vector<WiiBlockData> decryption_buffer(VolumeWii::BLOCKS_PER_GROUP);
+  std::vector<VolumeWii::HashBlock> hash_buffer(VolumeWii::BLOCKS_PER_GROUP);
+
   for (const DataEntry& data_entry : data_entries)
   {
     if (data_entry.is_partition)
@@ -1148,39 +1153,109 @@ WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile,
       {
         const u64 bytes_to_read = std::min<u64>(chunk_size, data_offset + data_size - bytes_read);
 
+        const u64 groups = Common::AlignUp(bytes_to_read, VolumeWii::GROUP_TOTAL_SIZE) /
+                           VolumeWii::GROUP_TOTAL_SIZE;
+
         ASSERT(bytes_to_read % VolumeWii::BLOCK_TOTAL_SIZE == 0);
-        const u64 bytes_to_write =
+        const u64 blocks = bytes_to_read / VolumeWii::BLOCK_TOTAL_SIZE;
-            bytes_to_read / VolumeWii::BLOCK_TOTAL_SIZE * VolumeWii::BLOCK_DATA_SIZE;
+        const u64 bytes_to_write = blocks * VolumeWii::BLOCK_DATA_SIZE;
 
         if (!infile->Read(bytes_read, bytes_to_read, buffer.data()))
           return ConversionResult::ReadFailed;
 
-        const u64 exception_lists = Common::AlignUp(bytes_to_read, VolumeWii::GROUP_TOTAL_SIZE) /
+        std::vector<std::vector<HashExceptionEntry>> exception_lists(exception_lists_per_chunk);
-                                    VolumeWii::GROUP_TOTAL_SIZE;
 
-        const u64 exceptions_size = Common::AlignUp(exception_lists * sizeof(u16), 4);
+        for (u64 j = 0; j < groups; ++j)
-        const u64 total_size = exceptions_size + bytes_to_write;
-        ASSERT((bytes_written & 3) == 0);
-        group_entries[i].data_offset = Common::swap32(static_cast<u32>(bytes_written >> 2));
-        group_entries[i].data_size = Common::swap32(static_cast<u32>(total_size));
-
-        for (u64 j = 0; j < exception_lists; ++j)
         {
-          const u16 exceptions = 0;
+          const u64 offset_of_group = j * VolumeWii::GROUP_TOTAL_SIZE;
+          const u64 write_offset_of_group = j * VolumeWii::GROUP_DATA_SIZE;
+
+          const u64 blocks_in_this_group =
+              std::min<u64>(VolumeWii::BLOCKS_PER_GROUP, blocks - j * VolumeWii::BLOCKS_PER_GROUP);
+
+          for (u32 k = 0; k < VolumeWii::BLOCKS_PER_GROUP; ++k)
+          {
+            if (k < blocks_in_this_group)
+            {
+              const u64 offset_of_block = offset_of_group + k * VolumeWii::BLOCK_TOTAL_SIZE;
+              VolumeWii::DecryptBlockData(buffer.data() + offset_of_block,
+                                          decryption_buffer[k].data(), &aes_context);
+            }
+            else
+            {
+              decryption_buffer[k].fill(0);
+            }
+          }
+
+          VolumeWii::HashGroup(decryption_buffer.data(), hash_buffer.data());
+
+          for (u64 k = 0; k < blocks_in_this_group; ++k)
+          {
+            const u64 offset_of_block = offset_of_group + k * VolumeWii::BLOCK_TOTAL_SIZE;
+            const u64 hash_offset_of_block = k * VolumeWii::BLOCK_HEADER_SIZE;
+
+            VolumeWii::HashBlock hashes;
+            VolumeWii::DecryptBlockHashes(buffer.data() + offset_of_block, &hashes, &aes_context);
+
+            const auto compare_hash = [&](size_t offset_in_block) {
+              ASSERT(offset_in_block + sizeof(SHA1) <= VolumeWii::BLOCK_HEADER_SIZE);
+
+              const u8* desired_hash = reinterpret_cast<u8*>(&hashes) + offset_in_block;
+              const u8* computed_hash = reinterpret_cast<u8*>(&hash_buffer[k]) + offset_in_block;
+
+              if (!std::equal(desired_hash, desired_hash + sizeof(SHA1), computed_hash))
+              {
+                const u64 hash_offset = hash_offset_of_block + offset_in_block;
+                ASSERT(hash_offset <= std::numeric_limits<u16>::max());
+
+                HashExceptionEntry& exception = exception_lists[j].emplace_back();
+                exception.offset = static_cast<u16>(Common::swap16(hash_offset));
+                std::memcpy(exception.hash.data(), desired_hash, sizeof(SHA1));
+              }
+            };
+
+            const auto compare_hashes = [&compare_hash](size_t offset, size_t size) {
+              for (size_t l = 0; l < size; l += sizeof(SHA1))
+                // The std::min is to ensure that we don't go beyond the end of HashBlock with
+                // padding_2, which is 32 bytes long (not divisible by sizeof(SHA1), which is 20).
+                compare_hash(offset + std::min(l, size - sizeof(SHA1)));
+            };
+
+            using HashBlock = VolumeWii::HashBlock;
+            compare_hashes(offsetof(HashBlock, h0), sizeof(HashBlock::h0));
+            compare_hashes(offsetof(HashBlock, padding_0), sizeof(HashBlock::padding_0));
+            compare_hashes(offsetof(HashBlock, h1), sizeof(HashBlock::h1));
+            compare_hashes(offsetof(HashBlock, padding_1), sizeof(HashBlock::padding_1));
+            compare_hashes(offsetof(HashBlock, h2), sizeof(HashBlock::h2));
+            compare_hashes(offsetof(HashBlock, padding_2), sizeof(HashBlock::padding_2));
+          }
+
+          for (u64 k = 0; k < blocks_in_this_group; ++k)
+          {
+            std::memcpy(buffer.data() + write_offset_of_group + k * VolumeWii::BLOCK_DATA_SIZE,
+                        decryption_buffer[k].data(), VolumeWii::BLOCK_DATA_SIZE);
+          }
+        }
+
+        const u64 write_offset = bytes_written;
+
+        for (const std::vector<HashExceptionEntry>& exception_list : exception_lists)
+        {
+          const u16 exceptions = Common::swap16(static_cast<u16>(exception_list.size()));
           if (!outfile->WriteArray(&exceptions, 1))
             return ConversionResult::WriteFailed;
-          bytes_written += sizeof(u16);
+
+          if (!outfile->WriteArray(exception_list.data(), exception_list.size()))
+            return ConversionResult::WriteFailed;
+
+          bytes_written += sizeof(u16) + exception_list.size() * sizeof(HashExceptionEntry);
         }
 
         if (!PadTo4(outfile, &bytes_written))
           return ConversionResult::WriteFailed;
 
-        for (u64 j = 0; j < bytes_to_read; j += VolumeWii::BLOCK_TOTAL_SIZE)
+        if (!outfile->WriteArray(buffer.data(), bytes_to_write))
-        {
+          return ConversionResult::WriteFailed;
-          VolumeWii::DecryptBlockData(buffer.data() + j, decryption_buffer.data(), &aes_context);
-          if (!outfile->WriteArray(decryption_buffer.data(), VolumeWii::BLOCK_DATA_SIZE))
-            return ConversionResult::WriteFailed;
-        }
 
         bytes_read += bytes_to_read;
         bytes_written += bytes_to_write;
@@ -1188,6 +1263,10 @@ WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile,
         if (!PadTo4(outfile, &bytes_written))
           return ConversionResult::WriteFailed;
 
+        ASSERT((write_offset & 3) == 0);
+        group_entries[i].data_offset = Common::swap32(static_cast<u32>(write_offset >> 2));
+        group_entries[i].data_size = Common::swap32(static_cast<u32>(bytes_written - write_offset));
+
         if (!run_callback())
           return ConversionResult::Canceled;
       }