Merge pull request #9409 from AdmiralCurtiss/wii-save-import-tmd

Make WiiSave::Import() behave closer to the Wii System Menu's SD Card save copying.

Source/Core/Core/HW/WiiSave.cpp

@@ -40,6 +40,7 @@
 #include "Core/IOS/IOS.h"
 #include "Core/IOS/IOSC.h"
 #include "Core/IOS/Uids.h"
+#include "Core/WiiUtils.h"
 
 namespace WiiSave
 {
@@ -68,9 +69,35 @@ class NandStorage final : public Storage
     ScanForFiles(m_data_dir);
   }
 
-  bool SaveExists() override
+  bool SaveExists() const override
   {
-    return m_uid && m_gid && m_fs->GetMetadata(*m_uid, *m_gid, m_data_dir + "/banner.bin");
+    return !m_files_list.empty() ||
+           (m_uid && m_gid && m_fs->GetMetadata(*m_uid, *m_gid, m_data_dir + "/banner.bin"));
+  }
+
+  bool EraseSave() override
+  {
+    // banner.bin is not in m_files_list, delete separately
+    const auto banner_delete_result =
+        m_fs->Delete(IOS::PID_KERNEL, IOS::PID_KERNEL, m_data_dir + "/banner.bin");
+    if (banner_delete_result != FS::ResultCode::Success)
+      return false;
+
+    for (const SaveFile& file : m_files_list)
+    {
+      // files in subdirs are deleted automatically when the subdir is deleted
+      if (file.path.find('/') != std::string::npos)
+        continue;
+
+      const auto result =
+          m_fs->Delete(IOS::PID_KERNEL, IOS::PID_KERNEL, m_data_dir + "/" + file.path);
+      if (result != FS::ResultCode::Success)
+        return false;
+    }
+
+    m_files_list.clear();
+    m_files_size = 0;
+    return true;
   }
 
   std::optional<Header> ReadHeader() override
@@ -245,6 +272,10 @@ class DataBinStorage final : public Storage
     m_file = File::IOFile{path, mode};
   }
 
