Re-implement FixCommonKeyIndex for WAD files

The old implementation of this was not able to distinguish between
a title that had the common key index set to 1 because it actually
was Korean and a title that had the common key index set to 1 due to
fakesigning. This new implementation solves the problem by
decrypting a content with each possible common key and checking
which result matches the provided SHA-1 hash.

The problem that the old implementation causes has only been reported
to affect a certain pirated WAD of Chronos Twins DX (WC6EUP), but it's
possible that the problem would start affecting more WADs if we add
support for the vWii common key (which uses index 2). Adding support
for the vWii common key would also prevent us from using the simpler
solution of always forcing the index to 0 if the title is not Korean.

Source/Core/Core/IOS/ES/Formats.cpp

@@ -455,7 +455,7 @@ std::array<u8, 16> TicketReader::GetTitleKey(const HLE::IOSC& iosc) const
   u8 index = m_bytes.at(offsetof(Ticket, common_key_index));
   if (index >= HLE::IOSC::COMMON_KEY_HANDLES.size())
   {
-    WARN_LOG(IOS_ES, "Bad common key index for title %016" PRIx64 ": %u -- using common key 0",
+    PanicAlert("Bad common key index for title %016" PRIx64 ": %u -- using common key 0",
                GetTitleId(), index);
     index = 0;
   }
@@ -533,11 +533,9 @@ HLE::ReturnCode TicketReader::Unpersonalise(HLE::IOSC& iosc)
   return ret;
 }
 
-void TicketReader::FixCommonKeyIndex()
+void TicketReader::OverwriteCommonKeyIndex(u8 index)
 {
-  u8& index = m_bytes[offsetof(Ticket, common_key_index)];
+  m_bytes[offsetof(Ticket, common_key_index)] = index;
-  // Assume the ticket is using the normal common key if it's an invalid value.
-  index = index < HLE::IOSC::COMMON_KEY_HANDLES.size() ? index : 0;
 }
 
 struct SharedContentMap::Entry

Source/Core/Core/IOS/ES/Formats.h

@@ -257,9 +257,8 @@ public:
   // and has a title key that must be decrypted first.
   HLE::ReturnCode Unpersonalise(HLE::IOSC& iosc);
 
-  // Reset the common key field back to 0 if it's an incorrect value.
   // Intended for use before importing fakesigned tickets, which tend to have a high bogus index.
-  void FixCommonKeyIndex();
+  void OverwriteCommonKeyIndex(u8 index);
 };
 
 class SharedContentMap final

Source/Core/Core/WiiUtils.cpp

@@ -61,9 +61,8 @@ static bool ImportWAD(IOS::HLE::Kernel& ios, const DiscIO::VolumeWAD& wad)
   IOS::HLE::ReturnCode ret;
   const bool checks_enabled = SConfig::GetInstance().m_enable_signature_checks;
 
-  IOS::ES::TicketReader ticket = wad.GetTicket();
   // Ensure the common key index is correct, as it's checked by IOS.
-  ticket.FixCommonKeyIndex();
+  IOS::ES::TicketReader ticket = wad.GetTicketWithFixedCommonKey();
 
   while ((ret = es->ImportTicket(ticket.GetBytes(), wad.GetCertificateChain(),
                                  IOS::HLE::Device::ES::TicketImportType::Unpersonalised)) < 0 ||

Source/Core/DiscIO/Volume.h

@@ -76,6 +76,12 @@ public:
   }
   virtual std::vector<u8> GetContent(u16 index) const { return {}; }
   virtual std::vector<u64> GetContentOffsets() const { return {}; }
+  virtual bool CheckContentIntegrity(const IOS::ES::Content& content, u64 content_offset,
+                                     const IOS::ES::TicketReader& ticket) const
+  {
+    return false;
+  }
+  virtual IOS::ES::TicketReader GetTicketWithFixedCommonKey() const { return {}; }
   // Returns a non-owning pointer. Returns nullptr if the file system couldn't be read.
   virtual const FileSystem* GetFileSystem(const Partition& partition) const = 0;
   virtual u64 PartitionOffsetToRawOffset(u64 offset, const Partition& partition) const

