VolumeVerifier: Verify WAD contents

Source/Core/DiscIO/Volume.h

@@ -73,6 +73,7 @@ public:
   {
     return INVALID_CERT_CHAIN;
   }
+  virtual std::vector<u64> GetContentOffsets() const { return {}; }
   // Returns a non-owning pointer. Returns nullptr if the file system couldn't be read.
   virtual const FileSystem* GetFileSystem(const Partition& partition) const = 0;
   virtual u64 PartitionOffsetToRawOffset(u64 offset, const Partition& partition) const

Source/Core/DiscIO/VolumeVerifier.cpp

@@ -11,6 +11,7 @@
 #include <string>
 #include <unordered_set>
 
+#include <mbedtls/aes.h>
 #include <mbedtls/md5.h>
 #include <mbedtls/sha1.h>
 #include <zlib.h>
@@ -630,7 +631,11 @@ void VolumeVerifier::CheckMisc()
 
 void VolumeVerifier::SetUpHashing()
 {
-  if (m_volume.GetVolumeType() == Platform::WiiDisc)
+  if (m_volume.GetVolumeType() == Platform::WiiWAD)
+  {
+    m_content_offsets = m_volume.GetContentOffsets();
+  }
+  else if (m_volume.GetVolumeType() == Platform::WiiDisc)
   {
     // Set up a DiscScrubber for checking whether blocks with errors are unused
     m_scrubber.SetupScrub(&m_volume, VolumeWii::BLOCK_TOTAL_SIZE);
@@ -663,14 +668,28 @@ void VolumeVerifier::Process()
   if (m_progress == m_max_progress)
     return;
 
+  IOS::ES::Content content;
+  bool content_read = false;
   u64 bytes_to_read = BLOCK_SIZE;
-  if (m_block_index < m_blocks.size() && m_blocks[m_block_index].offset == m_progress)
+  if (m_content_index < m_content_offsets.size() &&
+      m_content_offsets[m_content_index] == m_progress)
+  {
+    m_volume.GetTMD(PARTITION_NONE).GetContent(m_content_index, &content);
+    bytes_to_read = Common::AlignUp(content.size, 0x40);
+    content_read = true;
+  }
+  else if (m_content_index < m_content_offsets.size() &&
+           m_content_offsets[m_content_index] > m_progress)
+  {
+    bytes_to_read = std::min(bytes_to_read, m_content_offsets[m_content_index] - m_progress);
+  }
+  else if (m_block_index < m_blocks.size() && m_blocks[m_block_index].offset == m_progress)
   {
     bytes_to_read = VolumeWii::BLOCK_TOTAL_SIZE;
   }
-  else if (m_block_index + 1 < m_blocks.size() && m_blocks[m_block_index + 1].offset > m_progress)
+  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 + 1].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);
 
@@ -700,6 +719,17 @@ void VolumeVerifier::Process()
 
   m_progress += bytes_to_read;
 
