WIA: Add early support for WIA writing

4 files changed, 239 insertions, 2 deletions

diff --git a/Source/Core/DiscIO/Blob.h b/Source/Core/DiscIO/Blob.h
index 5371478d65..692ac57072 100644
--- a/Source/Core/DiscIO/Blob.h
+++ b/Source/Core/DiscIO/Blob.h
@@ -173,5 +173,8 @@ bool ConvertToGCZ(BlobReader* infile, const std::string& infile_path,
 bool ConvertToPlain(BlobReader* infile, const std::string& infile_path,
                     const std::string& outfile_path, CompressCB callback = nullptr,
                     void* arg = nullptr);
+bool ConvertToWIA(BlobReader* infile, const std::string& infile_path,
+                  const std::string& outfile_path, int chunk_size, CompressCB callback = nullptr,
+                  void* arg = nullptr);
 
 }  // namespace DiscIO
diff --git a/Source/Core/DiscIO/WIABlob.cpp b/Source/Core/DiscIO/WIABlob.cpp
index f97ceca75a..4a523a5c25 100644
--- a/Source/Core/DiscIO/WIABlob.cpp
+++ b/Source/Core/DiscIO/WIABlob.cpp
@@ -17,12 +17,16 @@
 #include <mbedtls/sha1.h>
 
 #include "Common/Align.h"
+#include "Common/Assert.h"
 #include "Common/CommonTypes.h"
 #include "Common/File.h"
+#include "Common/FileUtil.h"
 #include "Common/Logging/Log.h"
+#include "Common/MsgHandler.h"
 #include "Common/StringUtil.h"
 #include "Common/Swap.h"
 
+#include "DiscIO/Blob.h"
 #include "DiscIO/VolumeWii.h"
 #include "DiscIO/WiiEncryptionCache.h"
 
@@ -887,4 +891,195 @@ bool WIAFileReader::Chunk::ApplyHashExceptions(
   return true;
 }
 
