VolumeVerifier: Check whether invalid blocks are unused

5 files changed, 55 insertions, 26 deletions

diff --git a/Source/Core/DiscIO/CompressedBlob.cpp b/Source/Core/DiscIO/CompressedBlob.cpp
index 916773b90a..3dc9f86709 100644
--- a/Source/Core/DiscIO/CompressedBlob.cpp
+++ b/Source/Core/DiscIO/CompressedBlob.cpp
@@ -27,6 +27,7 @@
 #include "DiscIO/Blob.h"
 #include "DiscIO/CompressedBlob.h"
 #include "DiscIO/DiscScrubber.h"
+#include "DiscIO/Volume.h"
 
 namespace DiscIO
 {
@@ -181,9 +182,11 @@ bool CompressFileToBlob(const std::string& infile_path, const std::string& outfi
   }
 
   DiscScrubber disc_scrubber;
+  std::unique_ptr<Volume> volume;
   if (sub_type == 1)
   {
-    if (!disc_scrubber.SetupScrub(infile_path, block_size))
+    volume = CreateVolumeFromFilename(infile_path);
+    if (!volume || !disc_scrubber.SetupScrub(volume.get(), block_size))
     {
       PanicAlertT("\"%s\" failed to be scrubbed. Probably the image is corrupt.",
                   infile_path.c_str());
diff --git a/Source/Core/DiscIO/DiscScrubber.cpp b/Source/Core/DiscIO/DiscScrubber.cpp
index dc3c9bf590..0c944a782a 100644
--- a/Source/Core/DiscIO/DiscScrubber.cpp
+++ b/Source/Core/DiscIO/DiscScrubber.cpp
@@ -28,9 +28,11 @@ constexpr size_t CLUSTER_SIZE = 0x8000;
 DiscScrubber::DiscScrubber() = default;
 DiscScrubber::~DiscScrubber() = default;
 
-bool DiscScrubber::SetupScrub(const std::string& filename, int block_size)
+bool DiscScrubber::SetupScrub(const Volume* disc, int block_size)
 {
-  m_filename = filename;
+  if (!disc)
+    return false;
+  m_disc = disc;
   m_block_size = block_size;
 
   if (CLUSTER_SIZE % m_block_size != 0)
@@ -40,20 +42,13 @@ bool DiscScrubber::SetupScrub(const std::string& filename, int block_size)
     return false;
   }
 
-  m_disc = CreateVolumeFromFilename(filename);
-  if (!m_disc)
-    return false;
-
   m_file_size = m_disc->GetSize();
 
   const size_t num_clusters = static_cast<size_t>(m_file_size / CLUSTER_SIZE);
 
   // Warn if not DVD5 or DVD9 size
   if (num_clusters != 0x23048 && num_clusters != 0x46090)
-  {
-    WARN_LOG(DISCIO, "%s is not a standard sized Wii disc! (%zx blocks)", filename.c_str(),
-             num_clusters);
-  }
+    WARN_LOG(DISCIO, "Not a standard sized Wii disc! (%zx blocks)", num_clusters);
 
   // Table of free blocks
   m_free_table.resize(num_clusters, 1);
@@ -61,8 +56,6 @@ bool DiscScrubber::SetupScrub(const std::string& filename, int block_size)
   // Fill out table of free blocks
   const bool success = ParseDisc();
 
-  // Done with it; need it closed for the next part
-  m_disc.reset();
   m_block_count = 0;
 
   m_is_scrubbing = success;
@@ -72,10 +65,9 @@ bool DiscScrubber::SetupScrub(const std::string& filename, int block_size)
 size_t DiscScrubber::GetNextBlock(File::IOFile& in, u8* buffer)
 {
   const u64 current_offset = m_block_count * m_block_size;
-  const u64 i = current_offset / CLUSTER_SIZE;
 
   size_t read_bytes = 0;
-  if (m_is_scrubbing && m_free_table[i])
+  if (CanBlockBeScrubbed(current_offset))
   {
     DEBUG_LOG(DISCIO, "Freeing 0x%016" PRIx64, current_offset);
     std::fill(buffer, buffer + m_block_size, 0x00);
@@ -92,6 +84,11 @@ size_t DiscScrubber::GetNextBlock(File::IOFile& in, u8* buffer)
   return read_bytes;
 }
 
+bool DiscScrubber::CanBlockBeScrubbed(u64 offset) const
+{
+  return m_is_scrubbing && m_free_table[offset / CLUSTER_SIZE];
+}
+
 void DiscScrubber::MarkAsUsed(u64 offset, u64 size)
 {
   u64 current_offset = offset;
diff --git a/Source/Core/DiscIO/DiscScrubber.h b/Source/Core/DiscIO/DiscScrubber.h
index 29af40c73f..15e0cda7f7 100644
--- a/Source/Core/DiscIO/DiscScrubber.h
+++ b/Source/Core/DiscIO/DiscScrubber.h
@@ -13,7 +13,6 @@
 #pragma once
 
 #include <array>
-#include <memory>
 #include <string>
 #include <vector>
 #include "Common/CommonTypes.h"
@@ -35,8 +34,9 @@ public:
   DiscScrubber();
   ~DiscScrubber();
 
-  bool SetupScrub(const std::string& filename, int block_size);
+  bool SetupScrub(const Volume* disc, int block_size);
   size_t GetNextBlock(File::IOFile& in, u8* buffer);
+  bool CanBlockBeScrubbed(u64 offset) const;
 
 private:
   struct PartitionHeader final
@@ -68,8 +68,7 @@ private:
   bool ParsePartitionData(const Partition& partition, PartitionHeader* header);
   void ParseFileSystemData(u64 partition_data_offset, const FileInfo& directory);
 
-  std::string m_filename;
-  std::unique_ptr<Volume> m_disc;
+  const Volume* m_disc;
 
   std::vector<u8> m_free_table;
   u64 m_file_size = 0;
diff --git a/Source/Core/DiscIO/VolumeVerifier.cpp b/Source/Core/DiscIO/VolumeVerifier.cpp
index ac2f06cc15..74612656fa 100644
--- a/Source/Core/DiscIO/VolumeVerifier.cpp
+++ b/Source/Core/DiscIO/VolumeVerifier.cpp
@@ -29,6 +29,7 @@
 #include "Core/IOS/IOSC.h"
 #include "DiscIO/Blob.h"
 #include "DiscIO/DiscExtractor.h"
+#include "DiscIO/DiscScrubber.h"
 #include "DiscIO/Enums.h"
 #include "DiscIO/Filesystem.h"
 #include "DiscIO/Volume.h"
@@ -629,6 +630,12 @@ void VolumeVerifier::CheckMisc()
 
 void VolumeVerifier::SetUpHashing()
 {
+  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);
+  }
+
   std::sort(m_blocks.begin(), m_blocks.end(),
             [](const BlockToVerify& b1, const BlockToVerify& b2) { return b1.offset < b2.offset; });
 
@@ -698,9 +705,17 @@ void VolumeVerifier::Process()
     if (!m_volume.CheckBlockIntegrity(m_blocks[m_block_index].block_index,
                                       m_blocks[m_block_index].partition))
     {
-      WARN_LOG(DISCIO, "Integrity check failed for block at 0x%" PRIx64,
-               m_blocks[m_block_index].offset);
-      m_block_errors[m_blocks[m_block_index].partition]++;
+      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++;
   }
@@ -750,10 +765,22 @@ void VolumeVerifier::Finish()
     if (pair.second > 0)
     {
       const std::string name = GetPartitionName(m_volume.GetPartitionType(pair.first));
-      AddProblem(Severity::Medium,
-                 StringFromFormat(
-                     GetStringT("Errors were found in %zu blocks in the %s partition.").c_str(),
-                     pair.second, name.c_str()));
+      const std::string text = StringFromFormat(
+          GetStringT("Errors were found in %zu blocks in the %s partition.").c_str(), pair.second,
+          name.c_str());
+      AddProblem(Severity::Medium, text);
+    }
+  }
+
+  for (auto pair : m_unused_block_errors)
+  {
+    if (pair.second > 0)
+    {
+      const std::string name = GetPartitionName(m_volume.GetPartitionType(pair.first));
+      const std::string text = StringFromFormat(
+          GetStringT("Errors were found in %zu unused blocks in the %s partition.").c_str(),
+          pair.second, name.c_str());
+      AddProblem(Severity::Low, text);
     }
   }
 
diff --git a/Source/Core/DiscIO/VolumeVerifier.h b/Source/Core/DiscIO/VolumeVerifier.h
index 4337fdce46..263ef65ff3 100644
--- a/Source/Core/DiscIO/VolumeVerifier.h
+++ b/Source/Core/DiscIO/VolumeVerifier.h
@@ -13,6 +13,7 @@
 #include <mbedtls/sha1.h>
 
 #include "Common/CommonTypes.h"
+#include "DiscIO/DiscScrubber.h"
 #include "DiscIO/Volume.h"
 
 // To be used as follows:
@@ -114,9 +115,11 @@ private:
   mbedtls_md5_context m_md5_context;
   mbedtls_sha1_context m_sha1_context;
 
+  DiscScrubber m_scrubber;
   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;
+  std::map<Partition, size_t> m_unused_block_errors;
 
   bool m_started;
   bool m_done;