+  bool SaveExists() const override { return m_file.GetSize() > 0; }
+
+  bool EraseSave() override { return m_file.GetSize() == 0 || m_file.Resize(0); }
+
   std::optional<Header> ReadHeader() override
   {
     Header header;
@@ -446,42 +477,87 @@ StoragePointer MakeDataBinStorage(IOS::HLE::IOSC* iosc, const std::string& path,
   return StoragePointer{new DataBinStorage{iosc, path, mode}};
 }
 
-template <typename T>
-static bool Copy(std::string_view description, Storage* source,
-                 std::optional<T> (Storage::*read_fn)(), Storage* dest,
-                 bool (Storage::*write_fn)(const T&))
+CopyResult Copy(Storage* source, Storage* dest)
 {
-  const std::optional<T> data = (source->*read_fn)();
-  if (data && (dest->*write_fn)(*data))
-    return true;
-  ERROR_LOG_FMT(CORE, "WiiSave::Copy: Failed to {} {}", !data ? "read" : "write", description);
-  return false;
-}
+  // first make sure we can read all the data from the source
+  const auto header = source->ReadHeader();
+  if (!header)
+  {
+    ERROR_LOG_FMT(CORE, "WiiSave::Copy: Failed to read header");
+    return CopyResult::CorruptedSource;
+  }
 
-bool Copy(Storage* source, Storage* dest)
-{
-  return Copy("header", source, &Storage::ReadHeader, dest, &Storage::WriteHeader) &&
-         Copy("bk header", source, &Storage::ReadBkHeader, dest, &Storage::WriteBkHeader) &&
-         Copy("files", source, &Storage::ReadFiles, dest, &Storage::WriteFiles);
+  const auto bk_header = source->ReadBkHeader();
+  if (!bk_header)
+  {
+    ERROR_LOG_FMT(CORE, "WiiSave::Copy: Failed to read bk header");
+    return CopyResult::CorruptedSource;
+  }
+
+  const auto files = source->ReadFiles();
+  if (!files)
+  {
+    ERROR_LOG_FMT(CORE, "WiiSave::Copy: Failed to read files");
+    return CopyResult::CorruptedSource;
+  }
+
+  // once we have confirmed we can read the source, erase corresponding save in the destination
+  if (dest->SaveExists())
+  {
+    if (!dest->EraseSave())
+    {
+      ERROR_LOG_FMT(CORE, "WiiSave::Copy: Failed to erase existing save");
+      return CopyResult::Error;
+    }
+  }
+
+  // and then write it to the destination
+  if (!dest->WriteHeader(*header))
+  {
+    ERROR_LOG_FMT(CORE, "WiiSave::Copy: Failed to write header");
+    return CopyResult::Error;
+  }
+
+  if (!dest->WriteBkHeader(*bk_header))
+  {
+    ERROR_LOG_FMT(CORE, "WiiSave::Copy: Failed to write bk header");
+    return CopyResult::Error;
+  }
+
+  if (!dest->WriteFiles(*files))
+  {
+    ERROR_LOG_FMT(CORE, "WiiSave::Copy: Failed to write files");
+    return CopyResult::Error;
+  }
+
+  return CopyResult::Success;
 }
 
-bool Import(const std::string& data_bin_path, std::function<bool()> can_overwrite)
+CopyResult Import(const std::string& data_bin_path, std::function<bool()> can_overwrite)
 {
   IOS::HLE::Kernel ios;
   const auto data_bin = MakeDataBinStorage(&ios.GetIOSC(), data_bin_path, "rb");
   const std::optional<Header> header = data_bin->ReadHeader();
   if (!header)
   {
     ERROR_LOG_FMT(CORE, "WiiSave::Import: Failed to read header");
-    return false;
+    return CopyResult::CorruptedSource;
   }
+
+  if (!WiiUtils::EnsureTMDIsImported(*ios.GetFS(), *ios.GetES(), header->tid))
+  {
+    ERROR_LOG_FMT(CORE, "WiiSave::Import: Failed to find or import TMD for title {:16x}",
+                  header->tid);
+    return CopyResult::TitleMissing;
+  }
+
   const auto nand = MakeNandStorage(ios.GetFS().get(), header->tid);
   if (nand->SaveExists() && !can_overwrite())
-    return false;
+    return CopyResult::Cancelled;
   return Copy(data_bin.get(), nand.get());
 }
 
-static bool Export(u64 tid, std::string_view export_path, IOS::HLE::Kernel* ios)
+static CopyResult Export(u64 tid, std::string_view export_path, IOS::HLE::Kernel* ios)
 {
   const std::string path = fmt::format("{}/private/wii/title/{}{}{}{}/data.bin", export_path,
                                        static_cast<char>(tid >> 24), static_cast<char>(tid >> 16),
@@ -490,7 +566,7 @@ static bool Export(u64 tid, std::string_view export_path, IOS::HLE::Kernel* ios)
               MakeDataBinStorage(&ios->GetIOSC(), path, "w+b").get());
 }
 
-bool Export(u64 tid, std::string_view export_path)
+CopyResult Export(u64 tid, std::string_view export_path)
 {
   IOS::HLE::Kernel ios;
   return Export(tid, export_path, &ios);
@@ -502,7 +578,7 @@ size_t ExportAll(std::string_view export_path)
   size_t exported_save_count = 0;
   for (const u64 title : ios.GetES()->GetInstalledTitles())
   {
-    if (Export(title, export_path, &ios))
+    if (Export(title, export_path, &ios) == CopyResult::Success)
       ++exported_save_count;
   }
   return exported_save_count;

Source/Core/Core/HW/WiiSave.h

@@ -32,12 +32,21 @@ using StoragePointer = std::unique_ptr<Storage, StorageDeleter>;
 StoragePointer MakeNandStorage(IOS::HLE::FS::FileSystem* fs, u64 tid);
 StoragePointer MakeDataBinStorage(IOS::HLE::IOSC* iosc, const std::string& path, const char* mode);
 
-bool Copy(Storage* source, Storage* destination);
+enum class CopyResult
+{
+  Success,
+  Error,
+  Cancelled,
+  CorruptedSource,
+  TitleMissing,
+};
+
+CopyResult Copy(Storage* source, Storage* destination);
 
 /// Import a save into the NAND from a .bin file.
-bool Import(const std::string& data_bin_path, std::function<bool()> can_overwrite);
+CopyResult Import(const std::string& data_bin_path, std::function<bool()> can_overwrite);
 /// Export a save to a .bin file.
-bool Export(u64 tid, std::string_view export_path);
+CopyResult Export(u64 tid, std::string_view export_path);
 /// Export all saves that are in the NAND. Returns the number of exported saves.
 size_t ExportAll(std::string_view export_path);
 }  // namespace WiiSave

Source/Core/Core/HW/WiiSaveStructs.h

@@ -101,7 +101,8 @@ class Storage
   };
 
   virtual ~Storage() = default;
