PKG: Implement Multi-threaded installation

rpcs3/Crypto/unpkg.cpp

@@ -4,10 +4,13 @@
 #include "key_vault.h"
 #include "util/logs.hpp"
 #include "Utilities/StrUtil.h"
+#include "Utilities/Thread.h"
+#include "Utilities/mutex.h"
 #include "Emu/System.h"
 #include "Emu/system_utils.hpp"
 #include "Emu/VFS.h"
 #include "unpkg.h"
+#include "util/sysinfo.hpp"
 #include "Loader/PSF.h"
 
 LOG_CHANNEL(pkg_log, "PKG");
@@ -189,8 +192,7 @@ bool package_reader::read_metadata()
 
 	// Read title ID and use it as an installation directory
 	m_install_dir.resize(9);
-	archive_seek(55);
+	archive_read_block(55, &m_install_dir.front(), m_install_dir.size());
-	archive_read(&m_install_dir.front(), m_install_dir.size());
 
 	// Read package metadata
 
@@ -509,7 +511,7 @@ bool package_reader::read_param_sfo()
 
 	std::vector<PKGEntry> entries(m_header.file_count);
 
-	std::memcpy(entries.data(), m_buf.get(), entries.size() * sizeof(PKGEntry));
+	std::memcpy(entries.data(), m_bufs.back().get(), entries.size() * sizeof(PKGEntry));
 
 	for (const auto& entry : entries)
 	{
@@ -523,7 +525,7 @@ bool package_reader::read_param_sfo()
 
 		decrypt(entry.name_offset, entry.name_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data());
 
-		const std::string_view name{reinterpret_cast<char*>(m_buf.get()), entry.name_size};
+		const std::string_view name{reinterpret_cast<char*>(m_bufs.back().get()), entry.name_size};
 
 		// We're looking for the PARAM.SFO file, if there is any
 		if (usz ndelim = name.find_first_not_of('/'); ndelim == umax || name.substr(ndelim) != "PARAM.SFO")
@@ -538,13 +540,13 @@ bool package_reader::read_param_sfo()
 			{
 				const u64 block_size = std::min<u64>(BUF_SIZE, entry.file_size - pos);
 
-				if (decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data()) != block_size)
+				if (decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data()).size() != block_size)
 				{
 					pkg_log.error("Failed to decrypt PARAM.SFO file");
 					return false;
 				}
 
