From 47067f661acc0b39a6ceb54227d52c1b0ced5600 Mon Sep 17 00:00:00 2001
From: JosJuice <josjuice@gmail.com>
Date: Mon, 27 Jan 2020 14:26:12 +0100
Subject: [PATCH] WIA: Properly check for overlapping data

---
 Source/Core/DiscIO/WIABlob.cpp | 43 ++++++++++++++++++++++++++++++----
 Source/Core/DiscIO/WIABlob.h   |  1 +
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/Source/Core/DiscIO/WIABlob.cpp b/Source/Core/DiscIO/WIABlob.cpp
index d92aa8f0f1..5d71f24a16 100644
--- a/Source/Core/DiscIO/WIABlob.cpp
+++ b/Source/Core/DiscIO/WIABlob.cpp
@@ -142,10 +142,7 @@ bool WIAFileReader::Initialize(const std::string& path)
 
     if (non_empty_entries > 1)
     {
-      const u32 first_end =
-          Common::swap32(entries[0].first_sector) + Common::swap32(entries[0].number_of_sectors);
-      const u32 second_start = Common::swap32(entries[1].first_sector);
-      if (first_end > second_start)
+      if (Common::swap32(entries[0].first_sector) > Common::swap32(entries[1].first_sector))
         return false;
     }
   }
@@ -175,9 +172,47 @@ bool WIAFileReader::Initialize(const std::string& path)
   if (!group_entries.ReadAll(&m_group_entries))
     return false;
 
+  if (HasDataOverlap())
+    return false;
+
   return true;
 }
 
+bool WIAFileReader::HasDataOverlap() const
+{
+  for (size_t i = 0; i < m_partition_entries.size(); ++i)
+  {
+    const std::array<PartitionDataEntry, 2>& entries = m_partition_entries[i].data_entries;
+    for (size_t j = 0; j < entries.size(); ++j)
+    {
+      if (Common::swap32(entries[j].number_of_sectors) == 0)
+        continue;
+
+      const u64 data_offset = Common::swap32(entries[j].first_sector) * VolumeWii::BLOCK_TOTAL_SIZE;
+      const auto it = m_data_entries.upper_bound(data_offset);
+      if (it == m_data_entries.end())
+        return true;  // Not an overlap, but an error nonetheless
+      if (!it->second.is_partition || it->second.index != i || it->second.partition_data_index != j)
+        return true;  // Overlap
+    }
+  }
+
+  for (size_t i = 0; i < m_raw_data_entries.size(); ++i)
+  {
+    if (Common::swap64(m_raw_data_entries[i].data_size) == 0)
+      continue;
+
+    const u64 data_offset = Common::swap64(m_raw_data_entries[i].data_offset);
+    const auto it = m_data_entries.upper_bound(data_offset);
+    if (it == m_data_entries.end())
+      return true;  // Not an overlap, but an error nonetheless
+    if (it->second.is_partition || it->second.index != i)
+      return true;  // Overlap
+  }
+
+  return false;
+}
+
 std::unique_ptr<WIAFileReader> WIAFileReader::Create(File::IOFile file, const std::string& path)
 {
   std::unique_ptr<WIAFileReader> blob(new WIAFileReader(std::move(file), path));
diff --git a/Source/Core/DiscIO/WIABlob.h b/Source/Core/DiscIO/WIABlob.h
index 2360c7acf4..5bd0770719 100644
--- a/Source/Core/DiscIO/WIABlob.h
+++ b/Source/Core/DiscIO/WIABlob.h
@@ -261,6 +261,7 @@ private:
 
   explicit WIAFileReader(File::IOFile file, const std::string& path);
   bool Initialize(const std::string& path);
+  bool HasDataOverlap() const;
 
   bool ReadFromGroups(u64* offset, u64* size, u8** out_ptr, u64 chunk_size, u32 sector_size,
                       u64 data_offset, u64 data_size, u32 group_index, u32 number_of_groups,