DiscIO: Move scrubbing code out of ConvertToGCZ

This way, scrubbing can also be performed when converting
to other formats.

11 files changed, 194 insertions, 72 deletions

diff --git a/Source/Core/DiscIO/Blob.h b/Source/Core/DiscIO/Blob.h
index 9cb00087d2..1d23d7b6b0 100644
--- a/Source/Core/DiscIO/Blob.h
+++ b/Source/Core/DiscIO/Blob.h
@@ -165,9 +165,11 @@ std::unique_ptr<BlobReader> CreateBlobReader(const std::string& filename);
 
 typedef bool (*CompressCB)(const std::string& text, float percent, void* arg);
 
-bool ConvertToGCZ(const std::string& infile_path, const std::string& outfile_path, u32 sub_type = 0,
-                  int sector_size = 16384, CompressCB callback = nullptr, void* arg = nullptr);
-bool ConvertToPlain(const std::string& infile_path, const std::string& outfile_path,
-                    CompressCB callback = nullptr, void* arg = nullptr);
+bool ConvertToGCZ(BlobReader* infile, const std::string& infile_path,
+                  const std::string& outfile_path, u32 sub_type, int sector_size = 16384,
+                  CompressCB callback = nullptr, void* arg = nullptr);
+bool ConvertToPlain(BlobReader* infile, const std::string& infile_path,
+                    const std::string& outfile_path, CompressCB callback = nullptr,
+                    void* arg = nullptr);
 
 }  // namespace DiscIO
diff --git a/Source/Core/DiscIO/CMakeLists.txt b/Source/Core/DiscIO/CMakeLists.txt
index 279e72030d..4970f135c2 100644
--- a/Source/Core/DiscIO/CMakeLists.txt
+++ b/Source/Core/DiscIO/CMakeLists.txt
@@ -23,6 +23,8 @@ add_library(discio
   Filesystem.h
   NANDImporter.cpp
   NANDImporter.h
+  ScrubbedBlob.cpp
+  ScrubbedBlob.h
   TGCBlob.cpp
   TGCBlob.h
   Volume.cpp
diff --git a/Source/Core/DiscIO/CompressedBlob.cpp b/Source/Core/DiscIO/CompressedBlob.cpp
index 5e6723679a..cda3cad6d3 100644
--- a/Source/Core/DiscIO/CompressedBlob.cpp
+++ b/Source/Core/DiscIO/CompressedBlob.cpp
@@ -154,19 +154,11 @@ bool CompressedBlobReader::GetBlock(u64 block_num, u8* out_ptr)
   return true;
 }
 