+bool WIAFileReader::PadTo4(File::IOFile* file, u64* bytes_written)
+{
+  constexpr u32 ZEROES = 0;
+  const u64 bytes_to_write = Common::AlignUp(*bytes_written, 4) - *bytes_written;
+  if (bytes_to_write == 0)
+    return true;
+
+  *bytes_written += bytes_to_write;
+  return file->WriteBytes(&ZEROES, bytes_to_write);
+}
+
+WIAFileReader::ConversionResult WIAFileReader::ConvertToWIA(BlobReader* infile,
+                                                            File::IOFile* outfile, int chunk_size,
+                                                            CompressCB callback, void* arg)
+{
+  ASSERT(infile->IsDataSizeAccurate());
+  ASSERT(chunk_size > 0);
+
+  const u64 iso_size = infile->GetDataSize();
+
+  u64 bytes_read = 0;
+  u64 bytes_written = 0;
+
+  // These two headers will be filled in with proper values at the very end
+  WIAHeader1 header_1;
+  WIAHeader2 header_2;
+  if (!outfile->WriteArray(&header_1, 1) || !outfile->WriteArray(&header_2, 1))
+    return ConversionResult::WriteFailed;
+  bytes_written += sizeof(WIAHeader1) + sizeof(WIAHeader2);
+  if (!PadTo4(outfile, &bytes_written))
+    return ConversionResult::WriteFailed;
+
+  std::vector<GroupEntry> group_entries;
+  group_entries.resize(Common::AlignUp(iso_size, chunk_size) / chunk_size);
+
+  std::vector<RawDataEntry> raw_data_entries;
+  raw_data_entries.emplace_back(
+      RawDataEntry{Common::swap64(header_2.disc_header.size()),
+                   Common::swap64(iso_size - header_2.disc_header.size()), 0,
+                   Common::swap32(static_cast<u32>(group_entries.size()))});
+
+  std::vector<PartitionEntry> partition_entries;
+
+  const auto run_callback = [&](size_t groups_written) {
+    int ratio = 0;
+    if (bytes_read != 0)
+      ratio = static_cast<int>(100 * bytes_written / bytes_read);
+
+    const std::string temp =
+        StringFromFormat(Common::GetStringT("%i of %i blocks. Compression ratio %i%%").c_str(),
+                         groups_written, group_entries.size(), ratio);
+    return callback(temp, static_cast<float>(groups_written) / group_entries.size(), arg);
+  };
+
+  if (!infile->Read(0, header_2.disc_header.size(), header_2.disc_header.data()))
+    return ConversionResult::ReadFailed;
+  // We intentially do not increment bytes_read here, since these bytes will be read again
+
+  if (!run_callback(0))
+    return ConversionResult::Canceled;
+
+  std::vector<u8> buffer(chunk_size);
+  for (size_t i = 0; i < group_entries.size(); ++i)
+  {
+    const u64 bytes_to_read = std::min<u64>(chunk_size, iso_size - bytes_read);
+
+    if (bytes_written >> 2 > std::numeric_limits<u32>::max())
+      return ConversionResult::InternalError;
+
+    ASSERT((bytes_written & 3) == 0);
+    group_entries[i] = GroupEntry{Common::swap32(static_cast<u32>(bytes_written >> 2)),
+                                  Common::swap32(static_cast<u32>(bytes_to_read))};
+
+    if (!infile->Read(bytes_read, bytes_to_read, buffer.data()))
+      return ConversionResult::ReadFailed;
+    if (!outfile->WriteArray(buffer.data(), bytes_to_read))
+      return ConversionResult::WriteFailed;
+
+    bytes_read += bytes_to_read;
+    bytes_written += bytes_to_read;
+    if (!PadTo4(outfile, &bytes_written))
+      return ConversionResult::WriteFailed;
+
+    if (!run_callback(i))
+      return ConversionResult::Canceled;
+  }
+
+  const u64 partition_entries_offset = bytes_written;
+  const u64 partition_entries_size = partition_entries.size() * sizeof(PartitionEntry);
+  if (!outfile->WriteArray(partition_entries.data(), partition_entries.size()))
+    return ConversionResult::WriteFailed;
+  bytes_written += partition_entries_size;
+  if (!PadTo4(outfile, &bytes_written))
+    return ConversionResult::WriteFailed;
+
+  const u64 raw_data_entries_offset = bytes_written;
+  const u64 raw_data_entries_size = raw_data_entries.size() * sizeof(RawDataEntry);
+  if (!outfile->WriteArray(raw_data_entries.data(), raw_data_entries.size()))
+    return ConversionResult::WriteFailed;
+  bytes_written += raw_data_entries_size;
+  if (!PadTo4(outfile, &bytes_written))
+    return ConversionResult::WriteFailed;
+
+  const u64 group_entries_offset = bytes_written;
+  const u64 group_entries_size = group_entries.size() * sizeof(GroupEntry);
+  if (!outfile->WriteArray(group_entries.data(), group_entries.size()))
+    return ConversionResult::WriteFailed;
+  bytes_written += group_entries_size;
+  if (!PadTo4(outfile, &bytes_written))
+    return ConversionResult::WriteFailed;
+
+  header_2.disc_type = 0;  // TODO
+  header_2.compression_type = Common::swap32(static_cast<u32>(CompressionType::None));
+  header_2.compression_level = 0;
+  header_2.chunk_size = Common::swap32(static_cast<u32>(chunk_size));
+
+  header_2.number_of_partition_entries = Common::swap32(static_cast<u32>(partition_entries.size()));
+  header_2.partition_entry_size = Common::swap32(sizeof(PartitionEntry));
+  header_2.partition_entries_offset = Common::swap64(partition_entries_offset);
+
+  if (partition_entries.data() == nullptr)
+    partition_entries.reserve(1);  // Avoid a crash in mbedtls_sha1_ret
+  mbedtls_sha1_ret(reinterpret_cast<const u8*>(partition_entries.data()), partition_entries_size,
+                   header_2.partition_entries_hash.data());
+
+  header_2.number_of_raw_data_entries = Common::swap32(static_cast<u32>(raw_data_entries.size()));
+  header_2.raw_data_entries_offset = Common::swap64(raw_data_entries_offset);
+  header_2.raw_data_entries_size = Common::swap32(static_cast<u32>(raw_data_entries_size));
+
+  header_2.number_of_group_entries = Common::swap32(static_cast<u32>(group_entries.size()));
+  header_2.group_entries_offset = Common::swap64(group_entries_offset);
+  header_2.group_entries_size = Common::swap32(static_cast<u32>(group_entries_size));
+
+  header_2.compressor_data_size = 0;
+  std::fill(std::begin(header_2.compressor_data), std::end(header_2.compressor_data), 0);
+
+  header_1.magic = WIA_MAGIC;
+  header_1.version = Common::swap32(WIA_VERSION);
+  header_1.version_compatible = Common::swap32(WIA_VERSION_WRITE_COMPATIBLE);
+  header_1.header_2_size = Common::swap32(sizeof(WIAHeader2));
+  mbedtls_sha1_ret(reinterpret_cast<const u8*>(&header_2), sizeof(header_2),
+                   header_1.header_2_hash.data());
+  header_1.iso_file_size = Common::swap64(infile->GetDataSize());
+  header_1.wia_file_size = Common::swap64(bytes_written);
+  mbedtls_sha1_ret(reinterpret_cast<const u8*>(&header_1), offsetof(WIAHeader1, header_1_hash),
+                   header_1.header_1_hash.data());
+
+  if (!outfile->Seek(0, SEEK_SET))
+    return ConversionResult::WriteFailed;
+  if (!outfile->WriteArray(&header_1, 1) || !outfile->WriteArray(&header_2, 1))
+    return ConversionResult::WriteFailed;
+
+  return ConversionResult::Success;
+}
+
+bool ConvertToWIA(BlobReader* infile, const std::string& infile_path,
+                  const std::string& outfile_path, int chunk_size, CompressCB callback, void* arg)
+{
+  File::IOFile outfile(outfile_path, "wb");
+  if (!outfile)
+  {
+    PanicAlertT("Failed to open the output file \"%s\".\n"
+                "Check that you have permissions to write the target folder and that the media can "
+                "be written.",
+                outfile_path.c_str());
+    return false;
+  }
+
+  WIAFileReader::ConversionResult result =
+      WIAFileReader::ConvertToWIA(infile, &outfile, chunk_size, callback, arg);
+
+  if (result == WIAFileReader::ConversionResult::ReadFailed)
+    PanicAlertT("Failed to read from the input file \"%s\".", infile_path.c_str());
+
+  if (result == WIAFileReader::ConversionResult::WriteFailed)
+  {
+    PanicAlertT("Failed to write the output file \"%s\".\n"
+                "Check that you have enough space available on the target drive.",
+                outfile_path.c_str());
+  }
+
+  if (result != WIAFileReader::ConversionResult::Success)
+  {
+    // Remove the incomplete output file
+    outfile.Close();
+    File::Delete(outfile_path);
+  }
+
+  return result == WIAFileReader::ConversionResult::Success;
+}
+
 }  // namespace DiscIO
