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.

1 files changed, 31 insertions, 32 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);