From e15b957d531870bd67c2c0257b38004426f8d34b Mon Sep 17 00:00:00 2001
From: JosJuice <josjuice@gmail.com>
Date: Sat, 10 Jun 2017 19:07:12 +0200
Subject: [PATCH] DirectoryBlob: Implement GetDataSize()

---
 Source/Core/DiscIO/DirectoryBlob.cpp | 11 +++++++++--
 Source/Core/DiscIO/DirectoryBlob.h   |  5 +++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/Source/Core/DiscIO/DirectoryBlob.cpp b/Source/Core/DiscIO/DirectoryBlob.cpp
index 0b16c4da7b..4256d6cc9a 100644
--- a/Source/Core/DiscIO/DirectoryBlob.cpp
+++ b/Source/Core/DiscIO/DirectoryBlob.cpp
@@ -29,6 +29,7 @@
 #include "Common/Swap.h"
 #include "Core/Boot/DolReader.h"
 #include "DiscIO/Blob.h"
+#include "DiscIO/VolumeWii.h"
 
 namespace DiscIO
 {
@@ -165,11 +166,16 @@ DirectoryBlobReader::DirectoryBlobReader(const std::string& root_directory)
   if (!m_is_wii)
   {
     m_gamecube_pseudopartition = std::move(game_partition);
+    m_data_size = m_gamecube_pseudopartition.GetDataSize();
   }
   else
   {
     SetNonpartitionDiscHeader(game_partition.GetHeader());
 
+    const u64 unaligned_data_size = VolumeWii::PartitionOffsetToRawOffset(
+        game_partition.GetDataSize(), Partition(GAME_PARTITION_ADDRESS));
+    m_data_size = Common::AlignUp(unaligned_data_size, 0x8000ull);
+
     m_partitions.emplace(GAME_PARTITION_ADDRESS, std::move(game_partition));
 
     SetPartitionTable();
@@ -249,8 +255,7 @@ u64 DirectoryBlobReader::GetRawSize() const
 
 u64 DirectoryBlobReader::GetDataSize() const
 {
-  // Not implemented
-  return 0;
+  return m_data_size;
 }
 
 void DirectoryBlobReader::SetNonpartitionDiscHeader(const std::vector<u8>& partition_header)
@@ -451,6 +456,8 @@ void DirectoryBlobPartition::BuildFST(u64 fst_address)
   Write32((u32)(m_fst_data.size() >> m_address_shift), 0x042c, &m_disk_header);
 
   m_contents.emplace(fst_address, m_fst_data.size(), m_fst_data.data());
+
+  m_data_size = current_data_address;
 }
 
 void DirectoryBlobPartition::WriteEntryData(u32* entry_offset, u8 type, u32 name_offset,
diff --git a/Source/Core/DiscIO/DirectoryBlob.h b/Source/Core/DiscIO/DirectoryBlob.h
index 23048b64b0..c644db4692 100644
--- a/Source/Core/DiscIO/DirectoryBlob.h
+++ b/Source/Core/DiscIO/DirectoryBlob.h
@@ -66,6 +66,7 @@ public:
   DirectoryBlobPartition& operator=(DirectoryBlobPartition&&) = default;
 
   bool IsWii() const { return m_is_wii; }
+  u64 GetDataSize() const { return m_data_size; }
   const std::vector<u8>& GetHeader() const { return m_disk_header; }
   const std::set<DiscContent>& GetContents() const { return m_contents; }
 private:
@@ -95,6 +96,8 @@ private:
   bool m_is_wii = false;
   // GameCube has no shift, Wii has 2 bit shift
   u32 m_address_shift = 0;
+
+  u64 m_data_size;
 };
 
 class DirectoryBlobReader : public BlobReader
@@ -148,6 +151,8 @@ private:
   } m_tmd_header;
   static_assert(sizeof(TMDHeader) == 8, "Wrong size for TMDHeader");
 #pragma pack(pop)
+
+  u64 m_data_size;
 };
 
 }  // namespace