diff --git a/Source/Core/DiscIO/WIABlob.h b/Source/Core/DiscIO/WIABlob.h
index 63defd6055..d42aa92fef 100644
--- a/Source/Core/DiscIO/WIABlob.h
+++ b/Source/Core/DiscIO/WIABlob.h
@@ -44,6 +44,18 @@ public:
   bool SupportsReadWiiDecrypted() const override;
   bool ReadWiiDecrypted(u64 offset, u64 size, u8* out_ptr, u64 partition_data_offset) override;
 
+  enum class ConversionResult
+  {
+    Success,
+    Canceled,
+    ReadFailed,
+    WriteFailed,
+    InternalError,
+  };
+
+  static ConversionResult ConvertToWIA(BlobReader* infile, File::IOFile* outfile, int chunk_size,
+                                       CompressCB callback, void* arg);
+
 private:
   using SHA1 = std::array<u8, 20>;
   using WiiKey = std::array<u8, 16>;
@@ -275,6 +287,8 @@ private:
 
   static std::string VersionToString(u32 version);
 
+  static bool PadTo4(File::IOFile* file, u64* bytes_written);
+
   bool m_valid;
   CompressionType m_compression_type;
 
diff --git a/Source/Core/DolphinQt/ConvertDialog.cpp b/Source/Core/DolphinQt/ConvertDialog.cpp
index 0d01785dc4..831c0ec641 100644
--- a/Source/Core/DolphinQt/ConvertDialog.cpp
+++ b/Source/Core/DolphinQt/ConvertDialog.cpp
@@ -57,6 +57,7 @@ ConvertDialog::ConvertDialog(QList<std::shared_ptr<const UICommon::GameFile>> fi
   m_format = new QComboBox;
   m_format->addItem(QStringLiteral("ISO"), static_cast<int>(DiscIO::BlobType::PLAIN));
   m_format->addItem(QStringLiteral("GCZ"), static_cast<int>(DiscIO::BlobType::GCZ));
+  m_format->addItem(QStringLiteral("WIA"), static_cast<int>(DiscIO::BlobType::WIA));
   if (std::all_of(m_files.begin(), m_files.end(),
                   [](const auto& file) { return file->GetBlobType() == DiscIO::BlobType::PLAIN; }))
   {
@@ -88,7 +89,10 @@ ConvertDialog::ConvertDialog(QList<std::shared_ptr<const UICommon::GameFile>> fi
                     "It takes up more space than any other format.\n\n"
                     "GCZ: A basic compressed format which is compatible with most versions of "
                     "Dolphin and some other programs. It can't efficiently compress junk data "
-                    "(unless removed) or encrypted Wii data."));
+                    "(unless removed) or encrypted Wii data.\n\n"
+                    "WIA: An advanced compressed format which is compatible with recent versions "
+                    "of Dolphin and a few other programs. It can efficiently compress encrypted "
+                    "Wii data, but not junk data (unless removed)."));
   info_text->setWordWrap(true);
 
   QVBoxLayout* info_layout = new QVBoxLayout;
