1 files changed, 32 insertions, 27 deletions

diff --git a/Source/Core/DiscIO/WbfsBlob.cpp b/Source/Core/DiscIO/WbfsBlob.cpp
index 58ec6fc4c7..1bc7013aed 100644
--- a/Source/Core/DiscIO/WbfsBlob.cpp
+++ b/Source/Core/DiscIO/WbfsBlob.cpp
@@ -7,9 +7,11 @@
 #include <cstring>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "Common/Align.h"
+#include "Common/Assert.h"
 #include "Common/CommonFuncs.h"
 #include "Common/CommonTypes.h"
 #include "Common/FileUtil.h"
@@ -23,14 +25,15 @@ static const u64 WII_SECTOR_SIZE = 0x8000;
 static const u64 WII_SECTOR_COUNT = 143432 * 2;
 static const u64 WII_DISC_HEADER_SIZE = 256;
 
-WbfsFileReader::WbfsFileReader(const std::string& filename)
-    : m_total_files(0), m_size(0), m_good(true)
+WbfsFileReader::WbfsFileReader(File::IOFile file, const std::string& path)
+    : m_total_files(0), m_size(0), m_good(false)
 {
-  if (filename.length() < 4 || !OpenFiles(filename) || !ReadHeader())
-  {
-    m_good = false;
+  if (!AddFileToList(std::move(file)))
     return;
-  }
+  OpenAdditionalFiles(path);
+  if (!ReadHeader())
+    return;
+  m_good = true;
 
   // Grab disc info (assume slot 0, checked in ReadHeader())
   m_wlba_table.resize(m_blocks_per_disc);
@@ -50,38 +53,40 @@ u64 WbfsFileReader::GetDataSize() const
   return WII_SECTOR_COUNT * WII_SECTOR_SIZE;
 }
 
-bool WbfsFileReader::OpenFiles(const std::string& filename)
+void WbfsFileReader::OpenAdditionalFiles(const std::string& path)
 {
-  m_total_files = 0;
+  if (path.length() < 4)
+    return;
+
+  _assert_(m_total_files > 0);  // The code below gives .wbf0 for index 0, but it should be .wbfs
 
   while (true)
   {
-    auto new_entry = std::make_unique<file_entry>();
-
     // Replace last character with index (e.g. wbfs = wbf1)
-    std::string path = filename;
-    if (m_total_files != 0)
-    {
-      path[path.length() - 1] = '0' + m_total_files;
-    }
+    std::string current_path = path;
+    current_path.back() = '0' + m_total_files;
+    if (!AddFileToList(File::IOFile(current_path, "rb")))
+      return;
+  }
+}
 
-    if (!new_entry->file.Open(path, "rb"))
-    {
-      return m_total_files != 0;
-    }
+bool WbfsFileReader::AddFileToList(File::IOFile file)
+{
+  if (!file.IsOpen())
+    return false;
 
-    new_entry->base_address = m_size;
-    new_entry->size = new_entry->file.GetSize();
-    m_size += new_entry->size;
+  const u64 file_size = file.GetSize();
+  m_files.emplace_back(std::make_unique<file_entry>(std::move(file), m_size, file_size));
+  m_size += file_size;
+  m_total_files++;
 
-    m_total_files++;
-    m_files.emplace_back(std::move(new_entry));
-  }
+  return true;
 }
 
 bool WbfsFileReader::ReadHeader()
 {
   // Read hd size info
+  m_files[0]->file.Seek(0, SEEK_SET);
   m_files[0]->file.ReadBytes(&m_header, sizeof(WbfsHeader));
   if (m_header.magic != WBFS_MAGIC)
     return false;
@@ -164,9 +169,9 @@ File::IOFile& WbfsFileReader::SeekToCluster(u64 offset, u64* available)
   return m_files[0]->file;
 }
 
-std::unique_ptr<WbfsFileReader> WbfsFileReader::Create(const std::string& filename)
+std::unique_ptr<WbfsFileReader> WbfsFileReader::Create(File::IOFile file, const std::string& path)
 {
-  auto reader = std::unique_ptr<WbfsFileReader>(new WbfsFileReader(filename));
+  auto reader = std::unique_ptr<WbfsFileReader>(new WbfsFileReader(std::move(file), path));
 
   if (!reader->IsGood())
     reader.reset();