DiscIO: Don't keep volume pointer in DiscScrubber

Keeping the pointer creates use-after-free opportunities, and we don't
have much reason to keep it around anyway.

4 files changed, 39 insertions, 41 deletions

diff --git a/Source/Core/DiscIO/DiscScrubber.cpp b/Source/Core/DiscIO/DiscScrubber.cpp
index 22aa06571e..957fde32b9 100644
--- a/Source/Core/DiscIO/DiscScrubber.cpp
+++ b/Source/Core/DiscIO/DiscScrubber.cpp
@@ -24,13 +24,10 @@ namespace DiscIO
 DiscScrubber::DiscScrubber() = default;
 DiscScrubber::~DiscScrubber() = default;
 
-bool DiscScrubber::SetupScrub(const Volume* disc)
+bool DiscScrubber::SetupScrub(const Volume& disc)
 {
-  if (!disc)
-    return false;
-  m_disc = disc;
-
-  m_file_size = m_disc->GetDataSize();
+  m_file_size = disc.GetDataSize();
+  m_has_wii_hashes = disc.HasWiiHashes();
 
   // Round up when diving by CLUSTER_SIZE, otherwise MarkAsUsed might write out of bounds
   const size_t num_clusters = static_cast<size_t>((m_file_size + CLUSTER_SIZE - 1) / CLUSTER_SIZE);
@@ -39,7 +36,7 @@ bool DiscScrubber::SetupScrub(const Volume* disc)
   m_free_table.resize(num_clusters, 1);
 
   // Fill out table of free blocks
-  const bool success = ParseDisc();
+  const bool success = ParseDisc(disc);
 
   m_is_scrubbing = success;
   return success;
@@ -96,38 +93,40 @@ void DiscScrubber::MarkAsUsedE(u64 partition_data_offset, u64 offset, u64 size)
 // Compensate for 0x400 (SHA-1) per 0x8000 (cluster), and round to whole clusters
 u64 DiscScrubber::ToClusterOffset(u64 offset) const
 {
-  if (m_disc->HasWiiHashes())
+  if (m_has_wii_hashes)
     return offset / 0x7c00 * CLUSTER_SIZE;
   else
     return Common::AlignDown(offset, CLUSTER_SIZE);
 }
 
 // Helper functions for reading the BE volume
-bool DiscScrubber::ReadFromVolume(u64 offset, u32& buffer, const Partition& partition)
+bool DiscScrubber::ReadFromVolume(const Volume& disc, u64 offset, u32& buffer,
+                                  const Partition& partition)
 {
-  std::optional<u32> value = m_disc->ReadSwapped<u32>(offset, partition);
+  std::optional<u32> value = disc.ReadSwapped<u32>(offset, partition);
   if (value)
     buffer = *value;
   return value.has_value();
 }
 
-bool DiscScrubber::ReadFromVolume(u64 offset, u64& buffer, const Partition& partition)
+bool DiscScrubber::ReadFromVolume(const Volume& disc, u64 offset, u64& buffer,
+                                  const Partition& partition)
 {
-  std::optional<u64> value = m_disc->ReadSwappedAndShifted(offset, partition);
+  std::optional<u64> value = disc.ReadSwappedAndShifted(offset, partition);
   if (value)
     buffer = *value;
   return value.has_value();
 }
 
-bool DiscScrubber::ParseDisc()
+bool DiscScrubber::ParseDisc(const Volume& disc)
 {
-  if (m_disc->GetPartitions().empty())
-    return ParsePartitionData(PARTITION_NONE);
+  if (disc.GetPartitions().empty())
+    return ParsePartitionData(disc, PARTITION_NONE);
 
   // Mark the header as used - it's mostly 0s anyways
   MarkAsUsed(0, 0x50000);
 
-  for (const DiscIO::Partition& partition : m_disc->GetPartitions())
+  for (const DiscIO::Partition& partition : disc.GetPartitions())
   {
     u32 tmd_size;
     u64 tmd_offset;
@@ -136,15 +135,15 @@ bool DiscScrubber::ParseDisc()
     u64 h3_offset;
     // The H3 size is always 0x18000
 
-    if (!ReadFromVolume(partition.offset + WII_PARTITION_TMD_SIZE_ADDRESS, tmd_size,
-                        PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + WII_PARTITION_TMD_OFFSET_ADDRESS, tmd_offset,
+    if (!ReadFromVolume(disc, partition.offset + WII_PARTITION_TMD_SIZE_ADDRESS, tmd_size,
                         PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + WII_PARTITION_CERT_CHAIN_SIZE_ADDRESS, cert_chain_size,
+        !ReadFromVolume(disc, partition.offset + WII_PARTITION_TMD_OFFSET_ADDRESS, tmd_offset,
                         PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + WII_PARTITION_CERT_CHAIN_OFFSET_ADDRESS,
+        !ReadFromVolume(disc, partition.offset + WII_PARTITION_CERT_CHAIN_SIZE_ADDRESS,
+                        cert_chain_size, PARTITION_NONE) ||
+        !ReadFromVolume(disc, partition.offset + WII_PARTITION_CERT_CHAIN_OFFSET_ADDRESS,
                         cert_chain_offset, PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + WII_PARTITION_H3_OFFSET_ADDRESS, h3_offset,
+        !ReadFromVolume(disc, partition.offset + WII_PARTITION_H3_OFFSET_ADDRESS, h3_offset,
                         PARTITION_NONE))
     {
       return false;
@@ -157,7 +156,7 @@ bool DiscScrubber::ParseDisc()
     MarkAsUsed(partition.offset + h3_offset, WII_PARTITION_H3_SIZE);
 
     // Parse Data! This is where the big gain is
-    if (!ParsePartitionData(partition))
+    if (!ParsePartitionData(disc, partition))
       return false;
   }
 
@@ -165,9 +164,9 @@ bool DiscScrubber::ParseDisc()
 }
 
 // Operations dealing with encrypted space are done here
-bool DiscScrubber::ParsePartitionData(const Partition& partition)
+bool DiscScrubber::ParsePartitionData(const Volume& disc, const Partition& partition)
 {
-  const FileSystem* filesystem = m_disc->GetFileSystem(partition);
+  const FileSystem* filesystem = disc.GetFileSystem(partition);
   if (!filesystem)
   {
     ERROR_LOG_FMT(DISCIO, "Failed to read file system for the partition at {:#x}",
@@ -183,7 +182,7 @@ bool DiscScrubber::ParsePartitionData(const Partition& partition)
   else
   {
     u64 data_offset;
-    if (!ReadFromVolume(partition.offset + 0x2b8, data_offset, PARTITION_NONE))
+    if (!ReadFromVolume(disc, partition.offset + 0x2b8, data_offset, PARTITION_NONE))
       return false;
 
     partition_data_offset = partition.offset + data_offset;
@@ -193,25 +192,25 @@ bool DiscScrubber::ParsePartitionData(const Partition& partition)
   // Header, Header Information, Apploader
   u32 apploader_size;
   u32 apploader_trailer_size;
-  if (!ReadFromVolume(0x2440 + 0x14, apploader_size, partition) ||
-      !ReadFromVolume(0x2440 + 0x18, apploader_trailer_size, partition))
+  if (!ReadFromVolume(disc, 0x2440 + 0x14, apploader_size, partition) ||
+      !ReadFromVolume(disc, 0x2440 + 0x18, apploader_trailer_size, partition))
   {
     return false;
   }
   MarkAsUsedE(partition_data_offset, 0, 0x2440 + apploader_size + apploader_trailer_size);
 
   // DOL
-  const std::optional<u64> dol_offset = GetBootDOLOffset(*m_disc, partition);
+  const std::optional<u64> dol_offset = GetBootDOLOffset(disc, partition);
   if (!dol_offset)
     return false;
-  const std::optional<u64> dol_size = GetBootDOLSize(*m_disc, partition, *dol_offset);
+  const std::optional<u64> dol_size = GetBootDOLSize(disc, partition, *dol_offset);
   if (!dol_size)
     return false;
   MarkAsUsedE(partition_data_offset, *dol_offset, *dol_size);
 
   // FST
-  const std::optional<u64> fst_offset = GetFSTOffset(*m_disc, partition);
-  const std::optional<u64> fst_size = GetFSTSize(*m_disc, partition);
+  const std::optional<u64> fst_offset = GetFSTOffset(disc, partition);
+  const std::optional<u64> fst_size = GetFSTSize(disc, partition);
   if (!fst_offset || !fst_size)
     return false;
   MarkAsUsedE(partition_data_offset, *fst_offset, *fst_size);
diff --git a/Source/Core/DiscIO/DiscScrubber.h b/Source/Core/DiscIO/DiscScrubber.h
index 72d50bf402..87528fba3b 100644
--- a/Source/Core/DiscIO/DiscScrubber.h
+++ b/Source/Core/DiscIO/DiscScrubber.h
@@ -29,7 +29,7 @@ public:
   DiscScrubber();
   ~DiscScrubber();
 
-  bool SetupScrub(const Volume* disc);
+  bool SetupScrub(const Volume& disc);
 
   // Returns true if the specified 32 KiB block only contains unused data
   bool CanBlockBeScrubbed(u64 offset) const;
@@ -40,16 +40,15 @@ private:
   void MarkAsUsed(u64 offset, u64 size);
   void MarkAsUsedE(u64 partition_data_offset, u64 offset, u64 size);
   u64 ToClusterOffset(u64 offset) const;
-  bool ReadFromVolume(u64 offset, u32& buffer, const Partition& partition);
-  bool ReadFromVolume(u64 offset, u64& buffer, const Partition& partition);
-  bool ParseDisc();
-  bool ParsePartitionData(const Partition& partition);
+  bool ReadFromVolume(const Volume& disc, u64 offset, u32& buffer, const Partition& partition);
+  bool ReadFromVolume(const Volume& disc, u64 offset, u64& buffer, const Partition& partition);
+  bool ParseDisc(const Volume& disc);
+  bool ParsePartitionData(const Volume& disc, const Partition& partition);
   void ParseFileSystemData(u64 partition_data_offset, const FileInfo& directory);
 
-  const Volume* m_disc = nullptr;
-
   std::vector<u8> m_free_table;
   u64 m_file_size = 0;
+  bool m_has_wii_hashes = false;
   bool m_is_scrubbing = false;
 };
 
diff --git a/Source/Core/DiscIO/ScrubbedBlob.cpp b/Source/Core/DiscIO/ScrubbedBlob.cpp
index 2f4bc689c4..84a123e4de 100644
--- a/Source/Core/DiscIO/ScrubbedBlob.cpp
+++ b/Source/Core/DiscIO/ScrubbedBlob.cpp
@@ -27,7 +27,7 @@ std::unique_ptr<ScrubbedBlob> ScrubbedBlob::Create(const std::string& path)
     return nullptr;
 
   DiscScrubber scrubber;
-  if (!scrubber.SetupScrub(disc.get()))
+  if (!scrubber.SetupScrub(*disc))
     return nullptr;
 
   std::unique_ptr<BlobReader> blob = CreateBlobReader(path);
diff --git a/Source/Core/DiscIO/VolumeVerifier.cpp b/Source/Core/DiscIO/VolumeVerifier.cpp
index 66e7d00dc4..efbe129a5f 100644
--- a/Source/Core/DiscIO/VolumeVerifier.cpp
+++ b/Source/Core/DiscIO/VolumeVerifier.cpp
@@ -1059,7 +1059,7 @@ void VolumeVerifier::SetUpHashing()
   else if (m_volume.GetVolumeType() == Platform::WiiDisc)
   {
     // Set up a DiscScrubber for checking whether blocks with errors are unused
-    m_scrubber.SetupScrub(&m_volume);
+    m_scrubber.SetupScrub(m_volume);
   }
 
   std::sort(m_groups.begin(), m_groups.end(),