Skip to content

Commit

Permalink
PKG: Refactor and proper support for IDU
Browse files Browse the repository at this point in the history
  • Loading branch information
elad335 committed Feb 22, 2025
1 parent 1df4568 commit 38b5c84
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 89 deletions.
208 changes: 127 additions & 81 deletions rpcs3/Crypto/unpkg.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@

LOG_CHANNEL(pkg_log, "PKG");

#pragma optimize("", off)

package_reader::package_reader(const std::string& path)
: m_path(path)
{
Expand Down Expand Up @@ -504,45 +506,85 @@ bool package_reader::set_decryption_key()
return true;
}

switch (m_metadata.package_revision.data.make_package_npdrm_ver[0])
{
case 0x15:
std::memcpy(m_dec_key.data(), PKG_AES_KEY, m_dec_key.size());

if (std::vector<PKGEntry> entries; !read_entries(entries))
{
if (!!(read_from_ptr<u128>(m_metadata.qa_digest) | read_from_ptr<u64>(m_metadata.qa_digest + 16)))
{
pkg_log.error("IDU PKG contains QA Digest!");
}
pkg_log.notice("PKG may be IDU, retrying with IDU key.");

std::memcpy(m_dec_key.data(), PKG_AES_KEY_IDU, m_dec_key.size());
break;
}
default:
{
pkg_log.error("Unknown NPDRM package version: 0x%x", read_from_ptr<be_t<u16>>(m_metadata.package_revision.data.make_package_npdrm_ver));
[[fallthrough]];
}
case 0x19:
{
std::memcpy(m_dec_key.data(), PKG_AES_KEY, m_dec_key.size());
break;
}

if (!read_entries(entries))
{
pkg_log.error("PKG decryption failed!");
return false;
}
}

return true;
}

bool package_reader::read_entries(std::vector<PKGEntry>& entries)
{
const std::span<const char> data = decrypt(0, m_header.file_count * sizeof(PKGEntry), m_dec_key.data());
entries.clear();
entries.resize(m_header.file_count);

const usz read_size = decrypt(0, m_header.file_count * sizeof(PKGEntry), m_dec_key.data(), entries.data());

if (data.size() < m_header.file_count * sizeof(PKGEntry))
if (read_size < m_header.file_count * sizeof(PKGEntry))
{
return false;
}

entries.resize(m_header.file_count);
const usz fsz = m_file.size();

// Data integrity validation
for (const PKGEntry& entry : entries)
{
if (entry.name_size == 0)
{
continue;
}

if (entry.name_size > 256)
{
return false;
}

if (fsz < entry.name_size || fsz - entry.name_size < entry.name_offset)
{
// Name exceeds file(s)
return false;
}

if (entry.name_offset < sizeof(PKGHeader))
{
// Data points to header
return false;
}

if (entry.file_size)
{
if (fsz < entry.file_size || fsz - entry.file_size < entry.file_offset)
{
// Data exceeds file(s)
return false;
}

if (entry.name_offset == entry.file_offset)
{
// Repeated value: odd
return false;
}

if (entry.file_offset < sizeof(PKGHeader))
{
// Data points to header
return false;
}
}
}

std::memcpy(entries.data(), data.data(), data.size());
return true;
}

Expand All @@ -555,19 +597,27 @@ bool package_reader::read_param_sfo()
return false;
}

std::vector<u8> data_buf;

for (const PKGEntry& entry : entries)
{
if (entry.name_size > 256)
{
pkg_log.error("PKG name size is too big (0x%x)", entry.name_size);
pkg_log.error("PKG name size is too big (size=0x%x, offset=0x%x)", entry.name_size, entry.name_offset);
continue;
}

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());
std::string name(entry.name_size + BUF_PADDING, '\0');

if (usz read_size = decrypt(entry.name_offset, entry.name_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), name.data()); read_size < entry.name_size)
{
pkg_log.error("PKG name could not be read (size=0x%x, offset=0x%x)", entry.name_size, entry.name_offset);
continue;
}