+  if (content_read)
+  {
+    if (!CheckContentIntegrity(content))
+    {
+      AddProblem(Severity::High,
+                 StringFromFormat(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_volume.CheckBlockIntegrity(m_blocks[m_block_index].block_index,
@@ -721,6 +751,30 @@ void VolumeVerifier::Process()
   }
 }
 
+bool VolumeVerifier::CheckContentIntegrity(const IOS::ES::Content& content)
+{
+  const u64 padded_size = Common::AlignUp(content.size, 0x40);
+  std::vector<u8> encrypted_data(padded_size);
+  m_volume.Read(m_content_offsets[m_content_index], padded_size, encrypted_data.data(),
+                PARTITION_NONE);
+
+  mbedtls_aes_context context;
+  const std::array<u8, 16> key = m_volume.GetTicket(PARTITION_NONE).GetTitleKey();
+  mbedtls_aes_setkey_dec(&context, key.data(), 128);
+
+  std::array<u8, 16> iv{};
+  iv[0] = static_cast<u8>(content.index >> 8);
+  iv[1] = static_cast<u8>(content.index & 0xFF);
+
+  std::vector<u8> decrypted_data(padded_size);
+  mbedtls_aes_crypt_cbc(&context, MBEDTLS_AES_DECRYPT, padded_size, iv.data(),
+                        encrypted_data.data(), decrypted_data.data());
+
+  std::array<u8, 20> sha1;
+  mbedtls_sha1(decrypted_data.data(), content.size, sha1.data());
+  return sha1 == content.sha1;
+}
+
 u64 VolumeVerifier::GetBytesProcessed() const
 {
   return m_progress;

Source/Core/DiscIO/VolumeVerifier.h

@@ -31,6 +31,7 @@
 
 namespace IOS::ES
 {
+struct Content;
 class SignedBlobReader;
 }
 
@@ -100,6 +101,7 @@ private:
   u64 GetBiggestUsedOffset(const FileInfo& file_info) const;
   void CheckMisc();
   void SetUpHashing();
+  bool CheckContentIntegrity(const IOS::ES::Content& content);
 
   void AddProblem(Severity severity, const std::string& text);
 
@@ -116,6 +118,8 @@ private:
   mbedtls_sha1_context m_sha1_context;
 
   DiscScrubber m_scrubber;
+  std::vector<u64> m_content_offsets;
+  u16 m_content_index = 0;
   std::vector<BlockToVerify> m_blocks;
   size_t m_block_index = 0;  // Index in m_blocks, not index in a specific partition
   std::map<Partition, size_t> m_block_errors;

Source/Core/DiscIO/VolumeWad.cpp

@@ -40,8 +40,8 @@ VolumeWAD::VolumeWAD(std::unique_ptr<BlobReader> reader) : m_reader(std::move(re
   m_cert_chain_offset = Common::AlignUp(m_hdr_size, 0x40);
   m_ticket_offset = m_cert_chain_offset + Common::AlignUp(m_cert_chain_size, 0x40);
   m_tmd_offset = m_ticket_offset + Common::AlignUp(m_ticket_size, 0x40);
-  m_opening_bnr_offset =
-      m_tmd_offset + Common::AlignUp(m_tmd_size, 0x40) + Common::AlignUp(m_data_size, 0x40);
+  m_data_offset = m_tmd_offset + Common::AlignUp(m_tmd_size, 0x40);
+  m_opening_bnr_offset = m_data_offset + Common::AlignUp(m_data_size, 0x40);
 
   std::vector<u8> ticket_buffer(m_ticket_size);
   Read(m_ticket_offset, m_ticket_size, ticket_buffer.data());
@@ -118,6 +118,21 @@ const std::vector<u8>& VolumeWAD::GetCertificateChain(const Partition& partition
   return m_cert_chain;
 }
 
+std::vector<u64> VolumeWAD::GetContentOffsets() const
+{
+  const std::vector<IOS::ES::Content> contents = m_tmd.GetContents();
+  std::vector<u64> content_offsets;
+  content_offsets.reserve(contents.size());
+  u64 offset = m_data_offset;
+  for (const IOS::ES::Content& content : contents)
+  {
+    content_offsets.emplace_back(offset);
+    offset += Common::AlignUp(content.size, 0x40);
+  }
+
+  return content_offsets;
+}
+
 std::string VolumeWAD::GetGameID(const Partition& partition) const
 {
   return m_tmd.GetGameID();

Source/Core/DiscIO/VolumeWad.h

@@ -38,6 +38,7 @@ public:
   const IOS::ES::TMDReader& GetTMD(const Partition& partition = PARTITION_NONE) const override;
   const std::vector<u8>&
   GetCertificateChain(const Partition& partition = PARTITION_NONE) const override;
+  std::vector<u64> GetContentOffsets() const override;
   std::string GetGameID(const Partition& partition = PARTITION_NONE) const override;
   std::string GetGameTDBID(const Partition& partition = PARTITION_NONE) const override;
   std::string GetMakerID(const Partition& partition = PARTITION_NONE) const override;
@@ -69,6 +70,7 @@ private:
   u32 m_cert_chain_offset = 0;
   u32 m_ticket_offset = 0;
   u32 m_tmd_offset = 0;
+  u32 m_data_offset = 0;
   u32 m_opening_bnr_offset = 0;
   u32 m_hdr_size = 0;
   u32 m_cert_chain_size = 0;