diff --git a/Source/Core/Common/Common.vcxproj b/Source/Core/Common/Common.vcxproj index 5af496cc9f..9466fda774 100644 --- a/Source/Core/Common/Common.vcxproj +++ b/Source/Core/Common/Common.vcxproj @@ -126,6 +126,7 @@ + diff --git a/Source/Core/Common/Common.vcxproj.filters b/Source/Core/Common/Common.vcxproj.filters index 161557ddc0..4b48650ec1 100644 --- a/Source/Core/Common/Common.vcxproj.filters +++ b/Source/Core/Common/Common.vcxproj.filters @@ -253,6 +253,7 @@ GL\GLExtensions + GL\GLExtensions diff --git a/Source/Core/Common/Lazy.h b/Source/Core/Common/Lazy.h new file mode 100644 index 0000000000..322d4bfeac --- /dev/null +++ b/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 +#include +#include + +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 +class Lazy +{ +public: + Lazy() : m_value(T()) {} + Lazy(const std::variant>& value) : m_value(value) {} + Lazy(std::variant>&& 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(m_value)) + m_value = std::get>(m_value)(); + return &std::get(m_value); + } + + mutable std::variant> m_value; +}; +} diff --git a/Source/Core/DiscIO/VolumeWii.cpp b/Source/Core/DiscIO/VolumeWii.cpp index 7225a1e57f..7f7670d6d7 100644 --- a/Source/Core/DiscIO/VolumeWii.cpp +++ b/Source/Core/DiscIO/VolumeWii.cpp @@ -45,7 +45,6 @@ VolumeWii::VolumeWii(std::unique_ptr reader) return; } - // Get tickets, TMDs, and decryption keys for all partitions for (u32 partition_group = 0; partition_group < 4; ++partition_group) { const std::optional number_of_partitions = @@ -61,56 +60,61 @@ VolumeWii::VolumeWii(std::unique_ptr reader) for (u32 i = 0; i < number_of_partitions; i++) { - // Read the partition offset read_buffer = m_pReader->ReadSwapped(partition_table_offset + (i * 8)); if (!read_buffer) continue; const u64 partition_offset = static_cast(*read_buffer) << 2; + const Partition partition(partition_offset); - // Read the partition type const std::optional partition_type = m_pReader->ReadSwapped(partition_table_offset + (i * 8) + 4); if (!partition_type) continue; - // Read ticket - std::vector 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 tmd_size = m_pReader->ReadSwapped(partition_offset + 0x2a4); - std::optional tmd_address = m_pReader->ReadSwapped(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 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 key = ticket.GetTitleKey(); - std::unique_ptr aes_context = std::make_unique(); - 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 this is the game partition, set m_game_partition if (m_game_partition == PARTITION_NONE && *partition_type == 0) m_game_partition = partition; + + auto get_ticket = [this, partition]() -> IOS::ES::TicketReader { + std::vector 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 tmd_size = m_pReader->ReadSwapped(partition.offset + 0x2a4); + std::optional tmd_address = m_pReader->ReadSwapped(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 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 { + const IOS::ES::TicketReader& ticket = *m_partitions[partition].ticket; + if (!ticket.IsValid()) + return nullptr; + const std::array key = ticket.GetTitleKey(); + std::unique_ptr aes_context = std::make_unique(); + mbedtls_aes_setkey_dec(aes_context.get(), key.data(), 128); + return aes_context; + }; + + m_partitions.emplace( + partition, PartitionDetails{Common::Lazy>(get_key), + Common::Lazy(get_ticket), + Common::Lazy(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 read_buffer(BLOCK_TOTAL_SIZE); while (_Length > 0) @@ -202,13 +208,13 @@ std::optional 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; diff --git a/Source/Core/DiscIO/VolumeWii.h b/Source/Core/DiscIO/VolumeWii.h index d8cdfe8cb8..7ada5805c3 100644 --- a/Source/Core/DiscIO/VolumeWii.h +++ b/Source/Core/DiscIO/VolumeWii.h @@ -12,6 +12,7 @@ #include #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 key; - IOS::ES::TicketReader ticket; - IOS::ES::TMDReader tmd; + Common::Lazy> key; + Common::Lazy ticket; + Common::Lazy tmd; u32 type; };