const std::string_view name{reinterpret_cast<char*>(m_bufs.back().get()), entry.name_size};
fmt::trim_back(name, "\0"sv);

// 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")
Expand All @@ -582,13 +632,15 @@ 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()).size() != block_size)
data_buf.resize(block_size + BUF_PADDING);

if (decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), data_buf.data()) != block_size)
{
pkg_log.error("Failed to decrypt PARAM.SFO file");
return false;
}

if (tmp.write(m_bufs.back().get(), block_size) != block_size)
if (tmp.write(data_buf.data(), block_size) != block_size)
{
pkg_log.error("Failed to write to temporary PARAM.SFO file");
return false;
Expand Down Expand Up @@ -812,14 +864,23 @@ bool package_reader::fill_data(std::map<std::string, install_entry*>& all_instal
if (entry.name_size > PKG_MAX_FILENAME_SIZE)
{
num_failures++;
pkg_log.error("PKG name size is too big (0x%x)", entry.name_size);
pkg_log.error("PKG name size is too big (size=0x%x, offset=0x%x)", entry.name_size, entry.name_offset);
break;
}

std::string name(entry.name_size + BUF_PADDING, '\0');

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_view name{reinterpret_cast<char*>(m_bufs.back().get()), entry.name_size};
if (const usz read_size = decrypt(entry.name_offset, entry.name_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), name.data()); read_size < entry.name_size)
{
num_failures++;
pkg_log.error("PKG name could not be read (size=0x%x, offset=0x%x)", entry.name_size, entry.name_offset);
break;
}

fmt::trim_back(name, "\0"sv);

std::string path = m_install_path + vfs::escape(name);

if (entry.pad || (entry.type & ~PKG_FILE_ENTRY_KNOWN_BITS))
Expand Down Expand Up @@ -901,7 +962,7 @@ bool package_reader::fill_data(std::map<std::string, install_entry*>& all_instal

fs::file DecryptEDAT(const fs::file& input, const std::string& input_file_name, int mode, u8 *custom_klic);

void package_reader::extract_worker(thread_key thread_data_key)
void package_reader::extract_worker()
{
std::vector<u8> read_cache;

Expand Down Expand Up @@ -1097,41 +1158,37 @@ void package_reader::extract_worker(thread_key thread_data_key)
{
const u64 block_size = std::min<u64>({BUF_SIZE, std::max<u64>(size * 5 / 3, 65536), entry.file_size - pos});

read_cache.resize(block_size);
read_cache.resize(block_size + BUF_PADDING);
cache_off = pos;

const std::span<const char> data_span = decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), thread_data_key);
const usz advance_size = decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), read_cache.data());

if (data_span.empty())
if (!advance_size)
{
cache_off = umax;
read_cache.clear();
return 0;
}

read_cache.resize(data_span.size());
std::memcpy(read_cache.data(), data_span.data(), data_span.size());
read_cache.resize(advance_size);

size = std::min<usz>(data_span.size(), size);
std::memcpy(ptr, data_span.data(), size);
size = std::min<usz>(advance_size, size);
std::memcpy(ptr, read_cache.data(), size);
return size;
}

while (read_size < size)
{
const u64 block_size = std::min<u64>(BUF_SIZE, size - read_size);

const std::span<const char> data_span = decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), thread_data_key);
const usz advance_size = decrypt(entry.file_offset + pos, block_size, is_psp ? PKG_AES_KEY2 : m_dec_key.data(), static_cast<u8*>(ptr) + read_size);

if (data_span.empty())
if (!advance_size)
{
break;
}

std::memcpy(static_cast<u8*>(ptr) + read_size, data_span.data(), data_span.size());

read_size += data_span.size();
pos += data_span.size();
read_size += advance_size;
pos += advance_size;
}

return read_size + size_cache_end;
Expand Down Expand Up @@ -1159,9 +1216,9 @@ void package_reader::extract_worker(thread_key thread_data_key)
}