Source/Core/DiscIO/VolumeVerifier.cpp

@@ -11,7 +11,6 @@
 #include <string>
 #include <unordered_set>
 
-#include <mbedtls/aes.h>
 #include <mbedtls/md5.h>
 #include <mbedtls/sha1.h>
 #include <zlib.h>
@@ -604,29 +603,33 @@ void VolumeVerifier::CheckMisc()
     }
   }
 
-  const IOS::ES::TicketReader& ticket = m_volume.GetTicket(m_volume.GetGamePartition());
+  m_ticket = m_volume.GetTicket(m_volume.GetGamePartition());
-  if (ticket.IsValid())
+  if (m_ticket.IsValid())
   {
-    const u8 common_key = ticket.GetCommonKeyIndex();
+    const u8 specified_common_key_index = m_ticket.GetCommonKeyIndex();
 
-    if (common_key > IOS::HLE::IOSC::COMMON_KEY_HANDLES.size())
+    // Wii discs only use common key 0 (regular) and common key 1 (Korean), not common key 2 (vWii).
+    if (m_volume.GetVolumeType() == Platform::WiiDisc && specified_common_key_index > 1)
     {
-      // Many fakesigned WADs have the common key index set to a (random?) bogus value.
-      // For WADs, Dolphin will detect this and use common key 0 instead, making this low severity.
-      const Severity severity =
-          m_volume.GetVolumeType() == Platform::WiiWAD ? Severity::Low : Severity::High;
-      // i18n: This is "common" as in "shared", not the opposite of "uncommon"
-      AddProblem(severity, Common::GetStringT("This title is set to use an invalid common key."));
-    }
-
-    if (common_key == 1 && region != Region::NTSC_K)
-    {
-      // Apparently a certain pirate WAD of Chronos Twins DX unluckily got an index of 1,
-      // which Dolphin does not change to 0 because 1 is valid on Korean Wiis.
-      // https://forums.dolphin-emu.org/Thread-wiiware-chronos-twins-dx
       AddProblem(Severity::High,
                  // i18n: This is "common" as in "shared", not the opposite of "uncommon"
-                 Common::GetStringT("This non-Korean title is set to use the Korean common key."));
+                 Common::GetStringT("This title is set to use an invalid common key."));
+    }
+
+    if (m_volume.GetVolumeType() == Platform::WiiWAD)
+    {
+      m_ticket = m_volume.GetTicketWithFixedCommonKey();
+      const u8 fixed_common_key_index = m_ticket.GetCommonKeyIndex();
+      if (specified_common_key_index != fixed_common_key_index)
+      {
+        // Many fakesigned WADs have the common key index set to a (random?) bogus value.
+        // For WADs, Dolphin will detect this and use the correct key, making this low severity.
+        std::string text = StringFromFormat(
+            // i18n: This is "common" as in "shared", not the opposite of "uncommon"
+            Common::GetStringT("The specified common key index is %u but should be %u.").c_str(),
+            specified_common_key_index, fixed_common_key_index);
+        AddProblem(Severity::Low, std::move(text));
+      }
     }
   }
 
