DiscIO/DirectoryBlob: Add ability to have an offset for a partition ContentSource.

2 files changed, 32 insertions, 12 deletions

diff --git a/Source/Core/DiscIO/DirectoryBlob.cpp b/Source/Core/DiscIO/DirectoryBlob.cpp
index 0396e7ecf2..e76fa7ad27 100644
--- a/Source/Core/DiscIO/DirectoryBlob.cpp
+++ b/Source/Core/DiscIO/DirectoryBlob.cpp
@@ -112,16 +112,22 @@ bool DiscContent::Read(u64* offset, u64* length, u8** buffer) const
       const u8* const content_pointer = std::get<const u8*>(m_content_source) + offset_in_content;
       std::copy(content_pointer, content_pointer + bytes_to_read, *buffer);
     }
-    else
+    else if (std::holds_alternative<ContentPartition>(m_content_source))
     {
-      DirectoryBlobReader* blob = std::get<DirectoryBlobReader*>(m_content_source);
+      const auto& content = std::get<ContentPartition>(m_content_source);
+      DirectoryBlobReader* blob = content.m_reader;
       const u64 decrypted_size = m_size * VolumeWii::BLOCK_DATA_SIZE / VolumeWii::BLOCK_TOTAL_SIZE;
-      if (!blob->EncryptPartitionData(offset_in_content, bytes_to_read, *buffer, m_offset,
-                                      decrypted_size))
+      if (!blob->EncryptPartitionData(content.m_offset + offset_in_content, bytes_to_read, *buffer,
+                                      content.m_partition_data_offset, decrypted_size))
       {
         return false;
       }
     }
+    else
+    {
+      PanicAlertFmt("DirectoryBlob: Invalid content source in DiscContent.");
+      return false;
+    }
 
     *length -= bytes_to_read;
     *buffer += bytes_to_read;
@@ -545,9 +551,10 @@ void DirectoryBlobReader::SetPartitions(std::vector<PartitionWithType>&& partiti
     SetPartitionHeader(&partitions[i].partition, partition_address);
 
     const u64 data_size = partitions[i].partition.GetDataSize();
-    m_partitions.emplace(partition_address + PARTITION_DATA_OFFSET,
-                         std::move(partitions[i].partition));
-    m_nonpartition_contents.Add(partition_address + PARTITION_DATA_OFFSET, data_size, this);
+    const u64 partition_data_offset = partition_address + PARTITION_DATA_OFFSET;
+    m_partitions.emplace(partition_data_offset, std::move(partitions[i].partition));
+    m_nonpartition_contents.Add(partition_data_offset, data_size,
+                                ContentPartition{this, 0, partition_data_offset});
     const u64 unaligned_next_partition_address = VolumeWii::EncryptedPartitionOffsetToRawOffset(
         data_size, Partition(partition_address), PARTITION_DATA_OFFSET);
     partition_address = Common::AlignUp(unaligned_next_partition_address, 0x10000ull);
diff --git a/Source/Core/DiscIO/DirectoryBlob.h b/Source/Core/DiscIO/DirectoryBlob.h
index ff7c8e6182..e542705f9f 100644
--- a/Source/Core/DiscIO/DirectoryBlob.h
+++ b/Source/Core/DiscIO/DirectoryBlob.h
@@ -43,11 +43,24 @@ struct ContentFile
   u64 m_offset;
 };
 
-using ContentSource =
-    std::variant<ContentFile,          // File
-                 const u8*,            // Memory
-                 DirectoryBlobReader*  // Partition (which one it is is determined by m_offset)
-                 >;
+// Content chunk that loads data from a DirectoryBlobReader.
+// Intented for representing a partition within a disc.
+struct ContentPartition
+{
+  // The reader to read data from.
+  DirectoryBlobReader* m_reader;
+
+  // Offset from the start of the partition for the first byte represented by this chunk.
+  u64 m_offset;
+
+  // The value passed as partition_data_offset to EncryptPartitionData().
+  u64 m_partition_data_offset;
+};
+
+using ContentSource = std::variant<ContentFile,      // File
+                                   const u8*,        // Memory
+                                   ContentPartition  // Partition
+                                   >;
 
 class DiscContent
 {