-				if (tmp.write(m_buf.get(), block_size) != block_size)
+				if (tmp.write(m_bufs.back().get(), block_size) != block_size)
 				{
 					pkg_log.error("Failed to write to temporary PARAM.SFO file");
 					return false;
@@ -689,6 +691,159 @@ package_error package_reader::check_target_app_version() const
 
 fs::file DecryptEDAT(const fs::file& input, const std::string& input_file_name, int mode, u8 *custom_klic, bool verbose = false);
 
+usz package_reader::extract_worker(const std::string& dir, bool was_null, atomic_t<double>& sync, thread_key thread_data_key, atomic_t<usz>& entry_indexer, std::vector<PKGEntry>& entries)
+{
+	usz num_failures = 0;
+
+	for (usz entry_index = entry_indexer++; num_failures == 0 && entry_index < entries.size(); entry_index = entry_indexer++)
+	{
+		const auto& entry = ::at32(entries, entry_index);
+
+		if ((entry.type & 0xff) == PKG_FILE_ENTRY_FOLDER || (entry.type & 0xff) == 0x12u)
+		{
+			continue;
+		}
+
+		if (entry.name_size > PKG_MAX_FILENAME_SIZE)
+		{
+			num_failures++;
+			pkg_log.error("PKG name size is too big (0x%x)", entry.name_size);
+			continue;
+		}
+
+		const bool is_psp = (entry.type & PKG_FILE_ENTRY_PSP) != 0u;
+
+		std::span<const char> data_span = decrypt(entry.name_offset, entry.name_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), thread_data_key);
+
+		const std::string name{data_span.data(), entry.name_size};
+		const std::string path = dir + vfs::escape(name);
+
+		const bool log_error = entry.pad || (entry.type & ~PKG_FILE_ENTRY_KNOWN_BITS);
+		(log_error ? pkg_log.error : pkg_log.notice)("Entry 0x%08x: %s (pad=0x%x)", entry.type, name, entry.pad);
+
+		switch (const u8 entry_type = entry.type & 0xff)
+		{
+		case PKG_FILE_ENTRY_NPDRM:
+		case PKG_FILE_ENTRY_NPDRMEDAT:
+		case PKG_FILE_ENTRY_SDAT:
+		case PKG_FILE_ENTRY_REGULAR:
+		case PKG_FILE_ENTRY_UNK0:
+		case PKG_FILE_ENTRY_UNK1:
+		case 0xe:
+		case 0x10:
+		case 0x11:
+		case 0x13:
+		case 0x14:
+		case 0x15:
+		case 0x16:
+		case 0x18:
+		case 0x19:
+		{
+			const bool did_overwrite = fs::is_file(path);
+
+			if (did_overwrite && !(entry.type & PKG_FILE_ENTRY_OVERWRITE))
+			{
+				pkg_log.notice("Didn't overwrite %s", path);
+				break;
+			}
+
+			const bool is_buffered = entry_type == PKG_FILE_ENTRY_SDAT;
+
+			if (entry_type == PKG_FILE_ENTRY_NPDRMEDAT)
+			{
+				pkg_log.warning("NPDRM EDAT!");
+			}
+
+			if (fs::file out = is_buffered ? fs::make_stream<std::vector<u8>>() : fs::file{ path, did_overwrite ? fs::rewrite : fs::write_new })
+			{
+				bool extract_success = true;
+				for (u64 pos = 0; pos < entry.file_size; pos += BUF_SIZE)
+				{
+					const u64 block_size = std::min<u64>(BUF_SIZE, entry.file_size - pos);
+
+					data_span = decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), thread_data_key);
+
+					if (data_span.size() != block_size)
+					{
+						extract_success = false;
+						pkg_log.error("Failed to extract file %s", path);
+						break;
+					}
+
+					if (out.write(data_span.data(), block_size) != block_size)
+					{
+						extract_success = false;
+						pkg_log.error("Failed to write file %s", path);
+						break;
+					}
+
+					if (sync.fetch_add((block_size + 0.0) / m_header.data_size) < 0.)
+					{
+						// Cancel the installation
+						num_failures++;
+
+						if (was_null)
+						{
+							pkg_log.error("Package installation cancelled: %s", dir);
+							out.close();
+							return num_failures;
+						}
+					}
+				}
+
+				if (is_buffered)
+				{
+					out = DecryptEDAT(out, name, 1, reinterpret_cast<u8*>(&m_header.klicensee), true);
+					if (!out || !fs::write_file(path, fs::rewrite, static_cast<fs::container_stream<std::vector<u8>>*>(out.release().get())->obj))
+					{
+						num_failures++;
+						pkg_log.error("Failed to create file %s", path);
+						break;
+					}
+				}
+
+				if (extract_success)
+				{
+					if (did_overwrite)
+					{
+						pkg_log.warning("Overwritten file %s", path);
+					}
+					else
+					{
+						pkg_log.notice("Created file %s", path);
+
+						if (name == "USRDIR/EBOOT.BIN" && entry.file_size > 4)
+						{
+							// Expose the creation of a bootable file
+							m_bootable_file_path = path;
+						}
+					}
+				}
+				else
+				{
+					num_failures++;
+				}
+			}
+			else
+			{
+				num_failures++;
+				pkg_log.error("Failed to create file %s", path);
+			}
+
+			break;
+		}
+		default:
+		{
+			num_failures++;
+			pkg_log.error("Unknown PKG entry type (0x%x) %s", entry.type, name);
+			break;
+		}
+		}
+	}
+
+	return num_failures;
+}
+
 bool package_reader::extract_data(atomic_t<double>& sync)
 {
 	if (!m_is_valid)
@@ -744,20 +899,20 @@ bool package_reader::extract_data(atomic_t<double>& sync)
 		return false;
 	}
 
-	usz num_failures = 0;
+	atomic_t<usz> num_failures = 0;
 
 	std::vector<PKGEntry> entries(m_header.file_count);
 
