DiscIO: Move scrubbing code out of ConvertToGCZ

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

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,
+bool ConvertToGCZ(BlobReader* infile, const std::string& infile_path,
-                  int sector_size = 16384, CompressCB callback = nullptr, void* arg = nullptr);
+                  const std::string& outfile_path, u32 sub_type, int sector_size = 16384,
-bool ConvertToPlain(const std::string& infile_path, const std::string& outfile_path,
                   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

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

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,
+bool ConvertToGCZ(BlobReader* infile, const std::string& infile_path,
-                  int block_size, CompressCB callback, void* arg)
+                  const std::string& outfile_path, u32 sub_type, int block_size,
+                  CompressCB callback, void* arg)
 {
-  bool scrubbing = false;
+  ASSERT(infile->IsDataSizeAccurate());
-
-  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());
 
   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) ?
+    const u64 bytes_to_read = std::min<u64>(block_size, header.data_size - inpos);
-                                  0 :
-                                  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());

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" />

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" />

Source/Core/DiscIO/DiscScrubber.cpp

@@ -25,8 +25,6 @@
 
 namespace DiscIO
 {
-constexpr size_t CLUSTER_SIZE = 0x8000;
-
 DiscScrubber::DiscScrubber() = default;
 DiscScrubber::~DiscScrubber() = default;
 

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);

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,
+bool ConvertToPlain(BlobReader* infile, const std::string& infile_path,
-                    CompressCB callback, void* arg)
+                    const std::string& outfile_path, CompressCB callback, void* arg)
 {
-  std::unique_ptr<BlobReader> reader = CreateBlobReader(infile_path);
+  ASSERT(infile->IsDataSizeAccurate());
-  if (!reader)
-  {
-    PanicAlertT("Failed to open the input file \"%s\".", infile_path.c_str());
-    return false;
-  }
-
-  ASSERT(reader->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);
+    const u64 sz = std::min(buffer_size, infile->GetDataSize() - inpos);
-    if (!reader->Read(inpos, sz, buffer.data()))
+    if (!infile->Read(inpos, sz, buffer.data()))
     {
       PanicAlertT("Failed to read from the input file \"%s\".", infile_path.c_str());
       success = false;

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

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

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,13 +181,46 @@ void ConvertDialog::Convert()
           QFileInfo(QString::fromStdString(original_path)).fileName());
     }
 
+    std::unique_ptr<DiscIO::BlobReader> blob_reader;
+    bool scrub_current_file = scrub_wii && file->GetPlatform() == DiscIO::Platform::WiiDisc;
+
+    if (scrub_current_file)
+    {
+      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;
+      }
+    }
+
+    if (!scrub_current_file)
+      blob_reader = DiscIO::CreateBlobReader(original_path);
+
+    if (!blob_reader)
+    {
+      QErrorMessage(this).showMessage(
+          tr("Failed to open the input file \"%1\".").arg(QString::fromStdString(original_path)));
+    }
+    else
+    {
       std::future<bool> good;
 
       if (format == DiscIO::BlobType::PLAIN)
       {
         good = std::async(std::launch::async, [&] {
-        const bool good = DiscIO::ConvertToPlain(original_path, dst_path.toStdString(), &CompressCB,
+          const bool good =
-                                                 &progress_dialog);
+              DiscIO::ConvertToPlain(blob_reader.get(), original_path, dst_path.toStdString(),
+                                     &CompressCB, &progress_dialog);
           progress_dialog.Reset();
           return good;
         });
@@ -194,7 +229,7 @@ void ConvertDialog::Convert()
       {
         good = std::async(std::launch::async, [&] {
           const bool good =
-            DiscIO::ConvertToGCZ(original_path, dst_path.toStdString(),
+              DiscIO::ConvertToGCZ(blob_reader.get(), original_path, dst_path.toStdString(),
                                    file->GetPlatform() == DiscIO::Platform::WiiDisc ? 1 : 0, 16384,
                                    &CompressCB, &progress_dialog);
           progress_dialog.Reset();
@@ -209,6 +244,7 @@ void ConvertDialog::Convert()
         return;
       }
     }
+  }
 
   ModalMessageBox::information(this, tr("Success"),
                                tr("Successfully converted %n image(s).", "", m_files.size()));