Merge pull request #5871 from JosJuice/read-swapped-and-shifted

DiscIO: Add a Volume::ReadSwappedAndShifted function

7 files changed, 42 insertions, 39 deletions

diff --git a/Source/Core/DiscIO/DiscExtractor.cpp b/Source/Core/DiscIO/DiscExtractor.cpp
index 1220aadd50..4ebaf249fd 100644
--- a/Source/Core/DiscIO/DiscExtractor.cpp
+++ b/Source/Core/DiscIO/DiscExtractor.cpp
@@ -160,13 +160,13 @@ bool ExportTMD(const Volume& volume, const Partition& partition, const std::stri
   if (volume.GetVolumeType() != Platform::WII_DISC)
     return false;
 
-  std::optional<u32> size = volume.ReadSwapped<u32>(partition.offset + 0x2a4, PARTITION_NONE);
-  std::optional<u32> offset = volume.ReadSwapped<u32>(partition.offset + 0x2a8, PARTITION_NONE);
+  const std::optional<u32> size = volume.ReadSwapped<u32>(partition.offset + 0x2a4, PARTITION_NONE);
+  const std::optional<u64> offset =
+      volume.ReadSwappedAndShifted(partition.offset + 0x2a8, PARTITION_NONE);
   if (!size || !offset)
     return false;
 
-  const u64 actual_offset = partition.offset + (static_cast<u64>(*offset) << 2);
-  return ExportData(volume, PARTITION_NONE, actual_offset, *size, export_filename);
+  return ExportData(volume, PARTITION_NONE, *offset, *size, export_filename);
 }
 
 bool ExportCertificateChain(const Volume& volume, const Partition& partition,
@@ -175,13 +175,13 @@ bool ExportCertificateChain(const Volume& volume, const Partition& partition,
   if (volume.GetVolumeType() != Platform::WII_DISC)
     return false;
 
-  std::optional<u32> size = volume.ReadSwapped<u32>(partition.offset + 0x2ac, PARTITION_NONE);
-  std::optional<u32> offset = volume.ReadSwapped<u32>(partition.offset + 0x2b0, PARTITION_NONE);
+  const std::optional<u32> size = volume.ReadSwapped<u32>(partition.offset + 0x2ac, PARTITION_NONE);
+  const std::optional<u64> offset =
+      volume.ReadSwappedAndShifted(partition.offset + 0x2b0, PARTITION_NONE);
   if (!size || !offset)
     return false;
 
-  const u64 actual_offset = partition.offset + (static_cast<u64>(*offset) << 2);
-  return ExportData(volume, PARTITION_NONE, actual_offset, *size, export_filename);
+  return ExportData(volume, PARTITION_NONE, *offset, *size, export_filename);
 }
 
 bool ExportH3Hashes(const Volume& volume, const Partition& partition,
@@ -190,12 +190,12 @@ bool ExportH3Hashes(const Volume& volume, const Partition& partition,
   if (volume.GetVolumeType() != Platform::WII_DISC)
     return false;
 
-  std::optional<u32> offset = volume.ReadSwapped<u32>(partition.offset + 0x2b4, PARTITION_NONE);
+  const std::optional<u64> offset =
+      volume.ReadSwappedAndShifted(partition.offset + 0x2b4, PARTITION_NONE);
   if (!offset)
     return false;
 
-  const u64 actual_offset = partition.offset + (static_cast<u64>(*offset) << 2);
-  return ExportData(volume, PARTITION_NONE, actual_offset, 0x18000, export_filename);
+  return ExportData(volume, PARTITION_NONE, *offset, 0x18000, export_filename);
 }
 
 bool ExportHeader(const Volume& volume, const Partition& partition,
@@ -239,9 +239,7 @@ std::optional<u64> GetBootDOLOffset(const Volume& volume, const Partition& parti
   if (!IsDisc(volume_type))
     return {};
 
-  const std::optional<u32> offset = volume.ReadSwapped<u32>(0x420, partition);
-  const u8 offset_shift = volume_type == Platform::WII_DISC ? 2 : 0;
-  return offset ? static_cast<u64>(*offset) << offset_shift : std::optional<u64>();
+  return volume.ReadSwappedAndShifted(0x420, partition);
 }
 
 std::optional<u32> GetBootDOLSize(const Volume& volume, const Partition& partition, u64 dol_offset)
@@ -295,9 +293,7 @@ std::optional<u64> GetFSTOffset(const Volume& volume, const Partition& partition
   if (!IsDisc(volume_type))
     return {};
 
-  const std::optional<u32> offset = volume.ReadSwapped<u32>(0x424, partition);
-  const u8 offset_shift = volume_type == Platform::WII_DISC ? 2 : 0;
-  return offset ? static_cast<u64>(*offset) << offset_shift : std::optional<u64>();
+  return volume.ReadSwappedAndShifted(0x424, partition);
 }
 
 std::optional<u64> GetFSTSize(const Volume& volume, const Partition& partition)
@@ -306,9 +302,7 @@ std::optional<u64> GetFSTSize(const Volume& volume, const Partition& partition)
   if (!IsDisc(volume_type))
     return {};
 
-  const std::optional<u32> size = volume.ReadSwapped<u32>(0x428, partition);
-  const u8 offset_shift = volume_type == Platform::WII_DISC ? 2 : 0;
-  return size ? static_cast<u64>(*size) << offset_shift : std::optional<u64>();
+  return volume.ReadSwappedAndShifted(0x428, partition);
 }
 
 bool ExportFST(const Volume& volume, const Partition& partition, const std::string& export_filename)
diff --git a/Source/Core/DiscIO/DiscScrubber.cpp b/Source/Core/DiscIO/DiscScrubber.cpp
index 52d907be69..b2fc397eea 100644
--- a/Source/Core/DiscIO/DiscScrubber.cpp
+++ b/Source/Core/DiscIO/DiscScrubber.cpp
@@ -136,9 +136,9 @@ bool DiscScrubber::ReadFromVolume(u64 offset, u32& buffer, const Partition& part
 
 bool DiscScrubber::ReadFromVolume(u64 offset, u64& buffer, const Partition& partition)
 {
-  std::optional<u32> value = m_disc->ReadSwapped<u32>(offset, partition);
+  std::optional<u64> value = m_disc->ReadSwappedAndShifted(offset, partition);
   if (value)
-    buffer = static_cast<u64>(*value) << 2;
+    buffer = *value;
   return value.has_value();
 }
 
diff --git a/Source/Core/DiscIO/FileSystemGCWii.cpp b/Source/Core/DiscIO/FileSystemGCWii.cpp
index 9a7e24c288..c7cd7115cc 100644
--- a/Source/Core/DiscIO/FileSystemGCWii.cpp
+++ b/Source/Core/DiscIO/FileSystemGCWii.cpp
@@ -185,15 +185,16 @@ bool FileInfoGCWii::IsValid(u64 fst_size, const FileInfoGCWii& parent_directory)
 }
 
 FileSystemGCWii::FileSystemGCWii(const Volume* volume, const Partition& partition)
-    : FileSystem(volume, partition), m_valid(false), m_offset_shift(0), m_root(nullptr, 0, 0, 0)
+    : FileSystem(volume, partition), m_valid(false), m_root(nullptr, 0, 0, 0)
 {
+  u8 offset_shift;
   // Check if this is a GameCube or Wii disc
   if (m_volume->ReadSwapped<u32>(0x18, m_partition) == u32(0x5D1C9EA3))
-    m_offset_shift = 2;  // Wii file system
+    offset_shift = 2;  // Wii file system
   else if (m_volume->ReadSwapped<u32>(0x1c, m_partition) == u32(0xC2339F3D))
-    m_offset_shift = 0;  // GameCube file system
+    offset_shift = 0;  // GameCube file system
   else
-    return;
+    return;  // Invalid partition (maybe someone removed its data but not its partition table entry)
 
   const std::optional<u64> fst_offset = GetFSTOffset(*m_volume, m_partition);
   const std::optional<u64> fst_size = GetFSTSize(*m_volume, m_partition);
@@ -226,7 +227,7 @@ FileSystemGCWii::FileSystemGCWii(const Volume* volume, const Partition& partitio
   }
 
   // Create the root object
-  m_root = FileInfoGCWii(m_file_system_table.data(), m_offset_shift);
+  m_root = FileInfoGCWii(m_file_system_table.data(), offset_shift);
   if (!m_root.IsDirectory())
   {
     ERROR_LOG(DISCIO, "File system root is not a directory");
diff --git a/Source/Core/DiscIO/FileSystemGCWii.h b/Source/Core/DiscIO/FileSystemGCWii.h
index 8e393642bf..847bf2e639 100644
--- a/Source/Core/DiscIO/FileSystemGCWii.h
+++ b/Source/Core/DiscIO/FileSystemGCWii.h
@@ -95,7 +95,6 @@ public:
 
 private:
   bool m_valid;
-  u32 m_offset_shift;
   std::vector<u8> m_file_system_table;
   FileInfoGCWii m_root;
   // Maps the end offset of files to FST indexes
diff --git a/Source/Core/DiscIO/Volume.h b/Source/Core/DiscIO/Volume.h
index 979a63a23f..b4d2965f5a 100644
--- a/Source/Core/DiscIO/Volume.h
+++ b/Source/Core/DiscIO/Volume.h
@@ -51,6 +51,12 @@ public:
       return {};
     return Common::FromBigEndian(temp);
   }
+  std::optional<u64> ReadSwappedAndShifted(u64 offset, const Partition& partition) const
+  {
+    const std::optional<u32> temp = ReadSwapped<u32>(offset, partition);
+    return temp ? static_cast<u64>(*temp) << GetOffsetShift() : std::optional<u64>();
+  }
+
   virtual std::vector<Partition> GetPartitions() const { return {}; }
   virtual Partition GetGamePartition() const { return PARTITION_NONE; }
   virtual std::optional<u32> GetPartitionType(const Partition& partition) const { return {}; }
@@ -108,6 +114,7 @@ protected:
       return CP1252ToUTF8(string);
   }
 
+  virtual u32 GetOffsetShift() const { return 0; }
   static std::map<Language, std::string> ReadWiiNames(const std::vector<u8>& data);
 
   static const size_t NUMBER_OF_LANGUAGES = 10;
diff --git a/Source/Core/DiscIO/VolumeWii.cpp b/Source/Core/DiscIO/VolumeWii.cpp
index 4e8bba6b91..bce96822e7 100644
--- a/Source/Core/DiscIO/VolumeWii.cpp
+++ b/Source/Core/DiscIO/VolumeWii.cpp
@@ -52,22 +52,22 @@ VolumeWii::VolumeWii(std::unique_ptr<BlobReader> reader)
     if (!number_of_partitions)
       continue;
 
-    std::optional<u32> read_buffer =
-        m_pReader->ReadSwapped<u32>(0x40000 + (partition_group * 8) + 4);
-    if (!read_buffer)
+    const std::optional<u64> partition_table_offset =
+        ReadSwappedAndShifted(0x40000 + (partition_group * 8) + 4, PARTITION_NONE);
+    if (!partition_table_offset)
       continue;
-    const u64 partition_table_offset = static_cast<u64>(*read_buffer) << 2;
 
     for (u32 i = 0; i < number_of_partitions; i++)
     {
-      read_buffer = m_pReader->ReadSwapped<u32>(partition_table_offset + (i * 8));
-      if (!read_buffer)
+      const std::optional<u64> partition_offset =
+          ReadSwappedAndShifted(*partition_table_offset + (i * 8), PARTITION_NONE);
+      if (!partition_offset)
         continue;
-      const u64 partition_offset = static_cast<u64>(*read_buffer) << 2;
-      const Partition partition(partition_offset);
+
+      const Partition partition(*partition_offset);
 
       const std::optional<u32> partition_type =
-          m_pReader->ReadSwapped<u32>(partition_table_offset + (i * 8) + 4);
+          m_pReader->ReadSwapped<u32>(*partition_table_offset + (i * 8) + 4);
       if (!partition_type)
         continue;
 
@@ -84,10 +84,10 @@ VolumeWii::VolumeWii(std::unique_ptr<BlobReader> reader)
 
       auto get_tmd = [this, partition]() -> IOS::ES::TMDReader {
         const std::optional<u32> tmd_size = m_pReader->ReadSwapped<u32>(partition.offset + 0x2a4);
-        std::optional<u32> tmd_address = m_pReader->ReadSwapped<u32>(partition.offset + 0x2a8);
+        const std::optional<u64> tmd_address =
+            ReadSwappedAndShifted(partition.offset + 0x2a8, PARTITION_NONE);
         if (!tmd_size || !tmd_address)
           return INVALID_TMD;
-        *tmd_address <<= 2;
         if (!IOS::ES::IsValidTMDSize(*tmd_size))
         {
           // This check is normally done by ES in ES_DiVerify, but that would happen too late
diff --git a/Source/Core/DiscIO/VolumeWii.h b/Source/Core/DiscIO/VolumeWii.h
index 7ada5805c3..e702e316ae 100644
--- a/Source/Core/DiscIO/VolumeWii.h
+++ b/Source/Core/DiscIO/VolumeWii.h
@@ -64,6 +64,8 @@ public:
   static constexpr unsigned int BLOCK_DATA_SIZE = 0x7C00;
   static constexpr unsigned int BLOCK_TOTAL_SIZE = BLOCK_HEADER_SIZE + BLOCK_DATA_SIZE;
 
+protected:
+  u32 GetOffsetShift() const override { return 2; }
 private:
   struct PartitionDetails
   {