-	std::memcpy(entries.data(), m_buf.get(), entries.size() * sizeof(PKGEntry));
+	std::memcpy(entries.data(), m_bufs.back().get(), entries.size() * sizeof(PKGEntry));
 
 	// Create directories first
 	for (const auto& entry : entries)
 	{
-		if (entry.name_size > 256)
+		if (entry.name_size > PKG_MAX_FILENAME_SIZE)
 		{
 			num_failures++;
 			pkg_log.error("PKG name size is too big (0x%x)", entry.name_size);
-			continue;
+			break;
 		}
 
 		switch (const u8 entry_type = entry.type & 0xff)
@@ -768,7 +923,7 @@ bool package_reader::extract_data(atomic_t<double>& sync)
 			const bool is_psp = (entry.type & PKG_FILE_ENTRY_PSP) != 0u;
 			decrypt(entry.name_offset, entry.name_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data());
 
-			const std::string name{reinterpret_cast<char*>(m_buf.get()), entry.name_size};
+			const std::string name{reinterpret_cast<char*>(m_bufs.back().get()), entry.name_size};
 			const std::string path = dir + vfs::escape(name);
 
 			const bool log_error = entry.pad || (entry.type & ~PKG_FILE_ENTRY_KNOWN_BITS);
@@ -787,11 +942,11 @@ bool package_reader::extract_data(atomic_t<double>& sync)
 			{
 				num_failures++;
 				pkg_log.error("Failed to create directory %s", path);
+				break;
 			}
 
 			break;
 		}
-
 		default:
 		{
 			continue;
@@ -810,147 +965,19 @@ bool package_reader::extract_data(atomic_t<double>& sync)
 		return false;
 	}
 
-	for (const auto& entry : entries)
+	atomic_t<usz> entry_indexer = 0;
+	atomic_t<usz> thread_indexer = 0;
+
+	m_bufs.resize(std::min<usz>(utils::get_thread_count(), entries.size()));
+
+	named_thread_group workers("PKG Installer "sv, std::max<usz>(m_bufs.size(), 1) - 1, [&]()
 	{
-		if ((entry.type & 0xff) == PKG_FILE_ENTRY_FOLDER || (entry.type & 0xff) == 0x12)
+		num_failures += extract_worker(dir, was_null, sync, thread_key{thread_indexer++}, entry_indexer, entries);
-		{
+	});
-			continue;
-		}
 
-		if (entry.name_size > 256)
+	num_failures += extract_worker(dir, was_null, sync, thread_key{thread_indexer++}, entry_indexer, entries);
-		{
-			num_failures++;
-			pkg_log.error("PKG name size is too big (0x%x)", entry.name_size);
-			continue;
-		}
 
