Merge pull request #8299 from JosJuice/volumeverifier-performance

VolumeVerifier: Performance improvements

Source/Core/DiscIO/Volume.h

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

Source/Core/DiscIO/VolumeVerifier.cpp

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

Source/Core/DiscIO/VolumeVerifier.h

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

Source/Core/DiscIO/VolumeWad.cpp

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

Source/Core/DiscIO/VolumeWad.h

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

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

Source/Core/DiscIO/VolumeWii.h

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