1 files changed, 59 insertions, 39 deletions

diff --git a/Source/Core/DiscIO/DiscScrubber.cpp b/Source/Core/DiscIO/DiscScrubber.cpp
index 706a2cd00d..06c444c84d 100644
--- a/Source/Core/DiscIO/DiscScrubber.cpp
+++ b/Source/Core/DiscIO/DiscScrubber.cpp
@@ -72,20 +72,27 @@ void DiscScrubber::MarkAsUsed(u64 offset, u64 size)
 
 void DiscScrubber::MarkAsUsedE(u64 partition_data_offset, u64 offset, u64 size)
 {
-  u64 first_cluster_start = ToClusterOffset(offset) + partition_data_offset;
-
-  u64 last_cluster_end;
-  if (size == 0)
+  if (partition_data_offset == 0)
   {
-    // Without this special case, a size of 0 can be rounded to 1 cluster instead of 0
-    last_cluster_end = first_cluster_start;
+    MarkAsUsed(offset, size);
   }
   else
   {
-    last_cluster_end = ToClusterOffset(offset + size - 1) + CLUSTER_SIZE + partition_data_offset;
-  }
+    u64 first_cluster_start = ToClusterOffset(offset) + partition_data_offset;
+
+    u64 last_cluster_end;
+    if (size == 0)
+    {
+      // Without this special case, a size of 0 can be rounded to 1 cluster instead of 0
+      last_cluster_end = first_cluster_start;
+    }
+    else
+    {
+      last_cluster_end = ToClusterOffset(offset + size - 1) + CLUSTER_SIZE + partition_data_offset;
+    }
 
-  MarkAsUsed(first_cluster_start, last_cluster_end - first_cluster_start);
+    MarkAsUsed(first_cluster_start, last_cluster_end - first_cluster_start);
+  }
 }
 
 // Compensate for 0x400 (SHA-1) per 0x8000 (cluster), and round to whole clusters
