diff --git a/Source/Core/DiscIO/CMakeLists.txt b/Source/Core/DiscIO/CMakeLists.txt index 7dceb0b62453..22bc30185a68 100644 --- a/Source/Core/DiscIO/CMakeLists.txt +++ b/Source/Core/DiscIO/CMakeLists.txt @@ -48,5 +48,7 @@ target_link_libraries(discio PRIVATE minizip pugixml + BZip2::BZip2 + LibLZMA::LibLZMA ZLIB::ZLIB ) diff --git a/Source/Core/DiscIO/DiscIO.vcxproj b/Source/Core/DiscIO/DiscIO.vcxproj index f4e656986412..d6cab5e46607 100644 --- a/Source/Core/DiscIO/DiscIO.vcxproj +++ b/Source/Core/DiscIO/DiscIO.vcxproj @@ -107,6 +107,12 @@ {38fee76f-f347-484b-949c-b4649381cffb} + + {055a775f-b4f5-4970-9240-f6cf7661f37b} + + + {1d8c51d2-ffa4-418e-b183-9f42b6a6717e} + diff --git a/Source/Core/DiscIO/WIABlob.cpp b/Source/Core/DiscIO/WIABlob.cpp index 8caf0caf755a..dd9eb66d557a 100644 --- a/Source/Core/DiscIO/WIABlob.cpp +++ b/Source/Core/DiscIO/WIABlob.cpp @@ -9,6 +9,9 @@ #include #include +#include +#include + #include "Common/Align.h" #include "Common/CommonTypes.h" #include "Common/File.h" @@ -54,9 +57,17 @@ bool WIAFileReader::Initialize(const std::string& path) if (!m_file.ReadArray(&m_header_2, 1)) return false; + u8 compressor_data_size; + if (!m_file.ReadArray(&compressor_data_size, 1)) + return false; + + m_compressor_data.resize(compressor_data_size); + if (!m_file.ReadBytes(m_compressor_data.data(), m_compressor_data.size())) + return false; + const u32 compression_type = Common::swap32(m_header_2.compression_type); m_compression_type = static_cast(compression_type); - if (m_compression_type > CompressionType::Purge) + if (m_compression_type > CompressionType::LZMA2) { ERROR_LOG(DISCIO, "Unsupported WIA compression type %u in %s", compression_type, path.c_str()); return false; @@ -298,6 +309,58 @@ bool WIAFileReader::ReadCompressedData(u32 decompressed_data_size, u64 data_offs return true; } + + case CompressionType::Bzip2: + case CompressionType::LZMA: + case CompressionType::LZMA2: + { + std::vector compressed_data(data_size); + if (!m_file.Seek(data_offset, SEEK_SET) || !m_file.ReadBytes(compressed_data.data(), data_size)) + return false; + + std::unique_ptr decompressor; + switch (m_compression_type) + { + case CompressionType::Bzip2: + decompressor = std::make_unique(); + break; + case CompressionType::LZMA: + decompressor = std::make_unique(false, m_compressor_data); + break; + case CompressionType::LZMA2: + decompressor = std::make_unique(true, m_compressor_data); + break; + } + + if (!decompressor->Start(compressed_data.data(), compressed_data.size())) + return false; + + if (exception_list) + { + u16 exceptions; + if (decompressor->Read(reinterpret_cast(&exceptions), sizeof(exceptions)) != + sizeof(exceptions)) + { + return false; + } + + std::vector exception_entries(Common::swap16(exceptions)); + u8* exceptions_data = reinterpret_cast(exception_entries.data()); + const size_t exceptions_size = exception_entries.size() * sizeof(HashExceptionEntry); + if (decompressor->Read(exceptions_data, exceptions_size) != exceptions_size) + return false; + + // TODO: Actually handle the exceptions + } + + if (decompressor->Read(out_ptr, decompressed_data_size) != decompressed_data_size) + return false; + + if (!decompressor->DoneReading()) + return false; + + return true; + } } return false; @@ -369,4 +432,144 @@ std::string WIAFileReader::VersionToString(u32 version) else return StringFromFormat("%u.%02x.%02x.beta%u", a, b, c, d); } + +WIAFileReader::Decompressor::~Decompressor() = default; + +WIAFileReader::Bzip2Decompressor::~Bzip2Decompressor() +{ + End(); +} + +bool WIAFileReader::Bzip2Decompressor::Start(const u8* in_ptr, u64 size) +{ + if (m_started) + return false; + + m_stream.bzalloc = nullptr; + m_stream.bzfree = nullptr; + m_stream.opaque = nullptr; + + m_started = BZ2_bzDecompressInit(&m_stream, 0, 0) == BZ_OK; + + m_stream.next_in = reinterpret_cast(const_cast(in_ptr)); + m_stream.avail_in = size; + + return m_started; +} + +u64 WIAFileReader::Bzip2Decompressor::Read(u8* out_ptr, u64 size) +{ + if (!m_started || m_error_occurred || m_stream.avail_in == 0) + return 0; + + m_stream.next_out = reinterpret_cast(out_ptr); + m_stream.avail_out = size; + + const int result = BZ2_bzDecompress(&m_stream); + m_error_occurred = result != BZ_OK && result != BZ_STREAM_END; + + return m_error_occurred ? 0 : m_stream.next_out - reinterpret_cast(out_ptr); +} + +bool WIAFileReader::Bzip2Decompressor::DoneReading() const +{ + return m_started && !m_error_occurred && m_stream.avail_in == 0; +} + +void WIAFileReader::Bzip2Decompressor::End() +{ + if (m_started && !m_ended) + { + BZ2_bzDecompressEnd(&m_stream); + m_ended = true; + } +} + +WIAFileReader::LZMADecompressor::LZMADecompressor(bool lzma2, const std::vector& filter_options) +{ + m_options.preset_dict = nullptr; + + if (!lzma2 && filter_options.size() == 5) + { + // Little endian + m_options.dict_size = *reinterpret_cast(filter_options.data() + 1); + + const u8 d = filter_options[0]; + if (d >= (9 * 5 * 5)) + { + m_error_occurred = true; + } + else + { + m_options.lc = d % 9; + const u8 e = d / 9; + m_options.pb = e / 5; + m_options.lp = e % 5; + } + } + else if (lzma2 && filter_options.size() == 1) + { + const u8 d = filter_options[0]; + if (d > 40) + m_error_occurred = true; + else + m_options.dict_size = d == 40 ? 0xFFFFFFFF : (static_cast(2) | (d & 1)) << (d / 2 + 11); + } + else + { + m_error_occurred = true; + } + + m_filters[0].id = lzma2 ? LZMA_FILTER_LZMA2 : LZMA_FILTER_LZMA1; + m_filters[0].options = &m_options; + m_filters[1].id = LZMA_VLI_UNKNOWN; + m_filters[1].options = nullptr; +} + +WIAFileReader::LZMADecompressor::~LZMADecompressor() +{ + End(); +} + +bool WIAFileReader::LZMADecompressor::Start(const u8* in_ptr, u64 size) +{ + if (m_started || m_error_occurred) + return false; + + m_started = lzma_raw_decoder(&m_stream, m_filters) == LZMA_OK; + + m_stream.next_in = in_ptr; + m_stream.avail_in = size; + + return m_started; +} + +u64 WIAFileReader::LZMADecompressor::Read(u8* out_ptr, u64 size) +{ + if (!m_started || m_error_occurred || m_stream.avail_in == 0) + return 0; + + m_stream.next_out = out_ptr; + m_stream.avail_out = size; + + const lzma_ret result = lzma_code(&m_stream, LZMA_RUN); + m_error_occurred = result != LZMA_OK && result != LZMA_STREAM_END; + + return m_error_occurred ? 0 : m_stream.next_out - out_ptr; +} + +bool WIAFileReader::LZMADecompressor::DoneReading() const +{ + return m_started && !m_error_occurred && m_stream.avail_in == 0; +} + +void WIAFileReader::LZMADecompressor::End() +{ + if (m_started && !m_ended) + { + lzma_end(&m_stream); + m_ended = true; + } +} + } // namespace DiscIO diff --git a/Source/Core/DiscIO/WIABlob.h b/Source/Core/DiscIO/WIABlob.h index 8436667bcda7..01357805838c 100644 --- a/Source/Core/DiscIO/WIABlob.h +++ b/Source/Core/DiscIO/WIABlob.h @@ -8,6 +8,9 @@ #include #include +#include +#include + #include "Common/CommonTypes.h" #include "Common/File.h" #include "Common/Swap.h" @@ -146,6 +149,62 @@ class WIAFileReader : public BlobReader LZMA2 = 4, }; + class Decompressor + { + public: + virtual ~Decompressor(); + + // Specifies the compressed data to read. The data must still be in memory when calling Read. + virtual bool Start(const u8* in_ptr, u64 size) = 0; + + // Reads the specified number of bytes into out_ptr (or less, if there aren't that many bytes + // to output). Returns the number of bytes read. Start must be called before this. + virtual u64 Read(u8* out_ptr, u64 size) = 0; + + // Returns whether every byte of the input data has been read. + virtual bool DoneReading() const = 0; + + // Will be called automatically upon destruction, but can be called earlier if desired. + virtual void End() = 0; + }; + + class Bzip2Decompressor final : public Decompressor + { + public: + ~Bzip2Decompressor(); + + bool Start(const u8* in_ptr, u64 size) override; + u64 Read(u8* out_ptr, u64 size) override; + bool DoneReading() const override; + void End() override; + + private: + bz_stream m_stream; + bool m_started = false; + bool m_ended = false; + bool m_error_occurred = false; + }; + + class LZMADecompressor final : public Decompressor + { + public: + LZMADecompressor(bool lzma2, const std::vector& filter_options); + ~LZMADecompressor(); + + bool Start(const u8* in_ptr, u64 size) override; + u64 Read(u8* out_ptr, u64 size) override; + bool DoneReading() const override; + void End() override; + + private: + lzma_stream m_stream = LZMA_STREAM_INIT; + lzma_options_lzma m_options = {}; + lzma_filter m_filters[2]; + bool m_started = false; + bool m_ended = false; + bool m_error_occurred = false; + }; + bool m_valid; CompressionType m_compression_type; @@ -153,6 +212,7 @@ class WIAFileReader : public BlobReader WIAHeader1 m_header_1; WIAHeader2 m_header_2; + std::vector m_compressor_data; std::vector m_partition_entries; std::vector m_raw_data_entries; std::vector m_group_entries;