// 16MB buffer
std::vector<u8> buffer(std::min<usz>(entry.file_size, 1u << 24));
std::vector<u8> buffer(std::min<usz>(entry.file_size, 1u << 24) + BUF_PADDING);

while (usz read_size = final_data.read(buffer.data(), buffer.size()))
while (usz read_size = final_data.read(buffer.data(), buffer.size() - BUF_PADDING))
{
out.write(buffer.data(), read_size);
m_written_bytes += read_size;
Expand Down Expand Up @@ -1262,20 +1319,15 @@ package_install_result package_reader::extract_data(std::deque<package_reader>&

if (reader.m_num_failures == 0)
{
reader.m_bufs.resize(std::min<usz>(utils::get_thread_count(), reader.m_install_entries.size()));

atomic_t<usz> thread_indexer = 0;
const usz thread_count = std::min<usz>(utils::get_thread_count(), reader.m_install_entries.size());

named_thread_group workers("PKG Installer "sv, std::max<u32>(::narrow<u32>(reader.m_bufs.size()), 1) - 1, [&]()
named_thread_group workers("PKG Installer "sv, std::max<u32>(::narrow<u32>(thread_count), 1) - 1, [&]()
{
reader.extract_worker(thread_key{thread_indexer++});
reader.extract_worker();
});

reader.extract_worker(thread_key{thread_indexer++});
reader.extract_worker();
workers.join();

reader.m_bufs.clear();
reader.m_bufs.shrink_to_fit();
}

num_failures += reader.m_num_failures;
Expand Down Expand Up @@ -1350,33 +1402,20 @@ std::span<const char> package_reader::archive_read_block(u64 offset, void* data_
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)
usz package_reader::decrypt(u64 offset, u64 size, const uchar* key, void* local_buf)
{
if (!m_is_valid)
{
return {};
}

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
local_buf.reset(new u128[std::max<u64>(BUF_SIZE, sizeof(PKGEntry) * m_header.file_count) / sizeof(u128)]);
return 0;
}

// Read the data and set available 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()));
const auto data_span = archive_read_block(m_header.data_offset + offset, local_buf, size);
ensure(data_span.data() == static_cast<void*>(local_buf));

// Get block count
const u64 blocks = (data_span.size() + 15) / 16;
const auto out_data = reinterpret_cast<u8*>(local_buf);

if (m_header.pkg_type == PKG_RELEASE_TYPE_DEBUG)
{
Expand All @@ -1394,15 +1433,15 @@ std::span<const char> package_reader::decrypt(u64 offset, u64 size, const uchar*
// Initialize stream cipher for current position
input[7] = offset / 16 + i;

union sha1_hash
struct sha1_hash
{
u8 data[20];
u128 _v128;
} hash;
} hash{};

sha1(reinterpret_cast<const u8*>(input), sizeof(input), hash.data);

local_buf[i] ^= hash._v128;
const u128 v = read_from_ptr<u128>(out_data, i * 16);
write_to_ptr<u128>(out_data, i * 16, v ^ read_from_ptr<u128>(hash.data));
}
}
else if (m_header.pkg_type == PKG_RELEASE_TYPE_RELEASE)
Expand All @@ -1422,16 +1461,23 @@ std::span<const char> package_reader::decrypt(u64 offset, u64 size, const uchar*

aes_crypt_ecb(&ctx, AES_ENCRYPT, reinterpret_cast<const u8*>(&input), reinterpret_cast<u8*>(&key));

local_buf[i] ^= key;
const u128 v = read_from_ptr<u128>(out_data, i * 16);
write_to_ptr<u128>(out_data, i * 16, v ^ key);
}
}
else
{
pkg_log.error("Unknown release type (0x%x)", m_header.pkg_type);
}

if (blocks * 16 != size)
{
// Put NTS and other zeroes on unaligned reads
std::memset(out_data + size, 0, blocks * 16 - size);
}

// Return the amount of data written in buf
return data_span;
return std::min<usz>(size, data_span.size());
}

int package_reader::get_progress(int maximum) const
Expand Down
Loading

0 comments on commit 38b5c84

Please sign in to comment.