-bool ConvertToGCZ(const std::string& infile_path, const std::string& outfile_path, u32 sub_type,
-                  int block_size, CompressCB callback, void* arg)
+bool ConvertToGCZ(BlobReader* infile, const std::string& infile_path,
+                  const std::string& outfile_path, u32 sub_type, int block_size,
+                  CompressCB callback, void* arg)
 {
-  bool scrubbing = false;
-
-  std::unique_ptr<VolumeDisc> infile = CreateDisc(infile_path);
-  if (!infile)
-  {
-    PanicAlertT("Failed to open the input file \"%s\".", infile_path.c_str());
-    return false;
-  }
-
-  ASSERT(infile->IsSizeAccurate());
+  ASSERT(infile->IsDataSizeAccurate());
 
   File::IOFile outfile(outfile_path, "wb");
   if (!outfile)
@@ -178,19 +170,6 @@ bool ConvertToGCZ(const std::string& infile_path, const std::string& outfile_pat
     return false;
   }
 
-  DiscScrubber disc_scrubber;
-  if (sub_type == 1)
-  {
-    if (!disc_scrubber.SetupScrub(infile.get()))
-    {
-      PanicAlertT("\"%s\" failed to be scrubbed. Probably the image is corrupt.",
-                  infile_path.c_str());
-      return false;
-    }
-
-    scrubbing = true;
-  }
-
   z_stream z = {};
   if (deflateInit(&z, 9) != Z_OK)
     return false;
@@ -201,7 +180,7 @@ bool ConvertToGCZ(const std::string& infile_path, const std::string& outfile_pat
   header.magic_cookie = GCZ_MAGIC;
   header.sub_type = sub_type;
   header.block_size = block_size;
-  header.data_size = infile->GetSize();
+  header.data_size = infile->GetDataSize();
 
   // round upwards!
   header.num_blocks = (u32)((header.data_size + (block_size - 1)) / block_size);
@@ -245,11 +224,9 @@ bool ConvertToGCZ(const std::string& infile_path, const std::string& outfile_pat
 
     offsets[i] = position;
 
-    const u64 bytes_to_read = scrubbing && disc_scrubber.CanBlockBeScrubbed(inpos) ?
-                                  0 :
-                                  std::min<u64>(block_size, header.data_size - inpos);
+    const u64 bytes_to_read = std::min<u64>(block_size, header.data_size - inpos);
 
-    success = infile->Read(inpos, bytes_to_read, in_buf.data(), PARTITION_NONE);
+    success = infile->Read(inpos, bytes_to_read, in_buf.data());
     if (!success)
     {
       PanicAlertT("Failed to read from the input file \"%s\".", infile_path.c_str());
diff --git a/Source/Core/DiscIO/DiscIO.vcxproj b/Source/Core/DiscIO/DiscIO.vcxproj
index dd322dbb62..98bb500452 100644
--- a/Source/Core/DiscIO/DiscIO.vcxproj
+++ b/Source/Core/DiscIO/DiscIO.vcxproj
@@ -56,6 +56,7 @@
     <ClCompile Include="Filesystem.cpp" />
     <ClCompile Include="FileSystemGCWii.cpp" />
     <ClCompile Include="NANDImporter.cpp" />
+    <ClCompile Include="ScrubbedBlob.cpp" />
     <ClCompile Include="TGCBlob.cpp" />
     <ClCompile Include="Volume.cpp" />
     <ClCompile Include="VolumeFileBlobReader.cpp" />
@@ -80,6 +81,7 @@
     <ClInclude Include="Filesystem.h" />
     <ClInclude Include="FileSystemGCWii.h" />
     <ClInclude Include="NANDImporter.h" />
+    <ClInclude Include="ScrubbedBlob.h" />
     <ClInclude Include="TGCBlob.h" />
     <ClInclude Include="Volume.h" />
     <ClInclude Include="VolumeFileBlobReader.h" />
diff --git a/Source/Core/DiscIO/DiscIO.vcxproj.filters b/Source/Core/DiscIO/DiscIO.vcxproj.filters
index 5fb5891cfe..9b19fa8471 100644
--- a/Source/Core/DiscIO/DiscIO.vcxproj.filters
+++ b/Source/Core/DiscIO/DiscIO.vcxproj.filters
@@ -87,6 +87,9 @@
     <ClCompile Include="WiiEncryptionCache.cpp">
       <Filter>Volume\Blob</Filter>
     </ClCompile>
+    <ClCompile Include="ScrubbedBlob.cpp">
+      <Filter>Volume\Blob</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="DiscScrubber.h">
@@ -155,6 +158,9 @@
     <ClInclude Include="WiiEncryptionCache.h">
       <Filter>Volume\Blob</Filter>
     </ClInclude>
+    <ClInclude Include="ScrubbedBlob.h">
+      <Filter>Volume\Blob</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <Text Include="CMakeLists.txt" />
diff --git a/Source/Core/DiscIO/DiscScrubber.cpp b/Source/Core/DiscIO/DiscScrubber.cpp
index 06c444c84d..e7b098b374 100644
--- a/Source/Core/DiscIO/DiscScrubber.cpp
+++ b/Source/Core/DiscIO/DiscScrubber.cpp
@@ -25,8 +25,6 @@
 
 namespace DiscIO
 {
-constexpr size_t CLUSTER_SIZE = 0x8000;
-
 DiscScrubber::DiscScrubber() = default;
 DiscScrubber::~DiscScrubber() = default;
 
diff --git a/Source/Core/DiscIO/DiscScrubber.h b/Source/Core/DiscIO/DiscScrubber.h
index 4b46f6353a..5d64fa5376 100644
--- a/Source/Core/DiscIO/DiscScrubber.h
+++ b/Source/Core/DiscIO/DiscScrubber.h
@@ -35,6 +35,8 @@ public:
   // Returns true if the specified 32 KiB block only contains unused data
   bool CanBlockBeScrubbed(u64 offset) const;
 
+  static constexpr size_t CLUSTER_SIZE = 0x8000;
+
 private:
   void MarkAsUsed(u64 offset, u64 size);
   void MarkAsUsedE(u64 partition_data_offset, u64 offset, u64 size);
diff --git a/Source/Core/DiscIO/FileBlob.cpp b/Source/Core/DiscIO/FileBlob.cpp
index 95beeb1bc3..ca0aa2cfd0 100644
--- a/Source/Core/DiscIO/FileBlob.cpp
+++ b/Source/Core/DiscIO/FileBlob.cpp
@@ -41,17 +41,10 @@ bool PlainFileReader::Read(u64 offset, u64 nbytes, u8* out_ptr)
   }
 }
 
-bool ConvertToPlain(const std::string& infile_path, const std::string& outfile_path,
-                    CompressCB callback, void* arg)
+bool ConvertToPlain(BlobReader* infile, const std::string& infile_path,
+                    const std::string& outfile_path, CompressCB callback, void* arg)
 {
-  std::unique_ptr<BlobReader> reader = CreateBlobReader(infile_path);
-  if (!reader)
-  {
-    PanicAlertT("Failed to open the input file \"%s\".", infile_path.c_str());
-    return false;
-  }
-
-  ASSERT(reader->IsDataSizeAccurate());
+  ASSERT(infile->IsDataSizeAccurate());
 
   File::IOFile outfile(outfile_path, "wb");
   if (!outfile)
@@ -64,7 +57,7 @@ bool ConvertToPlain(const std::string& infile_path, const std::string& outfile_p
   }
 
   constexpr size_t DESIRED_BUFFER_SIZE = 0x80000;
-  u64 buffer_size = reader->GetBlockSize();
+  u64 buffer_size = infile->GetBlockSize();
   if (buffer_size == 0)
   {
     buffer_size = DESIRED_BUFFER_SIZE;
@@ -76,7 +69,7 @@ bool ConvertToPlain(const std::string& infile_path, const std::string& outfile_p
   }
 
   std::vector<u8> buffer(buffer_size);
-  const u64 num_buffers = (reader->GetDataSize() + buffer_size - 1) / buffer_size;
+  const u64 num_buffers = (infile->GetDataSize() + buffer_size - 1) / buffer_size;
   int progress_monitor = std::max<int>(1, num_buffers / 100);
   bool success = true;
 
@@ -93,8 +86,8 @@ bool ConvertToPlain(const std::string& infile_path, const std::string& outfile_p
       }
     }
     const u64 inpos = i * buffer_size;
-    const u64 sz = std::min(buffer_size, reader->GetDataSize() - inpos);
-    if (!reader->Read(inpos, sz, buffer.data()))
+    const u64 sz = std::min(buffer_size, infile->GetDataSize() - inpos);
+    if (!infile->Read(inpos, sz, buffer.data()))
     {
       PanicAlertT("Failed to read from the input file \"%s\".", infile_path.c_str());
       success = false;
diff --git a/Source/Core/DiscIO/ScrubbedBlob.cpp b/Source/Core/DiscIO/ScrubbedBlob.cpp
new file mode 100644
index 0000000000..255edebf57
--- /dev/null
+++ b/Source/Core/DiscIO/ScrubbedBlob.cpp
@@ -0,0 +1,67 @@
+// Copyright 2020 Dolphin Emulator Project
+// Licensed under GPLv2+
+// Refer to the license.txt file included.
+
+#include "DiscIO/ScrubbedBlob.h"
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "Common/Align.h"
+#include "DiscIO/Blob.h"
+#include "DiscIO/DiscScrubber.h"
+#include "DiscIO/Volume.h"
+
+namespace DiscIO
+{
+ScrubbedBlob::ScrubbedBlob(std::unique_ptr<BlobReader> blob_reader, DiscScrubber scrubber)
+    : m_blob_reader(std::move(blob_reader)), m_scrubber(std::move(scrubber))
+{
+}
+
+std::unique_ptr<ScrubbedBlob> ScrubbedBlob::Create(const std::string& path)
+{
+  std::unique_ptr<VolumeDisc> disc = CreateDisc(path);
+  if (!disc)
+    return nullptr;
+
+  DiscScrubber scrubber;
+  if (!scrubber.SetupScrub(disc.get()))
+    return nullptr;
+
+  std::unique_ptr<BlobReader> blob = CreateBlobReader(path);
+  if (!blob)
+    return nullptr;
+
+  return std::unique_ptr<ScrubbedBlob>(new ScrubbedBlob(std::move(blob), std::move(scrubber)));
+}
+
+bool ScrubbedBlob::Read(u64 offset, u64 size, u8* out_ptr)
+{
+  while (size > 0)
+  {
+    constexpr size_t CLUSTER_SIZE = DiscScrubber::CLUSTER_SIZE;
+    const u64 bytes_to_read =
+        std::min(Common::AlignDown(offset + CLUSTER_SIZE, CLUSTER_SIZE) - offset, size);
+
+    if (m_scrubber.CanBlockBeScrubbed(offset))
+    {
+      std::fill_n(out_ptr, bytes_to_read, 0);
+    }
+    else
+    {
+      if (!m_blob_reader->Read(offset, bytes_to_read, out_ptr))
+        return false;
+    }
+
+    offset += bytes_to_read;
+    size -= bytes_to_read;
+    out_ptr += bytes_to_read;
+  }
+
+  return true;
+}
+
+}  // namespace DiscIO
diff --git a/Source/Core/DiscIO/ScrubbedBlob.h b/Source/Core/DiscIO/ScrubbedBlob.h
new file mode 100644
index 0000000000..b98b78681e
--- /dev/null
+++ b/Source/Core/DiscIO/ScrubbedBlob.h
@@ -0,0 +1,37 @@
+// Copyright 2020 Dolphin Emulator Project
+// Licensed under GPLv2+
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include <memory>
+#include <string>
+
+#include "DiscIO/Blob.h"
+#include "DiscIO/DiscScrubber.h"
+
+namespace DiscIO
+{
+// This class wraps another BlobReader and zeroes out data that has been
+// identified by DiscScrubber as unused.
+class ScrubbedBlob : public BlobReader
+{
+public:
+  static std::unique_ptr<ScrubbedBlob> Create(const std::string& path);
+
+  BlobType GetBlobType() const override { return m_blob_reader->GetBlobType(); }
+  u64 GetRawSize() const override { return m_blob_reader->GetRawSize(); }
+  u64 GetDataSize() const override { return m_blob_reader->GetDataSize(); }
+  bool IsDataSizeAccurate() const override { return m_blob_reader->IsDataSizeAccurate(); }
+  u64 GetBlockSize() const override { return m_blob_reader->GetBlockSize(); }
+
+  bool Read(u64 offset, u64 size, u8* out_ptr) override;
+
+private:
+  ScrubbedBlob(std::unique_ptr<BlobReader> blob_reader, DiscScrubber scrubber);
+
+  std::unique_ptr<BlobReader> m_blob_reader;
+  DiscScrubber m_scrubber;
+};
+
+}  // namespace DiscIO
diff --git a/Source/Core/DolphinQt/ConvertDialog.cpp b/Source/Core/DolphinQt/ConvertDialog.cpp
index c3e102bd7c..f35569b96e 100644
--- a/Source/Core/DolphinQt/ConvertDialog.cpp
+++ b/Source/Core/DolphinQt/ConvertDialog.cpp
@@ -6,6 +6,7 @@
 
 #include <algorithm>
 #include <future>
+#include <memory>
 #include <utility>
 
 #include <QComboBox>
@@ -21,6 +22,7 @@
 
 #include "Common/Assert.h"
 #include "DiscIO/Blob.h"
+#include "DiscIO/ScrubbedBlob.h"
 #include "DolphinQt/QtUtils/ModalMessageBox.h"
 #include "DolphinQt/QtUtils/ParallelProgressDialog.h"
 #include "UICommon/GameFile.h"
@@ -179,34 +181,68 @@ void ConvertDialog::Convert()
           QFileInfo(QString::fromStdString(original_path)).fileName());
     }
 
-    std::future<bool> good;
+    std::unique_ptr<DiscIO::BlobReader> blob_reader;
+    bool scrub_current_file = scrub_wii && file->GetPlatform() == DiscIO::Platform::WiiDisc;
 
-    if (format == DiscIO::BlobType::PLAIN)
+    if (scrub_current_file)
     {
-      good = std::async(std::launch::async, [&] {
-        const bool good = DiscIO::ConvertToPlain(original_path, dst_path.toStdString(), &CompressCB,
-                                                 &progress_dialog);
-        progress_dialog.Reset();
-        return good;
-      });
+      blob_reader = DiscIO::ScrubbedBlob::Create(original_path);
+      if (!blob_reader)
+      {
+        const int result =
+            ModalMessageBox::warning(this, tr("Question"),
+                                     tr("Failed to remove junk data from file \"%1\".\n\n"
+                                        "Would you like to convert it without removing junk data?")
+                                         .arg(QString::fromStdString(original_path)),
+                                     QMessageBox::Ok | QMessageBox::Abort);
+
+        if (result == QMessageBox::Ok)
+          scrub_current_file = false;
+        else
+          return;
+      }
     }
-    else if (format == DiscIO::BlobType::GCZ)
+
+    if (!scrub_current_file)
+      blob_reader = DiscIO::CreateBlobReader(original_path);
+
+    if (!blob_reader)
     {
-      good = std::async(std::launch::async, [&] {
-        const bool good =
-            DiscIO::ConvertToGCZ(original_path, dst_path.toStdString(),
-                                 file->GetPlatform() == DiscIO::Platform::WiiDisc ? 1 : 0, 16384,
-                                 &CompressCB, &progress_dialog);
-        progress_dialog.Reset();
-        return good;
-      });
+      QErrorMessage(this).showMessage(
+          tr("Failed to open the input file \"%1\".").arg(QString::fromStdString(original_path)));
     }
-
-    progress_dialog.GetRaw()->exec();
-    if (!good.get())
+    else
     {
-      QErrorMessage(this).showMessage(tr("Dolphin failed to complete the requested action."));
-      return;
+      std::future<bool> good;
+
+      if (format == DiscIO::BlobType::PLAIN)
+      {
+        good = std::async(std::launch::async, [&] {
+          const bool good =
+              DiscIO::ConvertToPlain(blob_reader.get(), original_path, dst_path.toStdString(),
+                                     &CompressCB, &progress_dialog);
+          progress_dialog.Reset();
+          return good;
+        });
+      }
+      else if (format == DiscIO::BlobType::GCZ)
+      {
+        good = std::async(std::launch::async, [&] {
+          const bool good =
+              DiscIO::ConvertToGCZ(blob_reader.get(), original_path, dst_path.toStdString(),
+                                   file->GetPlatform() == DiscIO::Platform::WiiDisc ? 1 : 0, 16384,
+                                   &CompressCB, &progress_dialog);
+          progress_dialog.Reset();
+          return good;
+        });
+      }
+
+      progress_dialog.GetRaw()->exec();
+      if (!good.get())
+      {
+        QErrorMessage(this).showMessage(tr("Dolphin failed to complete the requested action."));
+        return;
+      }
     }
   }