-  virtual bool SaveExists() { return true; }
+  virtual bool SaveExists() const = 0;
+  virtual bool EraseSave() = 0;
   virtual std::optional<Header> ReadHeader() = 0;
   virtual std::optional<BkHeader> ReadBkHeader() = 0;
   virtual std::optional<std::vector<SaveFile>> ReadFiles() = 0;

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

@@ -126,7 +126,8 @@ class ES final : public Device
   ReturnCode ImportTicket(const std::vector<u8>& ticket_bytes, const std::vector<u8>& cert_chain,
                           TicketImportType type = TicketImportType::PossiblyPersonalised,
                           VerifySignature verify_signature = VerifySignature::Yes);
-  ReturnCode ImportTmd(Context& context, const std::vector<u8>& tmd_bytes);
+  ReturnCode ImportTmd(Context& context, const std::vector<u8>& tmd_bytes, u64 caller_title_id,
+                       u32 caller_title_flags);
   ReturnCode ImportTitleInit(Context& context, const std::vector<u8>& tmd_bytes,
                              const std::vector<u8>& cert_chain,
                              VerifySignature verify_signature = VerifySignature::Yes);
@@ -135,7 +136,8 @@ class ES final : public Device
   ReturnCode ImportContentEnd(Context& context, u32 content_fd);
   ReturnCode ImportTitleDone(Context& context);
   ReturnCode ImportTitleCancel(Context& context);
-  ReturnCode ExportTitleInit(Context& context, u64 title_id, u8* tmd, u32 tmd_size);
+  ReturnCode ExportTitleInit(Context& context, u64 title_id, u8* tmd, u32 tmd_size,
+                             u64 caller_title_id, u32 caller_title_flags);
   ReturnCode ExportContentBegin(Context& context, u64 title_id, u32 content_id);
   ReturnCode ExportContentData(Context& context, u32 content_fd, u8* data, u32 data_size);
   ReturnCode ExportContentEnd(Context& context, u32 content_fd);

Source/Core/Core/IOS/ES/TitleManagement.cpp

@@ -107,15 +107,14 @@ IPCCommandResult ES::ImportTicket(const IOCtlVRequest& request)
 constexpr std::array<u8, 16> NULL_KEY{};
 
 // Used for exporting titles and importing them back (ImportTmd and ExportTitleInit).