@@ -166,6 +170,13 @@ void ConvertDialog::OnFormatChanged()
 
     break;
   }
+  case DiscIO::BlobType::WIA:
+    m_block_size->setEnabled(true);
+
+    // This is the smallest block size supported by WIA. For performance, larger sizes are avoided.
+    AddToBlockSizeComboBox(0x200000);
+
+    break;
   default:
     break;
   }
@@ -224,7 +235,11 @@ void ConvertDialog::Convert()
     break;
   case DiscIO::BlobType::GCZ:
     extension = QStringLiteral(".gcz");
-    filter = tr("Compressed GC/Wii images (*.gcz)");
+    filter = tr("GCZ GC/Wii images (*.gcz)");
+    break;
+  case DiscIO::BlobType::WIA:
+    extension = QStringLiteral(".wia");
+    filter = tr("WIA GC/Wii images (*.wia)");
     break;
   default:
     ASSERT(false);
@@ -351,6 +366,16 @@ void ConvertDialog::Convert()
           return good;
         });
       }
+      else if (format == DiscIO::BlobType::WIA)
+      {
+        good = std::async(std::launch::async, [&] {
+          const bool good =
+              DiscIO::ConvertToWIA(blob_reader.get(), original_path, dst_path.toStdString(),
+                                   block_size, &CompressCB, &progress_dialog);
+          progress_dialog.Reset();
+          return good;
+        });
+      }
 
       progress_dialog.GetRaw()->exec();
       if (!good.get())