Merge pull request #5683 from JosJuice/volume-wii-defer

VolumeWii: Defer loading tickets, TMDs and keys until when needed

Source/Core/Common/Common.vcxproj

@@ -126,6 +126,7 @@
     <ClInclude Include="HttpRequest.h" />
     <ClInclude Include="IniFile.h" />
     <ClInclude Include="JitRegister.h" />
+    <ClInclude Include="Lazy.h" />
     <ClInclude Include="LdrWatcher.h" />
     <ClInclude Include="LinearDiskCache.h" />
     <ClInclude Include="MathUtil.h" />

Source/Core/Common/Common.vcxproj.filters

@@ -253,6 +253,7 @@
       <Filter>GL\GLExtensions</Filter>
     </ClInclude>
     <ClInclude Include="File.h" />
+    <ClInclude Include="Lazy.h" />
     <ClInclude Include="LdrWatcher.h" />
     <ClInclude Include="GL\GLExtensions\ARB_texture_compression_bptc.h">
       <Filter>GL\GLExtensions</Filter>

Source/Core/Common/Lazy.h

@@ -0,0 +1,38 @@
+// Copyright 2017 Dolphin Emulator Project
+// Licensed under GPLv2+
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include <functional>
+#include <utility>
+#include <variant>
+
+namespace Common
+{
+// A Lazy object holds a value. If a Lazy object is constructed using
+// a function as an argument, that function will be called to compute
+// the value the first time any code tries to access the value.
+
+template <typename T>
+class Lazy
+{
+public:
+  Lazy() : m_value(T()) {}
+  Lazy(const std::variant<T, std::function<T()>>& value) : m_value(value) {}
+  Lazy(std::variant<T, std::function<T()>>&& value) : m_value(std::move(value)) {}
+  const T& operator*() const { return *ComputeValue(); }
+  const T* operator->() const { return ComputeValue(); }
+  T& operator*() { return *ComputeValue(); }
+  T* operator->() { return ComputeValue(); }
+private:
+  T* ComputeValue() const
+  {
+    if (!std::holds_alternative<T>(m_value))
+      m_value = std::get<std::function<T()>>(m_value)();
+    return &std::get<T>(m_value);
+  }
+
+  mutable std::variant<T, std::function<T()>> m_value;
+};
+}

Source/Core/DiscIO/VolumeWii.cpp

@@ -45,7 +45,6 @@ VolumeWii::VolumeWii(std::unique_ptr<BlobReader> reader)
     return;
   }
 
-  // Get tickets, TMDs, and decryption keys for all partitions
   for (u32 partition_group = 0; partition_group < 4; ++partition_group)
   {
     const std::optional<u32> number_of_partitions =
@@ -61,56 +60,61 @@ VolumeWii::VolumeWii(std::unique_ptr<BlobReader> reader)
 
     for (u32 i = 0; i < number_of_partitions; i++)
     {
-      // Read the partition offset
       read_buffer = m_pReader->ReadSwapped<u32>(partition_table_offset + (i * 8));
       if (!read_buffer)
         continue;
       const u64 partition_offset = static_cast<u64>(*read_buffer) << 2;
+      const Partition partition(partition_offset);
 
-      // Read the partition type
       const std::optional<u32> partition_type =
           m_pReader->ReadSwapped<u32>(partition_table_offset + (i * 8) + 4);
       if (!partition_type)
         continue;
 
-      // Read ticket
+      // If this is the game partition, set m_game_partition
-      std::vector<u8> ticket_buffer(sizeof(IOS::ES::Ticket));
-      if (!m_pReader->Read(partition_offset, ticket_buffer.size(), ticket_buffer.data()))
-        continue;
-      IOS::ES::TicketReader ticket{std::move(ticket_buffer)};
-      if (!ticket.IsValid())
-        continue;
-
-      // Read TMD
-      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);
-      if (!tmd_size || !tmd_address)
-        continue;
-      *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
-        // (after allocating the buffer), so we do the check here.
-        PanicAlert("Invalid TMD size");
-        continue;
-      }
-      std::vector<u8> tmd_buffer(*tmd_size);
-      if (!m_pReader->Read(partition_offset + *tmd_address, *tmd_size, tmd_buffer.data()))
-        continue;
-      IOS::ES::TMDReader tmd{std::move(tmd_buffer)};
-
-      // Get the decryption key
-      const std::array<u8, 16> key = ticket.GetTitleKey();
-      std::unique_ptr<mbedtls_aes_context> aes_context = std::make_unique<mbedtls_aes_context>();
-      mbedtls_aes_setkey_dec(aes_context.get(), key.data(), 128);
-
-      // We've read everything. Time to store it! (The reason we don't store anything
-      // earlier is because we want to be able to skip adding the partition if an error occurs.)
-      const Partition partition(partition_offset);
-      m_partitions.emplace(partition, PartitionDetails{std::move(aes_context), std::move(ticket),
-                                                       std::move(tmd), *partition_type});
       if (m_game_partition == PARTITION_NONE && *partition_type == 0)
         m_game_partition = partition;