-static ReturnCode InitBackupKey(const IOS::ES::TMDReader& tmd, IOSC& iosc, IOSC::Handle* key)
+static ReturnCode InitBackupKey(u64 tid, u32 title_flags, IOSC& iosc, IOSC::Handle* key)
 {
   // Some versions of IOS have a bug that causes it to use a zeroed key instead of the PRNG key.
   // When Nintendo decided to fix it, they added checks to keep using the zeroed key only in
   // affected titles to avoid making existing exports useless.
 
   // Ignore the region byte.
-  const u64 title_id = tmd.GetTitleId() | 0xff;
-  const u32 title_flags = tmd.GetTitleFlags();
+  const u64 title_id = tid | 0xff;
   const u32 affected_type = IOS::ES::TITLE_TYPE_0x10 | IOS::ES::TITLE_TYPE_DATA;
   if (title_id == Titles::SYSTEM_MENU || (title_flags & affected_type) != affected_type ||
       !(title_id == 0x00010005735841ff || title_id - 0x00010005735a41ff <= 0x700))
@@ -136,7 +135,8 @@ static void ResetTitleImportContext(ES::Context* context, IOSC& iosc)
   context->title_import_export = {};
 }
 
-ReturnCode ES::ImportTmd(Context& context, const std::vector<u8>& tmd_bytes)
+ReturnCode ES::ImportTmd(Context& context, const std::vector<u8>& tmd_bytes, u64 caller_title_id,
+                         u32 caller_title_flags)
 {
   INFO_LOG_FMT(IOS_ES, "ImportTmd");
 
@@ -166,8 +166,8 @@ ReturnCode ES::ImportTmd(Context& context, const std::vector<u8>& tmd_bytes)
     return ES_EIO;
   }
 
-  ret =
-      InitBackupKey(m_title_context.tmd, m_ios.GetIOSC(), &context.title_import_export.key_handle);
+  ret = InitBackupKey(caller_title_id, caller_title_flags, m_ios.GetIOSC(),
+                      &context.title_import_export.key_handle);
   if (ret != IPC_SUCCESS)
   {
     ERROR_LOG_FMT(IOS_ES, "ImportTmd: InitBackupKey failed with error {}", ret);
@@ -189,7 +189,8 @@ IPCCommandResult ES::ImportTmd(Context& context, const IOCtlVRequest& request)
 
   std::vector<u8> tmd(request.in_vectors[0].size);
   Memory::CopyFromEmu(tmd.data(), request.in_vectors[0].address, request.in_vectors[0].size);
-  return GetDefaultReply(ImportTmd(context, tmd));
+  return GetDefaultReply(ImportTmd(context, tmd, m_title_context.tmd.GetTitleId(),
+                                   m_title_context.tmd.GetTitleFlags()));
 }
 
 static ReturnCode InitTitleImportKey(const std::vector<u8>& ticket_bytes, IOSC& iosc,
@@ -651,7 +652,8 @@ IPCCommandResult ES::DeleteContent(const IOCtlVRequest& request)
                                        Memory::Read_U32(request.in_vectors[1].address)));
 }
 
