From afcda22da9002f9ce70799b1a80af414b451ecc7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9o=20Lam?= <leo@innovatetechnologi.es>
Date: Sun, 14 May 2017 12:05:43 +0200
Subject: [PATCH] DiscIO: Add GetContent() for reading content from WADs

Direct access to the WAD bytes is required to read contents with proper
padding data (since they can sometimes end up being outside of the
data app section). Allowing the whole buffer to be accessed directly
would be error prone, so this commit adds GetContent() to WiiWAD
for getting raw content data by index.
---
 Source/Core/DiscIO/WiiWad.cpp | 39 +++++++++++++++++++++++++----------
 Source/Core/DiscIO/WiiWad.h   |  8 ++++++-
 2 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/Source/Core/DiscIO/WiiWad.cpp b/Source/Core/DiscIO/WiiWad.cpp
index 2576dac41f..89b687eb95 100644
--- a/Source/Core/DiscIO/WiiWad.cpp
+++ b/Source/Core/DiscIO/WiiWad.cpp
@@ -44,25 +44,24 @@ bool IsWiiWAD(const CBlobBigEndianReader& reader)
 }
 }  // Anonymous namespace
 
-WiiWAD::WiiWAD(const std::string& name)
+WiiWAD::WiiWAD(const std::string& name) : m_reader(CreateBlobReader(name))
 {
-  std::unique_ptr<IBlobReader> reader(CreateBlobReader(name));
-  if (reader == nullptr || File::IsDirectory(name))
+  if (m_reader == nullptr || File::IsDirectory(name))
   {
     m_valid = false;
     return;
   }
 
-  m_valid = ParseWAD(*reader);
+  m_valid = ParseWAD();
 }
 
 WiiWAD::~WiiWAD()
 {
 }
 
-bool WiiWAD::ParseWAD(IBlobReader& reader)
+bool WiiWAD::ParseWAD()
 {
-  CBlobBigEndianReader big_endian_reader(reader);
+  CBlobBigEndianReader big_endian_reader(*m_reader);
 
   if (!IsWiiWAD(big_endian_reader))
     return false;
@@ -86,18 +85,36 @@ bool WiiWAD::ParseWAD(IBlobReader& reader)
     _dbg_assert_msg_(BOOT, reserved == 0x00, "WiiWAD: Reserved must be 0x00");
 
   u32 offset = 0x40;
-  m_certificate_chain = CreateWADEntry(reader, certificate_chain_size, offset);
+  m_certificate_chain = CreateWADEntry(*m_reader, certificate_chain_size, offset);
   offset += Common::AlignUp(certificate_chain_size, 0x40);
-  m_ticket.SetBytes(CreateWADEntry(reader, ticket_size, offset));
+  m_ticket.SetBytes(CreateWADEntry(*m_reader, ticket_size, offset));
   offset += Common::AlignUp(ticket_size, 0x40);
-  m_tmd.SetBytes(CreateWADEntry(reader, tmd_size, offset));
+  m_tmd.SetBytes(CreateWADEntry(*m_reader, tmd_size, offset));
   offset += Common::AlignUp(tmd_size, 0x40);
-  m_data_app = CreateWADEntry(reader, data_app_size, offset);
+  m_data_app_offset = offset;
+  m_data_app = CreateWADEntry(*m_reader, data_app_size, offset);
   offset += Common::AlignUp(data_app_size, 0x40);
-  m_footer = CreateWADEntry(reader, footer_size, offset);
+  m_footer = CreateWADEntry(*m_reader, footer_size, offset);
   offset += Common::AlignUp(footer_size, 0x40);
 
   return true;
 }
 
+std::vector<u8> WiiWAD::GetContent(u16 index) const
+{
+  u64 offset = m_data_app_offset;
+  for (const IOS::ES::Content& content : m_tmd.GetContents())
+  {
+    const u64 aligned_size = Common::AlignUp(content.size, 0x40);
+    if (content.index == index)
+    {
+      std::vector<u8> data(aligned_size);
+      if (!m_reader->Read(offset, aligned_size, data.data()))
+        return {};
+      return data;
+    }
+    offset += aligned_size;
+  }
+  return {};
+}
 }  // namespace DiscIO
diff --git a/Source/Core/DiscIO/WiiWad.h b/Source/Core/DiscIO/WiiWad.h
index 46480fdaa6..fae5bfbf86 100644
--- a/Source/Core/DiscIO/WiiWad.h
+++ b/Source/Core/DiscIO/WiiWad.h
@@ -4,6 +4,7 @@
 
 #pragma once
 
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -27,11 +28,16 @@ public:
   const IOS::ES::TMDReader& GetTMD() const { return m_tmd; }
   const std::vector<u8>& GetDataApp() const { return m_data_app; }
   const std::vector<u8>& GetFooter() const { return m_footer; }
+  std::vector<u8> GetContent(u16 index) const;
+
 private:
-  bool ParseWAD(IBlobReader& reader);
+  bool ParseWAD();
 
   bool m_valid;
 
+  std::unique_ptr<IBlobReader> m_reader;
+
+  u64 m_data_app_offset = 0;
   std::vector<u8> m_certificate_chain;
   IOS::ES::TicketReader m_ticket;
   IOS::ES::TMDReader m_tmd;