+
+      auto get_ticket = [this, partition]() -> IOS::ES::TicketReader {
+        std::vector<u8> ticket_buffer(sizeof(IOS::ES::Ticket));
+        if (!m_pReader->Read(partition.offset, ticket_buffer.size(), ticket_buffer.data()))
+          return INVALID_TICKET;
+        return IOS::ES::TicketReader{std::move(ticket_buffer)};
+      };
+
+      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);
+        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
+          // (after allocating the buffer), so we do the check here.
+          PanicAlert("Invalid TMD size");
+          return INVALID_TMD;
+        }
+        std::vector<u8> tmd_buffer(*tmd_size);
+        if (!m_pReader->Read(partition.offset + *tmd_address, *tmd_size, tmd_buffer.data()))
+          return INVALID_TMD;
+        return IOS::ES::TMDReader{std::move(tmd_buffer)};
+      };
+
+      auto get_key = [this, partition]() -> std::unique_ptr<mbedtls_aes_context> {
+        const IOS::ES::TicketReader& ticket = *m_partitions[partition].ticket;
+        if (!ticket.IsValid())
+          return nullptr;
+        const std::array<u8, 16> key = ticket.GetTitleKey();
+        std::unique_ptr<mbedtls_aes_context> aes_context = std::make_unique<mbedtls_aes_context>();
+        mbedtls_aes_setkey_dec(aes_context.get(), key.data(), 128);
+        return aes_context;
+      };
+
+      m_partitions.emplace(
+          partition, PartitionDetails{Common::Lazy<std::unique_ptr<mbedtls_aes_context>>(get_key),
+                                      Common::Lazy<IOS::ES::TicketReader>(get_ticket),
+                                      Common::Lazy<IOS::ES::TMDReader>(get_tmd), *partition_type});
     }
   }
 }
@@ -128,7 +132,9 @@ bool VolumeWii::Read(u64 _ReadOffset, u64 _Length, u8* _pBuffer, const Partition
   auto it = m_partitions.find(partition);
   if (it == m_partitions.end())
     return false;
-  mbedtls_aes_context* aes_context = it->second.key.get();
+  mbedtls_aes_context* aes_context = it->second.key->get();
+  if (!aes_context)
+    return false;
 
   std::vector<u8> read_buffer(BLOCK_TOTAL_SIZE);
   while (_Length > 0)
@@ -202,13 +208,13 @@ std::optional<u64> VolumeWii::GetTitleID(const Partition& partition) const
 const IOS::ES::TicketReader& VolumeWii::GetTicket(const Partition& partition) const
 {
   auto it = m_partitions.find(partition);
-  return it != m_partitions.end() ? it->second.ticket : INVALID_TICKET;
+  return it != m_partitions.end() ? *it->second.ticket : INVALID_TICKET;
 }
 
 const IOS::ES::TMDReader& VolumeWii::GetTMD(const Partition& partition) const
 {
   auto it = m_partitions.find(partition);
-  return it != m_partitions.end() ? it->second.tmd : INVALID_TMD;
+  return it != m_partitions.end() ? *it->second.tmd : INVALID_TMD;
 }
 
 u64 VolumeWii::PartitionOffsetToRawOffset(u64 offset, const Partition& partition)
@@ -342,7 +348,9 @@ bool VolumeWii::CheckIntegrity(const Partition& partition) const
   auto it = m_partitions.find(partition);
   if (it == m_partitions.end())
     return false;
-  mbedtls_aes_context* aes_context = it->second.key.get();
+  mbedtls_aes_context* aes_context = it->second.key->get();
+  if (!aes_context)
+    return false;
 
   // Get partition data size
   u32 partSizeDiv4;

Source/Core/DiscIO/VolumeWii.h

@@ -12,6 +12,7 @@
 #include <vector>
 
 #include "Common/CommonTypes.h"
+#include "Common/Lazy.h"
 #include "Core/IOS/ES/Formats.h"
 #include "DiscIO/Volume.h"
 
@@ -66,9 +67,9 @@ public:
 private:
   struct PartitionDetails
   {
-    std::unique_ptr<mbedtls_aes_context> key;
+    Common::Lazy<std::unique_ptr<mbedtls_aes_context>> key;
-    IOS::ES::TicketReader ticket;
+    Common::Lazy<IOS::ES::TicketReader> ticket;
-    IOS::ES::TMDReader tmd;
+    Common::Lazy<IOS::ES::TMDReader> tmd;
     u32 type;
   };