-ReturnCode ES::ExportTitleInit(Context& context, u64 title_id, u8* tmd_bytes, u32 tmd_size)
+ReturnCode ES::ExportTitleInit(Context& context, u64 title_id, u8* tmd_bytes, u32 tmd_size,
+                               u64 caller_title_id, u32 caller_title_flags)
 {
   // No concurrent title import/export is allowed.
   if (context.title_import_export.valid)
@@ -664,8 +666,8 @@ ReturnCode ES::ExportTitleInit(Context& context, u64 title_id, u8* tmd_bytes, u3
   ResetTitleImportContext(&context, m_ios.GetIOSC());
   context.title_import_export.tmd = tmd;
 
-  const ReturnCode ret =
-      InitBackupKey(m_title_context.tmd, m_ios.GetIOSC(), &context.title_import_export.key_handle);
+  const ReturnCode ret = InitBackupKey(caller_title_id, caller_title_flags, m_ios.GetIOSC(),
+                                       &context.title_import_export.key_handle);
   if (ret != IPC_SUCCESS)
     return ret;
 
@@ -688,7 +690,9 @@ IPCCommandResult ES::ExportTitleInit(Context& context, const IOCtlVRequest& requ
   u8* tmd_bytes = Memory::GetPointer(request.io_vectors[0].address);
   const u32 tmd_size = request.io_vectors[0].size;
 
-  return GetDefaultReply(ExportTitleInit(context, title_id, tmd_bytes, tmd_size));
+  return GetDefaultReply(ExportTitleInit(context, title_id, tmd_bytes, tmd_size,
+                                         m_title_context.tmd.GetTitleId(),
+                                         m_title_context.tmd.GetTitleFlags()));
 }
 
 ReturnCode ES::ExportContentBegin(Context& context, u64 title_id, u32 content_id)

Source/Core/Core/WiiUtils.cpp

@@ -19,6 +19,7 @@
 #include <fmt/format.h>
 #include <pugixml.hpp>
 
+#include "Common/Align.h"
 #include "Common/Assert.h"
 #include "Common/CommonTypes.h"
 #include "Common/FileUtil.h"
@@ -223,6 +224,67 @@ bool IsTitleInstalled(u64 title_id)
                      [](const std::string& file) { return file != "title.tmd"; });
 }
 
+bool IsTMDImported(IOS::HLE::FS::FileSystem& fs, u64 title_id)
+{
+  const auto entries = fs.ReadDirectory(0, 0, Common::GetTitleContentPath(title_id));
+  return entries && std::any_of(entries->begin(), entries->end(),
+                                [](const std::string& file) { return file == "title.tmd"; });
+}
+
+IOS::ES::TMDReader FindBackupTMD(IOS::HLE::FS::FileSystem& fs, u64 title_id)
+{
+  auto file = fs.OpenFile(IOS::PID_KERNEL, IOS::PID_KERNEL,
+                          "/title/00000001/00000002/data/tmds.sys", IOS::HLE::FS::Mode::Read);
+  if (!file)
+    return {};
+
+  // structure of this file is as follows:
+  // - 32 bytes descriptor of a TMD, which contains a title ID and a length
+  // - the TMD, with padding aligning to 32 bytes
+  // - repeat for as many TMDs as stored
+  while (true)
+  {
+    std::array<u8, 32> descriptor;
+    if (!file->Read(descriptor.data(), descriptor.size()))
+      return {};
+
+    const u64 tid = Common::swap64(descriptor.data());
+    const u32 tmd_length = Common::swap32(descriptor.data() + 8);
+    if (tid == title_id)
+    {
+      // found the right TMD
+      std::vector<u8> tmd_bytes(tmd_length);
+      if (!file->Read(tmd_bytes.data(), tmd_length))
+        return {};
+      return IOS::ES::TMDReader(std::move(tmd_bytes));
+    }
+
+    // not the right TMD, skip this one and go to the next
+    if (!file->Seek(Common::AlignUp(tmd_length, 32), IOS::HLE::FS::SeekMode::Current))
+      return {};
+  }
+}
+
+bool EnsureTMDIsImported(IOS::HLE::FS::FileSystem& fs, IOS::HLE::Device::ES& es, u64 title_id)
+{
+  if (IsTMDImported(fs, title_id))
+    return true;
+
+  auto tmd = FindBackupTMD(fs, title_id);
+  if (!tmd.IsValid())
+    return false;
+
+  IOS::HLE::Device::ES::Context context;
+  context.uid = IOS::SYSMENU_UID;
+  context.gid = IOS::SYSMENU_GID;
+  const auto import_result =
+      es.ImportTmd(context, tmd.GetBytes(), Titles::SYSTEM_MENU, IOS::ES::TITLE_TYPE_DEFAULT);
+  if (import_result != IOS::HLE::IPC_SUCCESS)
+    return false;
+
+  return es.ImportTitleDone(context) == IOS::HLE::IPC_SUCCESS;
+}
+
 // Common functionality for system updaters.
 class SystemUpdater
 {