VolumeVerifier: Don't read data multiple times

6 files changed, 69 insertions, 23 deletions

diff --git a/Source/Core/DiscIO/Volume.h b/Source/Core/DiscIO/Volume.h
index 05325aced2..81d986f547 100644
--- a/Source/Core/DiscIO/Volume.h
+++ b/Source/Core/DiscIO/Volume.h
@@ -76,6 +76,12 @@ public:
   }
   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 @@ public:
   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;
diff --git a/Source/Core/DiscIO/VolumeVerifier.cpp b/Source/Core/DiscIO/VolumeVerifier.cpp
index 3850adad22..33a3a03e98 100644
--- a/Source/Core/DiscIO/VolumeVerifier.cpp
+++ b/Source/Core/DiscIO/VolumeVerifier.cpp
@@ -691,8 +691,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,6 +710,7 @@ 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)
   {
@@ -716,10 +718,17 @@ void VolumeVerifier::Process()
   }
   bytes_to_read = std::min(bytes_to_read, m_max_progress - m_progress);
 
+  bool read_succeeded = false;
+  std::vector<u8> data(bytes_to_read);
+  if (m_calculating_any_hash || content_read || block_read)
+  {
+    if (m_volume.Read(m_progress, bytes_to_read, data.data(), PARTITION_NONE))
+      read_succeeded = true;
+  }
+
   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;
     }
@@ -740,11 +749,9 @@ void VolumeVerifier::Process()
     }
   }
 
-  m_progress += bytes_to_read;
-
   if (content_read)
   {
-    if (!m_volume.CheckContentIntegrity(content, m_content_offsets[m_content_index], m_ticket))
+    if (!read_succeeded || !m_volume.CheckContentIntegrity(content, data, m_ticket))
     {
       AddProblem(
           Severity::High,
@@ -754,10 +761,16 @@ void VolumeVerifier::Process()
     m_content_index++;
   }
 
-  while (m_block_index < m_blocks.size() && m_blocks[m_block_index].offset < m_progress)
+  while (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))
+    const bool success = m_blocks[m_block_index].offset == m_progress ?
+                             read_succeeded && m_volume.CheckBlockIntegrity(
+                                                   m_blocks[m_block_index].block_index, data,
+                                                   m_blocks[m_block_index].partition) :
+                             m_volume.CheckBlockIntegrity(m_blocks[m_block_index].block_index,
+                                                          m_blocks[m_block_index].partition);
+    if (!success)
     {
       const u64 offset = m_blocks[m_block_index].offset;
       if (m_scrubber.CanBlockBeScrubbed(offset))
@@ -773,6 +786,8 @@ void VolumeVerifier::Process()
     }
     m_block_index++;
   }
+
+  m_progress += bytes_to_read;
 }
 
 u64 VolumeVerifier::GetBytesProcessed() const
@@ -791,6 +806,9 @@ void VolumeVerifier::Finish()
     return;
   m_done = true;
 
+  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)
diff --git a/Source/Core/DiscIO/VolumeWad.cpp b/Source/Core/DiscIO/VolumeWad.cpp
index cc83f59081..e471edcd1d 100644
--- a/Source/Core/DiscIO/VolumeWad.cpp
+++ b/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);
diff --git a/Source/Core/DiscIO/VolumeWad.h b/Source/Core/DiscIO/VolumeWad.h
index 02bac94e62..85df789228 100644
--- a/Source/Core/DiscIO/VolumeWad.h
+++ b/Source/Core/DiscIO/VolumeWad.h
@@ -39,6 +39,8 @@ public:
   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 @@ public:
   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;
diff --git a/Source/Core/DiscIO/VolumeWii.cpp b/Source/Core/DiscIO/VolumeWii.cpp
index 8ab3722873..5b6524f807 100644
--- a/Source/Core/DiscIO/VolumeWii.cpp
+++ b/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
diff --git a/Source/Core/DiscIO/VolumeWii.h b/Source/Core/DiscIO/VolumeWii.h
index 38dfc1290b..fb55169ad0 100644
--- a/Source/Core/DiscIO/VolumeWii.h
+++ b/Source/Core/DiscIO/VolumeWii.h
@@ -58,6 +58,8 @@ public:
   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;