-		const bool is_psp = (entry.type & PKG_FILE_ENTRY_PSP) != 0u;
+	workers.join();
-
-		decrypt(entry.name_offset, entry.name_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data());
-
-		const std::string name{reinterpret_cast<char*>(m_buf.get()), entry.name_size};
-		const std::string path = dir + vfs::escape(name);
-
-		const bool log_error = entry.pad || (entry.type & ~PKG_FILE_ENTRY_KNOWN_BITS);
-		(log_error ? pkg_log.error : pkg_log.notice)("Entry 0x%08x: %s (pad=0x%x)", entry.type, name, entry.pad);
-
-		switch (const u8 entry_type = entry.type & 0xff)
-		{
-		case PKG_FILE_ENTRY_NPDRM:
-		case PKG_FILE_ENTRY_NPDRMEDAT:
-		case PKG_FILE_ENTRY_SDAT:
-		case PKG_FILE_ENTRY_REGULAR:
-		case PKG_FILE_ENTRY_UNK0:
-		case PKG_FILE_ENTRY_UNK1:
-		case 0xe:
-		case 0x10:
-		case 0x11:
-		case 0x13:
-		case 0x14:
-		case 0x15:
-		case 0x16:
-		case 0x18:
-		case 0x19:
-		{
-			const bool did_overwrite = fs::is_file(path);
-
-			if (did_overwrite && !(entry.type & PKG_FILE_ENTRY_OVERWRITE))
-			{
-				pkg_log.notice("Didn't overwrite %s", path);
-				break;
-			}
-
-			const bool is_buffered = entry_type == PKG_FILE_ENTRY_SDAT;
-
-			if (entry_type == PKG_FILE_ENTRY_NPDRMEDAT)
-			{
-				pkg_log.todo("NPDRM EDAT!");
-			}
-
-			if (fs::file out = is_buffered ? fs::make_stream<std::vector<u8>>() : fs::file{ path, fs::rewrite })
-			{
-				bool extract_success = true;
-				for (u64 pos = 0; pos < entry.file_size; pos += BUF_SIZE)
-				{
-					const u64 block_size = std::min<u64>(BUF_SIZE, entry.file_size - pos);
-
-					if (decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data()) != block_size)
-					{
-						extract_success = false;
-						pkg_log.error("Failed to extract file %s", path);
-						break;
-					}
-
-					if (out.write(m_buf.get(), block_size) != block_size)
-					{
-						extract_success = false;
-						pkg_log.error("Failed to write file %s", path);
-						break;
-					}
-
-					if (sync.fetch_add((block_size + 0.0) / m_header.data_size) < 0.)
-					{
-						if (was_null)
-						{
-							pkg_log.error("Package installation cancelled: %s", dir);
-							out.close();
-							fs::remove_all(dir, true);
-							return false;
-						}
-
-						// Cannot cancel the installation
-						sync += 1.;
-					}
-				}
-
-				if (is_buffered)
-				{
-					out = DecryptEDAT(out, name, 1, reinterpret_cast<u8*>(&m_header.klicensee), true);
-					if (!out || !fs::write_file(path, fs::rewrite, static_cast<fs::container_stream<std::vector<u8>>*>(out.release().get())->obj))
-					{
-						num_failures++;
-						pkg_log.error("Failed to create file %s", path);
-						break;
-					}
-				}
-
-				if (extract_success)
-				{
-					if (did_overwrite)
-					{
-						pkg_log.warning("Overwritten file %s", path);
-					}
-					else
-					{
-						pkg_log.notice("Created file %s", path);
-
-						if (name == "USRDIR/EBOOT.BIN" && entry.file_size > 4)
-						{
-							// Expose the creation of a bootable file
-							m_bootable_file_path = path;
-						}
-					}
-				}
-				else
-				{
-					num_failures++;
-				}
-			}
-			else
-			{
-				num_failures++;
-				pkg_log.error("Failed to create file %s", path);
-			}
-
-			break;
-		}
-		default:
-		{
-			num_failures++;
-			pkg_log.error("Unknown PKG entry type (0x%x) %s", entry.type, name);
-		}
-		}
-	}
 
 	if (num_failures == 0)
 	{
@@ -966,6 +993,7 @@ bool package_reader::extract_data(atomic_t<double>& sync)
 		pkg_log.error("Package installation failed: %s", dir);
 	}
 
+	m_bufs.clear();
 	return num_failures == 0;
 }
 
@@ -977,28 +1005,42 @@ void package_reader::archive_seek(const s64 new_offset, const fs::seek_mode damo
 u64 package_reader::archive_read(void* data_ptr, const u64 num_bytes)
 {
 	return m_file ? m_file.read(data_ptr, num_bytes) : 0;
-};
+}
 
-u64 package_reader::decrypt(u64 offset, u64 size, const uchar* key)
+std::span<const char> package_reader::archive_read_block(u64 offset, void* data_ptr, u64 num_bytes)
+{
+	const usz read_n = m_file.read_at(offset, data_ptr, num_bytes);
+
+	return {static_cast<const char*>(data_ptr), read_n};
+}
+
+std::span<const char> package_reader::decrypt(u64 offset, u64 size, const uchar* key, thread_key thread_data_key)
 {
 	if (!m_is_valid)
 	{
-		return 0;
+		return {};
 	}
 
-	if (!m_buf)
+	if (m_bufs.empty())
+	{
+		// Assume in single-threaded mode still
+		m_bufs.resize(1);
+	}
+
+	auto& local_buf = ::at32(m_bufs, thread_data_key.unique_num);
+
+	if (!local_buf)
 	{
 		// Allocate buffer with BUF_SIZE size or more if required
-		m_buf.reset(new u128[std::max<u64>(BUF_SIZE, sizeof(PKGEntry) * m_header.file_count) / sizeof(u128)]);
+		local_buf.reset(new u128[std::max<u64>(BUF_SIZE, sizeof(PKGEntry) * m_header.file_count) / sizeof(u128)]);
 	}
 
-	archive_seek(m_header.data_offset + offset);
-
 	// Read the data and set available size
-	const u64 read = archive_read(m_buf.get(), size);
+	const auto data_span = archive_read_block(m_header.data_offset + offset, local_buf.get(), size);
+	ensure(data_span.data() == static_cast<void*>(local_buf.get()));
 
 	// Get block count
-	const u64 blocks = (read + 15) / 16;
+	const u64 blocks = (data_span.size() + 15) / 16;
 
 	if (m_header.pkg_type == PKG_RELEASE_TYPE_DEBUG)
 	{
@@ -1024,7 +1066,7 @@ u64 package_reader::decrypt(u64 offset, u64 size, const uchar* key)
 
 			sha1(reinterpret_cast<const u8*>(input), sizeof(input), hash.data);
 
-			m_buf[i] ^= hash._v128;
+			local_buf[i] ^= hash._v128;
 		}
 	}
 	else if (m_header.pkg_type == PKG_RELEASE_TYPE_RELEASE)
