From 3613f33c9b2d6c924f309a1901132e4c20c8cf60 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9o=20Lam?= <leo@innovatetechnologi.es>
Date: Tue, 20 Jun 2017 12:15:48 +0200
Subject: [PATCH] IOS/ES: Make sure all contents are imported

This adds a check to ImportTitleDone to make sure all required contents
that are listed in the TMD have been imported before allowing to finish
the import. Not checking for this could allow titles to be left in an
inconsistent state.
---
 Source/Core/Core/IOS/ES/Formats.cpp         |  5 +++++
 Source/Core/Core/IOS/ES/Formats.h           |  1 +
 Source/Core/Core/IOS/ES/TitleManagement.cpp | 20 +++++++++++++++++++-
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/Source/Core/Core/IOS/ES/Formats.cpp b/Source/Core/Core/IOS/ES/Formats.cpp
index 8790dd2ff1..214fd8a31f 100644
--- a/Source/Core/Core/IOS/ES/Formats.cpp
+++ b/Source/Core/Core/IOS/ES/Formats.cpp
@@ -61,6 +61,11 @@ bool Content::IsShared() const
   return (type & 0x8000) != 0;
 }
 
+bool Content::IsOptional() const
+{
+  return (type & 0x4000) != 0;
+}
+
 SignedBlobReader::SignedBlobReader(const std::vector<u8>& bytes) : m_bytes(bytes)
 {
 }
diff --git a/Source/Core/Core/IOS/ES/Formats.h b/Source/Core/Core/IOS/ES/Formats.h
index 943f93d8da..dd1352b367 100644
--- a/Source/Core/Core/IOS/ES/Formats.h
+++ b/Source/Core/Core/IOS/ES/Formats.h
@@ -83,6 +83,7 @@ static_assert(sizeof(TMDHeader) == 0x1e4, "TMDHeader has the wrong size");
 struct Content
 {
   bool IsShared() const;
+  bool IsOptional() const;
   u32 id;
   u16 index;
   u16 type;
diff --git a/Source/Core/Core/IOS/ES/TitleManagement.cpp b/Source/Core/Core/IOS/ES/TitleManagement.cpp
index 968e672e3e..143c4a86d3 100644
--- a/Source/Core/Core/IOS/ES/TitleManagement.cpp
+++ b/Source/Core/Core/IOS/ES/TitleManagement.cpp
@@ -353,13 +353,31 @@ ReturnCode ES::ImportTitleDone(Context& context)
   if (!context.title_import.tmd.IsValid() || context.title_import.content_id != 0xFFFFFFFF)
     return ES_EINVAL;
 
+  // Make sure all listed, non-optional contents have been imported.
+  const u64 title_id = context.title_import.tmd.GetTitleId();
+  const std::vector<IOS::ES::Content> contents = context.title_import.tmd.GetContents();
+  const IOS::ES::SharedContentMap shared_content_map{Common::FROM_SESSION_ROOT};
+  const bool has_all_required_contents =
+      std::all_of(contents.cbegin(), contents.cend(), [&](const IOS::ES::Content& content) {
+        if (content.IsOptional())
+          return true;
+
+        if (content.IsShared())
+          return shared_content_map.GetFilenameFromSHA1(content.sha1).has_value();
+
+        return File::Exists(Common::GetTitleContentPath(title_id, Common::FROM_SESSION_ROOT) +
+                            StringFromFormat("%08x.app", content.id));
+      });
+  if (!has_all_required_contents)
+    return ES_EINVAL;
+
   if (!WriteImportTMD(context.title_import.tmd))
     return ES_EIO;
 
   if (!FinishImport(context.title_import.tmd))
     return ES_EIO;
 
-  INFO_LOG(IOS_ES, "ImportTitleDone: title %016" PRIx64, context.title_import.tmd.GetTitleId());
+  INFO_LOG(IOS_ES, "ImportTitleDone: title %016" PRIx64, title_id);
   context.title_import.tmd.SetBytes({});
   return IPC_SUCCESS;
 }