@@ -737,7 +740,7 @@ void VolumeVerifier::Process()
 
   if (content_read)
   {
-    if (!CheckContentIntegrity(content))
+    if (!m_volume.CheckContentIntegrity(content, m_content_offsets[m_content_index], m_ticket))
     {
       AddProblem(
           Severity::High,
@@ -768,27 +771,6 @@ void VolumeVerifier::Process()
   }
 }
 
-bool VolumeVerifier::CheckContentIntegrity(const IOS::ES::Content& content)
-{
-  std::vector<u8> encrypted_data = m_volume.GetContent(content.index);
-
-  mbedtls_aes_context context;
-  const std::array<u8, 16> key = m_volume.GetTicket(PARTITION_NONE).GetTitleKey();
-  mbedtls_aes_setkey_dec(&context, key.data(), 128);
-
-  std::array<u8, 16> iv{};
-  iv[0] = static_cast<u8>(content.index >> 8);
-  iv[1] = static_cast<u8>(content.index & 0xFF);
-
-  std::vector<u8> decrypted_data(Common::AlignUp(content.size, 0x40));
-  mbedtls_aes_crypt_cbc(&context, MBEDTLS_AES_DECRYPT, decrypted_data.size(), iv.data(),
-                        encrypted_data.data(), decrypted_data.data());
-
-  std::array<u8, 20> sha1;
-  mbedtls_sha1_ret(decrypted_data.data(), content.size, sha1.data());
-  return sha1 == content.sha1;
-}
-
 u64 VolumeVerifier::GetBytesProcessed() const
 {
   return m_progress;

Source/Core/DiscIO/VolumeVerifier.h

@@ -13,6 +13,7 @@
 #include <mbedtls/sha1.h>
 
 #include "Common/CommonTypes.h"
+#include "Core/IOS/ES/Formats.h"
 #include "DiscIO/DiscScrubber.h"
 #include "DiscIO/Volume.h"
 
@@ -29,12 +30,6 @@
 //
 // GetResult() can be called before the processing is finished, but the result will be incomplete.
 
-namespace IOS::ES
-{
-struct Content;
-class SignedBlobReader;
-}  // namespace IOS::ES
-
 namespace DiscIO
 {
 class FileInfo;
@@ -103,7 +98,6 @@ private:
   u64 GetBiggestUsedOffset(const FileInfo& file_info) const;
   void CheckMisc();
   void SetUpHashing();
-  bool CheckContentIntegrity(const IOS::ES::Content& content);
 
   void AddProblem(Severity severity, std::string text);
 
@@ -120,6 +114,7 @@ private:
   mbedtls_sha1_context m_sha1_context;
 
   DiscScrubber m_scrubber;
+  IOS::ES::TicketReader m_ticket;
   std::vector<u64> m_content_offsets;
   u16 m_content_index = 0;
   std::vector<BlockToVerify> m_blocks;

Source/Core/DiscIO/VolumeWad.cpp

@@ -2,6 +2,7 @@
 // Licensed under GPLv2+
 // Refer to the license.txt file included.
 
+#include <algorithm>
 #include <cstddef>
 #include <cstring>
 #include <locale>
@@ -12,12 +13,16 @@
 #include <utility>
 #include <vector>
 
+#include <mbedtls/aes.h>
+#include <mbedtls/sha1.h>
+
 #include "Common/Align.h"
 #include "Common/Assert.h"
 #include "Common/CommonTypes.h"
 #include "Common/Logging/Log.h"
 #include "Common/MsgHandler.h"
 #include "Common/StringUtil.h"
+#include "Core/IOS/IOSC.h"
 #include "DiscIO/Blob.h"
 #include "DiscIO/Enums.h"
 #include "DiscIO/Volume.h"
@@ -147,6 +152,95 @@ std::vector<u64> VolumeWAD::GetContentOffsets() const
   return content_offsets;
 }
 
+bool VolumeWAD::CheckContentIntegrity(const IOS::ES::Content& content,
+                                      const std::vector<u8>& encrypted_data,
+                                      const IOS::ES::TicketReader& ticket) const
+{
+  mbedtls_aes_context context;
+  const std::array<u8, 16> key = ticket.GetTitleKey();
+  mbedtls_aes_setkey_dec(&context, key.data(), 128);
+
+  std::array<u8, 16> iv{};
+  iv[0] = static_cast<u8>(content.index >> 8);
+  iv[1] = static_cast<u8>(content.index & 0xFF);
+
+  std::vector<u8> decrypted_data(encrypted_data.size());
+  mbedtls_aes_crypt_cbc(&context, MBEDTLS_AES_DECRYPT, decrypted_data.size(), iv.data(),
+                        encrypted_data.data(), decrypted_data.data());
+
+  std::array<u8, 20> sha1;
+  mbedtls_sha1_ret(decrypted_data.data(), content.size, sha1.data());
+  return sha1 == content.sha1;
+}
+
+bool VolumeWAD::CheckContentIntegrity(const IOS::ES::Content& content, u64 content_offset,
+                                      const IOS::ES::TicketReader& ticket) const
+{
+  std::vector<u8> encrypted_data(Common::AlignUp(content.size, 0x40));
+  if (!m_reader->Read(content_offset, encrypted_data.size(), encrypted_data.data()))
+    return false;
+  return CheckContentIntegrity(content, encrypted_data, ticket);
+}
+
+IOS::ES::TicketReader VolumeWAD::GetTicketWithFixedCommonKey() const
+{
+  if (!m_ticket.IsValid() || !m_tmd.IsValid())
+    return m_ticket;
+
+  const std::vector<u8> sig = m_ticket.GetSignatureData();
+  if (!std::all_of(sig.cbegin(), sig.cend(), [](u8 a) { return a == 0; }))
+  {
+    // This does not look like a typical "invalid common key index" ticket, so let's assume
+    // the index is correct. This saves some time when reading properly signed titles.
+    return m_ticket;
+  }
+
+  const std::vector<IOS::ES::Content> contents = m_tmd.GetContents();
+  if (contents.empty())
+    return m_ticket;
+
+  // Find the smallest content so that we spend as little time as possible in CheckContentIntegrity
+  IOS::ES::Content smallest_content = contents[0];
+  u64 offset_of_smallest_content = m_data_offset;
+
+  u64 offset = m_data_offset;
+  for (const IOS::ES::Content& content : contents)
+  {
+    if (content.size < smallest_content.size)
+    {
+      smallest_content = content;
+      offset_of_smallest_content = offset;
+    }
+    offset += Common::AlignUp(content.size, 0x40);
+  }
+
+  std::vector<u8> content_data(Common::AlignUp(smallest_content.size, 0x40));
+  if (!m_reader->Read(offset_of_smallest_content, content_data.size(), content_data.data()))
+    return m_ticket;
+
+  const u8 specified_index = m_ticket.GetCommonKeyIndex();
+  if (specified_index < IOS::HLE::IOSC::COMMON_KEY_HANDLES.size() &&
+      CheckContentIntegrity(smallest_content, content_data, m_ticket))
+  {
+    return m_ticket;  // The common key index is already correct
+  }
+
+  // Try every common key index except the one we already tried
+  IOS::ES::TicketReader new_ticket = m_ticket;
+  for (u8 i = 0; i < IOS::HLE::IOSC::COMMON_KEY_HANDLES.size(); ++i)
+  {
+    if (i != specified_index)
+    {
+      new_ticket.OverwriteCommonKeyIndex(i);
+      if (CheckContentIntegrity(smallest_content, content_data, new_ticket))
+        return new_ticket;  // We've found the common key index that should be used
+    }
+  }
+
+  ERROR_LOG(DISCIO, "Couldn't find valid common key for WAD file (%u specified)", specified_index);
+  return m_ticket;
+}
+
 std::string VolumeWAD::GetGameID(const Partition& partition) const
 {
   return m_tmd.GetGameID();

Source/Core/DiscIO/VolumeWad.h

@@ -39,6 +39,9 @@ public:
   GetCertificateChain(const Partition& partition = PARTITION_NONE) const override;
   std::vector<u8> GetContent(u16 index) const override;
   std::vector<u64> GetContentOffsets() const override;
+  bool CheckContentIntegrity(const IOS::ES::Content& content, u64 content_offset,
+                             const IOS::ES::TicketReader& ticket) const override;
+  IOS::ES::TicketReader GetTicketWithFixedCommonKey() const override;
   std::string GetGameID(const Partition& partition = PARTITION_NONE) const override;
   std::string GetGameTDBID(const Partition& partition = PARTITION_NONE) const override;
   std::string GetMakerID(const Partition& partition = PARTITION_NONE) const override;
@@ -63,6 +66,9 @@ public:
   u64 GetRawSize() const override;
 
 private:
+  bool CheckContentIntegrity(const IOS::ES::Content& content, const std::vector<u8>& encrypted_data,
+                             const IOS::ES::TicketReader& ticket) const;
+
   std::unique_ptr<BlobReader> m_reader;
   IOS::ES::TicketReader m_ticket;
   IOS::ES::TMDReader m_tmd;