@@ -116,35 +123,38 @@ bool DiscScrubber::ReadFromVolume(u64 offset, u64& buffer, const Partition& part
 
 bool DiscScrubber::ParseDisc()
 {
+  if (m_disc->GetPartitions().empty())
+    return ParsePartitionData(PARTITION_NONE);
+
   // Mark the header as used - it's mostly 0s anyways
   MarkAsUsed(0, 0x50000);
 
   for (const DiscIO::Partition& partition : m_disc->GetPartitions())
   {
-    PartitionHeader header;
-
-    if (!ReadFromVolume(partition.offset + 0x2a4, header.tmd_size, PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + 0x2a8, header.tmd_offset, PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + 0x2ac, header.cert_chain_size, PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + 0x2b0, header.cert_chain_offset, PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + 0x2b4, header.h3_offset, PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + 0x2b8, header.data_offset, PARTITION_NONE) ||
-        !ReadFromVolume(partition.offset + 0x2bc, header.data_size, PARTITION_NONE))
+    u32 tmd_size;
+    u64 tmd_offset;
+    u32 cert_chain_size;
+    u64 cert_chain_offset;
+    u64 h3_offset;
+    // The H3 size is always 0x18000
+
+    if (!ReadFromVolume(partition.offset + 0x2a4, tmd_size, PARTITION_NONE) ||
+        !ReadFromVolume(partition.offset + 0x2a8, tmd_offset, PARTITION_NONE) ||
+        !ReadFromVolume(partition.offset + 0x2ac, cert_chain_size, PARTITION_NONE) ||
+        !ReadFromVolume(partition.offset + 0x2b0, cert_chain_offset, PARTITION_NONE) ||
+        !ReadFromVolume(partition.offset + 0x2b4, h3_offset, PARTITION_NONE))
     {
       return false;
     }
 
     MarkAsUsed(partition.offset, 0x2c0);
 
-    MarkAsUsed(partition.offset + header.tmd_offset, header.tmd_size);
-    MarkAsUsed(partition.offset + header.cert_chain_offset, header.cert_chain_size);
-    MarkAsUsed(partition.offset + header.h3_offset, 0x18000);
-    // This would mark the whole (encrypted) data area
-    // we need to parse FST and other crap to find what's free within it!
-    // MarkAsUsed(partition.offset + header.data_offset, header.data_size);
+    MarkAsUsed(partition.offset + tmd_offset, tmd_size);
+    MarkAsUsed(partition.offset + cert_chain_offset, cert_chain_size);
+    MarkAsUsed(partition.offset + h3_offset, 0x18000);
 
     // Parse Data! This is where the big gain is
-    if (!ParsePartitionData(partition, &header))
+    if (!ParsePartitionData(partition))
       return false;
   }
 
@@ -152,7 +162,7 @@ bool DiscScrubber::ParseDisc()
 }
 
 // Operations dealing with encrypted space are done here
-bool DiscScrubber::ParsePartitionData(const Partition& partition, PartitionHeader* header)
+bool DiscScrubber::ParsePartitionData(const Partition& partition)
 {
   const FileSystem* filesystem = m_disc->GetFileSystem(partition);
   if (!filesystem)
@@ -162,17 +172,30 @@ bool DiscScrubber::ParsePartitionData(const Partition& partition, PartitionHeade
     return false;
   }
 
-  const u64 partition_data_offset = partition.offset + header->data_offset;
+  u64 partition_data_offset;
+  if (partition == PARTITION_NONE)
+  {
+    partition_data_offset = 0;
+  }
+  else
+  {
+    u64 data_offset;
+    if (!ReadFromVolume(partition.offset + 0x2b8, data_offset, PARTITION_NONE))
+      return false;
+
+    partition_data_offset = partition.offset + data_offset;
+  }
 
   // Mark things as used which are not in the filesystem
   // Header, Header Information, Apploader
-  if (!ReadFromVolume(0x2440 + 0x14, header->apploader_size, partition) ||
-      !ReadFromVolume(0x2440 + 0x18, header->apploader_size, partition))
+  u32 apploader_size;
+  u32 apploader_trailer_size;
+  if (!ReadFromVolume(0x2440 + 0x14, apploader_size, partition) ||
+      !ReadFromVolume(0x2440 + 0x18, apploader_trailer_size, partition))
   {
     return false;
   }
-  MarkAsUsedE(partition_data_offset, 0,
-              0x2440 + header->apploader_size + header->apploader_trailer_size);
+  MarkAsUsedE(partition_data_offset, 0, 0x2440 + apploader_size + apploader_trailer_size);
 
   // DOL
   const std::optional<u64> dol_offset = GetBootDOLOffset(*m_disc, partition);
@@ -181,17 +204,14 @@ bool DiscScrubber::ParsePartitionData(const Partition& partition, PartitionHeade
   const std::optional<u64> dol_size = GetBootDOLSize(*m_disc, partition, *dol_offset);
   if (!dol_size)
     return false;
-  header->dol_offset = *dol_offset;
-  header->dol_size = *dol_size;
-  MarkAsUsedE(partition_data_offset, header->dol_offset, header->dol_size);
+  MarkAsUsedE(partition_data_offset, *dol_offset, *dol_size);
 
   // FST
-  if (!ReadFromVolume(0x424, header->fst_offset, partition) ||
-      !ReadFromVolume(0x428, header->fst_size, partition))
-  {
+  const std::optional<u64> fst_offset = GetFSTOffset(*m_disc, partition);
+  const std::optional<u64> fst_size = GetFSTSize(*m_disc, partition);
+  if (!fst_offset || !fst_size)
     return false;
-  }
-  MarkAsUsedE(partition_data_offset, header->fst_offset, header->fst_size);
+  MarkAsUsedE(partition_data_offset, *fst_offset, *fst_size);
 
   // Go through the filesystem and mark entries as used
   ParseFileSystemData(partition_data_offset, filesystem->GetRoot());