@@ -1044,7 +1086,7 @@ u64 package_reader::decrypt(u64 offset, u64 size, const uchar* key)
 
 			aes_crypt_ecb(&ctx, AES_ENCRYPT, reinterpret_cast<const u8*>(&input), reinterpret_cast<u8*>(&key));
 
-			m_buf[i] ^= key;
+			local_buf[i] ^= key;
 		}
 	}
 	else
@@ -1053,5 +1095,5 @@ u64 package_reader::decrypt(u64 offset, u64 size, const uchar* key)
 	}
 
 	// Return the amount of data written in buf
-	return read;
+	return data_span;
 };

rpcs3/Crypto/unpkg.h

@@ -6,12 +6,14 @@
 #include "Utilities/File.h"
 #include <sstream>
 #include <iomanip>
+#include <span>
 
 // Constants
-enum
+enum : u32
 {
 	PKG_HEADER_SIZE  = 0xC0, // sizeof(pkg_header) + sizeof(pkg_unk_checksum)
 	PKG_HEADER_SIZE2 = 0x280,
+	PKG_MAX_FILENAME_SIZE = 256,
 };
 
 enum : u16
@@ -317,14 +319,21 @@ public:
 		return m_bootable_file_path;
 	}
 
+	struct thread_key
+	{
+		const usz unique_num = umax;
+	};
+
 private:
 	bool read_header();
 	bool read_metadata();
 	bool read_param_sfo();
 	bool decrypt_data();
-	void archive_seek(const s64 new_offset, const fs::seek_mode damode = fs::seek_set);
+	void archive_seek(s64 new_offset, const fs::seek_mode damode = fs::seek_set);
-	u64 archive_read(void* data_ptr, const u64 num_bytes);
+	u64 archive_read(void* data_ptr, u64 num_bytes);
-	u64 decrypt(u64 offset, u64 size, const uchar* key);
+	std::span<const char> archive_read_block(u64 offset, void* data_ptr, u64 num_bytes);
+	std::span<const char> decrypt(u64 offset, u64 size, const uchar* key, thread_key thread_data_key = {0});
+	usz extract_worker(const std::string& dir, bool was_null, atomic_t<double>& sync, thread_key thread_data_key, atomic_t<usz>& entry_indexer, std::vector<PKGEntry>& entries);
 
 	const usz BUF_SIZE = 8192 * 1024; // 8 MB
 
@@ -333,7 +342,7 @@ private:
 	std::string m_path{};
 	std::string m_install_dir{};
 	fs::file m_file{};
-	std::unique_ptr<u128[]> m_buf{};
+	std::vector<std::unique_ptr<u128[]>> m_bufs{};
 	std::array<uchar, 16> m_dec_key{};
 
 	PKGHeader m_header{};

rpcs3/rpcs3qt/main_window.cpp

@@ -921,6 +921,7 @@ void main_window::HandlePackageInstallation(QStringList file_paths)
 			// Update progress window
 			double pval = progress;
 			if (pval < 0.) pval += 1.;
+
 			pdlg.SetValue(static_cast<int>(pval * pdlg.maximum()));
 			QCoreApplication::processEvents();
 		}