WIA/RVZ: Move (de)compression to a different file

This commit is contained in:
JosJuice 2020-05-17 15:21:34 +02:00
parent 39caac925a
commit ca4e4a6207
7 changed files with 1054 additions and 1007 deletions

View file

@ -46,6 +46,8 @@ add_library(discio
WbfsBlob.h
WIABlob.cpp
WIABlob.h
WIACompression.cpp
WIACompression.h
WiiEncryptionCache.cpp
WiiEncryptionCache.h
WiiSaveBanner.cpp

View file

@ -67,6 +67,7 @@
<ClCompile Include="VolumeWii.cpp" />
<ClCompile Include="WbfsBlob.cpp" />
<ClCompile Include="WIABlob.cpp" />
<ClCompile Include="WIACompression.cpp" />
<ClCompile Include="WiiEncryptionCache.cpp" />
<ClCompile Include="WiiSaveBanner.cpp" />
</ItemGroup>
@ -95,6 +96,7 @@
<ClInclude Include="VolumeWii.h" />
<ClInclude Include="WbfsBlob.h" />
<ClInclude Include="WIABlob.h" />
<ClInclude Include="WIACompression.h" />
<ClInclude Include="WiiEncryptionCache.h" />
<ClInclude Include="WiiSaveBanner.h" />
</ItemGroup>

View file

@ -96,6 +96,9 @@
<ClCompile Include="LaggedFibonacciGenerator.cpp">
<Filter>Volume\Blob</Filter>
</ClCompile>
<ClCompile Include="WIACompression.cpp">
<Filter>Volume\Blob</Filter>
</ClCompile>
</ItemGroup>
<ItemGroup>
<ClInclude Include="DiscScrubber.h">
@ -176,6 +179,9 @@
<ClInclude Include="LaggedFibonacciGenerator.h">
<Filter>Volume\Blob</Filter>
</ClInclude>
<ClInclude Include="WIACompression.h">
<Filter>Volume\Blob</Filter>
</ClInclude>
</ItemGroup>
<ItemGroup>
<Text Include="CMakeLists.txt" />

View file

@ -16,8 +16,6 @@
#include <type_traits>
#include <utility>
#include <bzlib.h>
#include <lzma.h>
#include <mbedtls/sha1.h>
#include <zstd.h>
@ -39,6 +37,7 @@
#include "DiscIO/MultithreadedCompressor.h"
#include "DiscIO/Volume.h"
#include "DiscIO/VolumeWii.h"
#include "DiscIO/WIACompression.h"
#include "DiscIO/WiiEncryptionCache.h"
namespace DiscIO
@ -526,777 +525,6 @@ std::string WIAFileReader::VersionToString(u32 version)
return StringFromFormat("%u.%02x.%02x.beta%u", a, b, c, d);
}
u32 WIAFileReader::LZMA2DictionarySize(u8 p)
{
return (static_cast<u32>(2) | (p & 1)) << (p / 2 + 11);
}
WIAFileReader::Decompressor::~Decompressor() = default;
bool WIAFileReader::NoneDecompressor::Decompress(const DecompressionBuffer& in,
DecompressionBuffer* out, size_t* in_bytes_read)
{
const size_t length =
std::min(in.bytes_written - *in_bytes_read, out->data.size() - out->bytes_written);
std::memcpy(out->data.data() + out->bytes_written, in.data.data() + *in_bytes_read, length);
*in_bytes_read += length;
out->bytes_written += length;
m_done = in.data.size() == *in_bytes_read;
return true;
}
WIAFileReader::PurgeDecompressor::PurgeDecompressor(u64 decompressed_size)
: m_decompressed_size(decompressed_size)
{
mbedtls_sha1_init(&m_sha1_context);
}
bool WIAFileReader::PurgeDecompressor::Decompress(const DecompressionBuffer& in,
DecompressionBuffer* out, size_t* in_bytes_read)
{
if (!m_started)
{
mbedtls_sha1_starts_ret(&m_sha1_context);
// Include the exception lists in the SHA-1 calculation (but not in the compression...)
mbedtls_sha1_update_ret(&m_sha1_context, in.data.data(), *in_bytes_read);
m_started = true;
}
while (!m_done && in.bytes_written != *in_bytes_read &&
(m_segment_bytes_written < sizeof(m_segment) || out->data.size() != out->bytes_written))
{
if (m_segment_bytes_written == 0 && *in_bytes_read == in.data.size() - sizeof(SHA1))
{
const size_t zeroes_to_write = std::min<size_t>(m_decompressed_size - m_out_bytes_written,
out->data.size() - out->bytes_written);
std::memset(out->data.data() + out->bytes_written, 0, zeroes_to_write);
out->bytes_written += zeroes_to_write;
m_out_bytes_written += zeroes_to_write;
if (m_out_bytes_written == m_decompressed_size && in.bytes_written == in.data.size())
{
SHA1 actual_hash;
mbedtls_sha1_finish_ret(&m_sha1_context, actual_hash.data());
SHA1 expected_hash;
std::memcpy(expected_hash.data(), in.data.data() + *in_bytes_read, expected_hash.size());
*in_bytes_read += expected_hash.size();
m_done = true;
if (actual_hash != expected_hash)
return false;
}
return true;
}
if (m_segment_bytes_written < sizeof(m_segment))
{
const size_t bytes_to_copy =
std::min(in.bytes_written - *in_bytes_read, sizeof(m_segment) - m_segment_bytes_written);
std::memcpy(reinterpret_cast<u8*>(&m_segment) + m_segment_bytes_written,
in.data.data() + *in_bytes_read, bytes_to_copy);
mbedtls_sha1_update_ret(&m_sha1_context, in.data.data() + *in_bytes_read, bytes_to_copy);
*in_bytes_read += bytes_to_copy;
m_bytes_read += bytes_to_copy;
m_segment_bytes_written += bytes_to_copy;
}
if (m_segment_bytes_written < sizeof(m_segment))
return true;
const size_t offset = Common::swap32(m_segment.offset);
const size_t size = Common::swap32(m_segment.size);
if (m_out_bytes_written < offset)
{
const size_t zeroes_to_write =
std::min(offset - m_out_bytes_written, out->data.size() - out->bytes_written);
std::memset(out->data.data() + out->bytes_written, 0, zeroes_to_write);
out->bytes_written += zeroes_to_write;
m_out_bytes_written += zeroes_to_write;
}
if (m_out_bytes_written >= offset && m_out_bytes_written < offset + size)
{
const size_t bytes_to_copy = std::min(
std::min(offset + size - m_out_bytes_written, out->data.size() - out->bytes_written),
in.bytes_written - *in_bytes_read);
std::memcpy(out->data.data() + out->bytes_written, in.data.data() + *in_bytes_read,
bytes_to_copy);
mbedtls_sha1_update_ret(&m_sha1_context, in.data.data() + *in_bytes_read, bytes_to_copy);
*in_bytes_read += bytes_to_copy;
m_bytes_read += bytes_to_copy;
out->bytes_written += bytes_to_copy;
m_out_bytes_written += bytes_to_copy;
}
if (m_out_bytes_written >= offset + size)
m_segment_bytes_written = 0;
}
return true;
}
WIAFileReader::Bzip2Decompressor::~Bzip2Decompressor()
{
if (m_started)
BZ2_bzDecompressEnd(&m_stream);
}
bool WIAFileReader::Bzip2Decompressor::Decompress(const DecompressionBuffer& in,
DecompressionBuffer* out, size_t* in_bytes_read)
{
if (!m_started)
{
if (BZ2_bzDecompressInit(&m_stream, 0, 0) != BZ_OK)
return false;
m_started = true;
}
constexpr auto clamped_cast = [](size_t x) {
return static_cast<unsigned int>(
std::min<size_t>(std::numeric_limits<unsigned int>().max(), x));
};
char* const in_ptr = reinterpret_cast<char*>(const_cast<u8*>(in.data.data() + *in_bytes_read));
m_stream.next_in = in_ptr;
m_stream.avail_in = clamped_cast(in.bytes_written - *in_bytes_read);
char* const out_ptr = reinterpret_cast<char*>(out->data.data() + out->bytes_written);
m_stream.next_out = out_ptr;
m_stream.avail_out = clamped_cast(out->data.size() - out->bytes_written);
const int result = BZ2_bzDecompress(&m_stream);
*in_bytes_read += m_stream.next_in - in_ptr;
out->bytes_written += m_stream.next_out - out_ptr;
m_done = result == BZ_STREAM_END;
return result == BZ_OK || result == BZ_STREAM_END;
}
WIAFileReader::LZMADecompressor::LZMADecompressor(bool lzma2, const u8* filter_options,
size_t filter_options_size)
{
m_options.preset_dict = nullptr;
if (!lzma2 && filter_options_size == 5)
{
// The dictionary size is stored as a 32-bit little endian unsigned integer
static_assert(sizeof(m_options.dict_size) == sizeof(u32));
std::memcpy(&m_options.dict_size, filter_options + 1, sizeof(u32));
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 : LZMA2DictionarySize(d);
}
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()
{
if (m_started)
lzma_end(&m_stream);
}
bool WIAFileReader::LZMADecompressor::Decompress(const DecompressionBuffer& in,
DecompressionBuffer* out, size_t* in_bytes_read)
{
if (!m_started)
{
if (m_error_occurred || lzma_raw_decoder(&m_stream, m_filters) != LZMA_OK)
return false;
m_started = true;
}
const u8* const in_ptr = in.data.data() + *in_bytes_read;
m_stream.next_in = in_ptr;
m_stream.avail_in = in.bytes_written - *in_bytes_read;
u8* const out_ptr = out->data.data() + out->bytes_written;
m_stream.next_out = out_ptr;
m_stream.avail_out = out->data.size() - out->bytes_written;
const lzma_ret result = lzma_code(&m_stream, LZMA_RUN);
*in_bytes_read += m_stream.next_in - in_ptr;
out->bytes_written += m_stream.next_out - out_ptr;
m_done = result == LZMA_STREAM_END;
return result == LZMA_OK || result == LZMA_STREAM_END;
}
WIAFileReader::ZstdDecompressor::ZstdDecompressor()
{
m_stream = ZSTD_createDStream();
}
WIAFileReader::ZstdDecompressor::~ZstdDecompressor()
{
ZSTD_freeDStream(m_stream);
}
bool WIAFileReader::ZstdDecompressor::Decompress(const DecompressionBuffer& in,
DecompressionBuffer* out, size_t* in_bytes_read)
{
if (!m_stream)
return false;
ZSTD_inBuffer in_buffer{in.data.data(), in.bytes_written, *in_bytes_read};
ZSTD_outBuffer out_buffer{out->data.data(), out->data.size(), out->bytes_written};
const size_t result = ZSTD_decompressStream(m_stream, &out_buffer, &in_buffer);
*in_bytes_read = in_buffer.pos;
out->bytes_written = out_buffer.pos;
m_done = result == 0;
return !ZSTD_isError(result);
}
WIAFileReader::RVZPackDecompressor::RVZPackDecompressor(std::unique_ptr<Decompressor> decompressor,
DecompressionBuffer decompressed,
u64 data_offset)
: m_decompressor(std::move(decompressor)), m_decompressed(std::move(decompressed)),
m_data_offset(data_offset)
{
}
std::optional<bool> WIAFileReader::RVZPackDecompressor::ReadToDecompressed(
const DecompressionBuffer& in, size_t* in_bytes_read, size_t decompressed_bytes_read,
size_t bytes_to_read)
{
if (m_decompressed.data.size() < decompressed_bytes_read + bytes_to_read)
m_decompressed.data.resize(decompressed_bytes_read + bytes_to_read);
if (m_decompressed.bytes_written < decompressed_bytes_read + bytes_to_read)
{
if (!m_decompressor->Decompress(in, &m_decompressed, in_bytes_read))
return false;
if (m_decompressed.bytes_written < decompressed_bytes_read + bytes_to_read)
return true;
}
return std::nullopt;
}
bool WIAFileReader::RVZPackDecompressor::Decompress(const DecompressionBuffer& in,
DecompressionBuffer* out, size_t* in_bytes_read)
{
while (out->data.size() != out->bytes_written && !Done())
{
if (m_size == 0)
{
if (m_decompressed.bytes_written == m_decompressed_bytes_read)
{
m_decompressed.data.resize(sizeof(u32));
m_decompressed.bytes_written = 0;
m_decompressed_bytes_read = 0;
}
std::optional<bool> result =
ReadToDecompressed(in, in_bytes_read, m_decompressed_bytes_read, sizeof(u32));
if (result)
return *result;
m_size = Common::swap32(m_decompressed.data.data() + m_decompressed_bytes_read);
m_junk = m_size & 0x80000000;
if (m_junk)
{
m_size &= 0x7FFFFFFF;
constexpr size_t SEED_SIZE = LaggedFibonacciGenerator::SEED_SIZE * sizeof(u32);
result = ReadToDecompressed(in, in_bytes_read, m_decompressed_bytes_read + sizeof(u32),
SEED_SIZE);
if (result)
return *result;
m_lfg.SetSeed(m_decompressed.data.data() + m_decompressed_bytes_read + sizeof(u32));
m_lfg.Forward(m_data_offset % VolumeWii::BLOCK_TOTAL_SIZE);
m_decompressed_bytes_read += SEED_SIZE;
}
m_decompressed_bytes_read += sizeof(u32);
}
size_t bytes_to_write = std::min<size_t>(m_size, out->data.size() - out->bytes_written);
if (m_junk)
{
m_lfg.GetBytes(bytes_to_write, out->data.data() + out->bytes_written);
out->bytes_written += bytes_to_write;
}
else
{
if (m_decompressed.bytes_written != m_decompressed_bytes_read)
{
bytes_to_write =
std::min(bytes_to_write, m_decompressed.bytes_written - m_decompressed_bytes_read);
std::memcpy(out->data.data() + out->bytes_written,
m_decompressed.data.data() + m_decompressed_bytes_read, bytes_to_write);
m_decompressed_bytes_read += bytes_to_write;
out->bytes_written += bytes_to_write;
}
else
{
const size_t prev_out_bytes_written = out->bytes_written;
const size_t old_out_size = out->data.size();
const size_t new_out_size = out->bytes_written + bytes_to_write;
if (new_out_size < old_out_size)
out->data.resize(new_out_size);
if (!m_decompressor->Decompress(in, out, in_bytes_read))
return false;
out->data.resize(old_out_size);
bytes_to_write = out->bytes_written - prev_out_bytes_written;
if (bytes_to_write == 0)
return true;
}
}
m_data_offset += bytes_to_write;
m_size -= static_cast<u32>(bytes_to_write);
}
// If out is full but not all data has been read from in, give the decompressor a chance to read
// from in anyway. This is needed for the case where zstd has read everything except the checksum.
if (out->data.size() == out->bytes_written && in.bytes_written != *in_bytes_read)
{
if (!m_decompressor->Decompress(in, out, in_bytes_read))
return false;
}
return true;
}
bool WIAFileReader::RVZPackDecompressor::Done() const
{
return m_size == 0 && m_decompressed.bytes_written == m_decompressed_bytes_read &&
m_decompressor->Done();
}
WIAFileReader::Compressor::~Compressor() = default;
WIAFileReader::PurgeCompressor::PurgeCompressor()
{
mbedtls_sha1_init(&m_sha1_context);
}
WIAFileReader::PurgeCompressor::~PurgeCompressor() = default;
bool WIAFileReader::PurgeCompressor::Start()
{
m_buffer.clear();
m_bytes_written = 0;
mbedtls_sha1_starts_ret(&m_sha1_context);
return true;
}
bool WIAFileReader::PurgeCompressor::AddPrecedingDataOnlyForPurgeHashing(const u8* data,
size_t size)
{
mbedtls_sha1_update_ret(&m_sha1_context, data, size);
return true;
}
bool WIAFileReader::PurgeCompressor::Compress(const u8* data, size_t size)
{
// We could add support for calling this twice if we're fine with
// making the code more complicated, but there's no need to support it
ASSERT_MSG(DISCIO, m_bytes_written == 0,
"Calling PurgeCompressor::Compress() twice is not supported");
m_buffer.resize(size + sizeof(PurgeSegment) + sizeof(SHA1));
size_t bytes_read = 0;
while (true)
{
const auto first_non_zero =
std::find_if(data + bytes_read, data + size, [](u8 x) { return x != 0; });
const u32 non_zero_data_start = static_cast<u32>(first_non_zero - data);
if (non_zero_data_start == size)
break;
size_t non_zero_data_end = non_zero_data_start;
size_t sequence_length = 0;
for (size_t i = non_zero_data_start; i < size; ++i)
{
if (data[i] == 0)
{
++sequence_length;
}
else
{
sequence_length = 0;
non_zero_data_end = i + 1;
}
// To avoid wasting space, only count runs of zeroes that are of a certain length
// (unless there is nothing after the run of zeroes, then we might as well always count it)
if (sequence_length > sizeof(PurgeSegment))
break;
}
const u32 non_zero_data_length = static_cast<u32>(non_zero_data_end - non_zero_data_start);
const PurgeSegment segment{Common::swap32(non_zero_data_start),
Common::swap32(non_zero_data_length)};
std::memcpy(m_buffer.data() + m_bytes_written, &segment, sizeof(segment));
m_bytes_written += sizeof(segment);
std::memcpy(m_buffer.data() + m_bytes_written, data + non_zero_data_start,
non_zero_data_length);
m_bytes_written += non_zero_data_length;
bytes_read = non_zero_data_end;
}
return true;
}
bool WIAFileReader::PurgeCompressor::End()
{
mbedtls_sha1_update_ret(&m_sha1_context, m_buffer.data(), m_bytes_written);
mbedtls_sha1_finish_ret(&m_sha1_context, m_buffer.data() + m_bytes_written);
m_bytes_written += sizeof(SHA1);
ASSERT(m_bytes_written <= m_buffer.size());
return true;
}
const u8* WIAFileReader::PurgeCompressor::GetData() const
{
return m_buffer.data();
}
size_t WIAFileReader::PurgeCompressor::GetSize() const
{
return m_bytes_written;
}
WIAFileReader::Bzip2Compressor::Bzip2Compressor(int compression_level)
: m_compression_level(compression_level)
{
}
WIAFileReader::Bzip2Compressor::~Bzip2Compressor()
{
BZ2_bzCompressEnd(&m_stream);
}
bool WIAFileReader::Bzip2Compressor::Start()
{
ASSERT_MSG(DISCIO, m_stream.state == nullptr,
"Called Bzip2Compressor::Start() twice without calling Bzip2Compressor::End()");
m_buffer.clear();
m_stream.next_out = reinterpret_cast<char*>(m_buffer.data());
return BZ2_bzCompressInit(&m_stream, m_compression_level, 0, 0) == BZ_OK;
}
bool WIAFileReader::Bzip2Compressor::Compress(const u8* data, size_t size)
{
m_stream.next_in = reinterpret_cast<char*>(const_cast<u8*>(data));
m_stream.avail_in = static_cast<unsigned int>(size);
ExpandBuffer(size);
while (m_stream.avail_in != 0)
{
if (m_stream.avail_out == 0)
ExpandBuffer(0x100);
if (BZ2_bzCompress(&m_stream, BZ_RUN) != BZ_RUN_OK)
return false;
}
return true;
}
bool WIAFileReader::Bzip2Compressor::End()
{
bool success = true;
while (true)
{
if (m_stream.avail_out == 0)
ExpandBuffer(0x100);
const int result = BZ2_bzCompress(&m_stream, BZ_FINISH);
if (result != BZ_FINISH_OK && result != BZ_STREAM_END)
success = false;
if (result != BZ_FINISH_OK)
break;
}
if (BZ2_bzCompressEnd(&m_stream) != BZ_OK)
success = false;
return success;
}
void WIAFileReader::Bzip2Compressor::ExpandBuffer(size_t bytes_to_add)
{
const size_t bytes_written = GetSize();
m_buffer.resize(m_buffer.size() + bytes_to_add);
m_stream.next_out = reinterpret_cast<char*>(m_buffer.data()) + bytes_written;
m_stream.avail_out = static_cast<unsigned int>(m_buffer.size() - bytes_written);
}
const u8* WIAFileReader::Bzip2Compressor::GetData() const
{
return m_buffer.data();
}
size_t WIAFileReader::Bzip2Compressor::GetSize() const
{
return static_cast<size_t>(reinterpret_cast<u8*>(m_stream.next_out) - m_buffer.data());
}
WIAFileReader::LZMACompressor::LZMACompressor(bool lzma2, int compression_level,
u8 compressor_data_out[7],
u8* compressor_data_size_out)
{
// lzma_lzma_preset returns false on success for some reason
if (lzma_lzma_preset(&m_options, static_cast<uint32_t>(compression_level)))
{
m_initialization_failed = true;
return;
}
if (!lzma2)
{
if (compressor_data_size_out)
*compressor_data_size_out = 5;
if (compressor_data_out)
{
ASSERT(m_options.lc < 9);
ASSERT(m_options.lp < 5);
ASSERT(m_options.pb < 5);
compressor_data_out[0] =
static_cast<u8>((m_options.pb * 5 + m_options.lp) * 9 + m_options.lc);
// The dictionary size is stored as a 32-bit little endian unsigned integer
static_assert(sizeof(m_options.dict_size) == sizeof(u32));
std::memcpy(compressor_data_out + 1, &m_options.dict_size, sizeof(u32));
}
}
else
{
if (compressor_data_size_out)
*compressor_data_size_out = 1;
if (compressor_data_out)
{
u8 encoded_dict_size = 0;
while (encoded_dict_size < 40 && m_options.dict_size > LZMA2DictionarySize(encoded_dict_size))
++encoded_dict_size;
compressor_data_out[0] = encoded_dict_size;
}
}
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::LZMACompressor::~LZMACompressor()
{
lzma_end(&m_stream);
}
bool WIAFileReader::LZMACompressor::Start()
{
if (m_initialization_failed)
return false;
m_buffer.clear();
m_stream.next_out = m_buffer.data();
return lzma_raw_encoder(&m_stream, m_filters) == LZMA_OK;
}
bool WIAFileReader::LZMACompressor::Compress(const u8* data, size_t size)
{
m_stream.next_in = data;
m_stream.avail_in = size;
ExpandBuffer(size);
while (m_stream.avail_in != 0)
{
if (m_stream.avail_out == 0)
ExpandBuffer(0x100);
if (lzma_code(&m_stream, LZMA_RUN) != LZMA_OK)
return false;
}
return true;
}
bool WIAFileReader::LZMACompressor::End()
{
while (true)
{
if (m_stream.avail_out == 0)
ExpandBuffer(0x100);
switch (lzma_code(&m_stream, LZMA_FINISH))
{
case LZMA_OK:
break;
case LZMA_STREAM_END:
return true;
default:
return false;
}
}
}
void WIAFileReader::LZMACompressor::ExpandBuffer(size_t bytes_to_add)
{
const size_t bytes_written = GetSize();
m_buffer.resize(m_buffer.size() + bytes_to_add);
m_stream.next_out = m_buffer.data() + bytes_written;
m_stream.avail_out = m_buffer.size() - bytes_written;
}
const u8* WIAFileReader::LZMACompressor::GetData() const
{
return m_buffer.data();
}
size_t WIAFileReader::LZMACompressor::GetSize() const
{
return static_cast<size_t>(m_stream.next_out - m_buffer.data());
}
WIAFileReader::ZstdCompressor::ZstdCompressor(int compression_level)
{
m_stream = ZSTD_createCStream();
if (ZSTD_isError(ZSTD_CCtx_setParameter(m_stream, ZSTD_c_compressionLevel, compression_level)))
m_stream = nullptr;
}
WIAFileReader::ZstdCompressor::~ZstdCompressor()
{
ZSTD_freeCStream(m_stream);
}
bool WIAFileReader::ZstdCompressor::Start()
{
if (!m_stream)
return false;
m_buffer.clear();
m_out_buffer = {};
return !ZSTD_isError(ZSTD_CCtx_reset(m_stream, ZSTD_reset_session_only));
}
bool WIAFileReader::ZstdCompressor::Compress(const u8* data, size_t size)
{
ZSTD_inBuffer in_buffer{data, size, 0};
ExpandBuffer(size);
while (in_buffer.size != in_buffer.pos)
{
if (m_out_buffer.size == m_out_buffer.pos)
ExpandBuffer(0x100);
if (ZSTD_isError(ZSTD_compressStream(m_stream, &m_out_buffer, &in_buffer)))
return false;
}
return true;
}
bool WIAFileReader::ZstdCompressor::End()
{
while (true)
{
if (m_out_buffer.size == m_out_buffer.pos)
ExpandBuffer(0x100);
const size_t result = ZSTD_endStream(m_stream, &m_out_buffer);
if (ZSTD_isError(result))
return false;
if (result == 0)
return true;
}
}
void WIAFileReader::ZstdCompressor::ExpandBuffer(size_t bytes_to_add)
{
m_buffer.resize(m_buffer.size() + bytes_to_add);
m_out_buffer.dst = m_buffer.data();
m_out_buffer.size = m_buffer.size();
}
WIAFileReader::Chunk::Chunk() = default;
WIAFileReader::Chunk::Chunk(File::IOFile* file, u64 offset_in_file, u64 compressed_size,

View file

@ -9,21 +9,15 @@
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <type_traits>
#include <utility>
#include <bzlib.h>
#include <lzma.h>
#include <mbedtls/sha1.h>
#include <zstd.h>
#include "Common/CommonTypes.h"
#include "Common/File.h"
#include "Common/Swap.h"
#include "DiscIO/Blob.h"
#include "DiscIO/LaggedFibonacciGenerator.h"
#include "DiscIO/MultithreadedCompressor.h"
#include "DiscIO/WIACompression.h"
#include "DiscIO/WiiEncryptionCache.h"
namespace DiscIO
@ -157,13 +151,6 @@ private:
SHA1 hash;
};
static_assert(sizeof(HashExceptionEntry) == 0x16, "Wrong size for WIA hash exception entry");
struct PurgeSegment
{
u32 offset;
u32 size;
};
static_assert(sizeof(PurgeSegment) == 0x08, "Wrong size for WIA purge segment");
#pragma pack(pop)
struct DataEntry
@ -180,224 +167,6 @@ private:
}
};
struct DecompressionBuffer
{
std::vector<u8> data;
size_t bytes_written = 0;
};
class Decompressor
{
public:
virtual ~Decompressor();
virtual bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) = 0;
virtual bool Done() const { return m_done; };
protected:
bool m_done = false;
};
class NoneDecompressor final : public Decompressor
{
public:
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
};
// This class assumes that more bytes won't be added to in once in.bytes_written == in.data.size()
// and that *in_bytes_read initially will be equal to the size of the exception lists
class PurgeDecompressor final : public Decompressor
{
public:
PurgeDecompressor(u64 decompressed_size);
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
private:
const u64 m_decompressed_size;
PurgeSegment m_segment = {};
size_t m_bytes_read = 0;
size_t m_segment_bytes_written = 0;
size_t m_out_bytes_written = 0;
bool m_started = false;
mbedtls_sha1_context m_sha1_context;
};
class Bzip2Decompressor final : public Decompressor
{
public:
~Bzip2Decompressor();
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
private:
bz_stream m_stream = {};
bool m_started = false;
};
class LZMADecompressor final : public Decompressor
{
public:
LZMADecompressor(bool lzma2, const u8* filter_options, size_t filter_options_size);
~LZMADecompressor();
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) 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_error_occurred = false;
};
class ZstdDecompressor final : public Decompressor
{
public:
ZstdDecompressor();
~ZstdDecompressor();
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
private:
ZSTD_DStream* m_stream;
};
class RVZPackDecompressor final : public Decompressor
{
public:
RVZPackDecompressor(std::unique_ptr<Decompressor> decompressor,
DecompressionBuffer decompressed, u64 data_offset);
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
bool Done() const override;
private:
std::optional<bool> ReadToDecompressed(const DecompressionBuffer& in, size_t* in_bytes_read,
size_t decompressed_bytes_read, size_t bytes_to_read);
std::unique_ptr<Decompressor> m_decompressor;
DecompressionBuffer m_decompressed;
size_t m_decompressed_bytes_read = 0;
u64 m_data_offset;
u32 m_size = 0;
bool m_junk;
LaggedFibonacciGenerator m_lfg;
};
class Compressor
{
public:
virtual ~Compressor();
// First call Start, then AddDataOnlyForPurgeHashing/Compress any number of times,
// then End, then GetData/GetSize any number of times.
virtual bool Start() = 0;
virtual bool AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) { return true; }
virtual bool Compress(const u8* data, size_t size) = 0;
virtual bool End() = 0;
virtual const u8* GetData() const = 0;
virtual size_t GetSize() const = 0;
};
class PurgeCompressor final : public Compressor
{
public:
PurgeCompressor();
~PurgeCompressor();
bool Start() override;
bool AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) override;
bool Compress(const u8* data, size_t size) override;
bool End() override;
const u8* GetData() const override;
size_t GetSize() const override;
private:
std::vector<u8> m_buffer;
size_t m_bytes_written;
mbedtls_sha1_context m_sha1_context;
};
class Bzip2Compressor final : public Compressor
{
public:
Bzip2Compressor(int compression_level);
~Bzip2Compressor();
bool Start() override;
bool Compress(const u8* data, size_t size) override;
bool End() override;
const u8* GetData() const override;
size_t GetSize() const override;
private:
void ExpandBuffer(size_t bytes_to_add);
bz_stream m_stream = {};
std::vector<u8> m_buffer;
int m_compression_level;
};
class LZMACompressor final : public Compressor
{
public:
LZMACompressor(bool lzma2, int compression_level, u8 compressor_data_out[7],
u8* compressor_data_size_out);
~LZMACompressor();
bool Start() override;
bool Compress(const u8* data, size_t size) override;
bool End() override;
const u8* GetData() const override;
size_t GetSize() const override;
private:
void ExpandBuffer(size_t bytes_to_add);
lzma_stream m_stream = LZMA_STREAM_INIT;
lzma_options_lzma m_options = {};
lzma_filter m_filters[2];
std::vector<u8> m_buffer;
bool m_initialization_failed = false;
};
class ZstdCompressor final : public Compressor
{
public:
ZstdCompressor(int compression_level);
~ZstdCompressor();
bool Start() override;
bool Compress(const u8* data, size_t size) override;
bool End() override;
const u8* GetData() const override { return m_buffer.data(); }
size_t GetSize() const override { return m_out_buffer.pos; }
private:
void ExpandBuffer(size_t bytes_to_add);
ZSTD_CStream* m_stream;
ZSTD_outBuffer m_out_buffer;
std::vector<u8> m_buffer;
};
class Chunk
{
public:
@ -455,8 +224,6 @@ private:
static std::string VersionToString(u32 version);
static u32 LZMA2DictionarySize(u8 p);
struct ReuseID
{
bool operator==(const ReuseID& other) const

View file

@ -0,0 +1,793 @@
// Copyright 2020 Dolphin Emulator Project
// Licensed under GPLv2+
// Refer to the license.txt file included.
#include "DiscIO/WIACompression.h"
#include <algorithm>
#include <cstddef>
#include <cstring>
#include <memory>
#include <optional>
#include <vector>
#include <bzlib.h>
#include <lzma.h>
#include <mbedtls/sha1.h>
#include <zstd.h>
#include "Common/Assert.h"
#include "Common/CommonTypes.h"
#include "Common/Swap.h"
#include "DiscIO/LaggedFibonacciGenerator.h"
namespace DiscIO
{
static u32 LZMA2DictionarySize(u8 p)
{
return (static_cast<u32>(2) | (p & 1)) << (p / 2 + 11);
}
Decompressor::~Decompressor() = default;
bool NoneDecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read)
{
const size_t length =
std::min(in.bytes_written - *in_bytes_read, out->data.size() - out->bytes_written);
std::memcpy(out->data.data() + out->bytes_written, in.data.data() + *in_bytes_read, length);
*in_bytes_read += length;
out->bytes_written += length;
m_done = in.data.size() == *in_bytes_read;
return true;
}
PurgeDecompressor::PurgeDecompressor(u64 decompressed_size) : m_decompressed_size(decompressed_size)
{
mbedtls_sha1_init(&m_sha1_context);
}
bool PurgeDecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read)
{
if (!m_started)
{
mbedtls_sha1_starts_ret(&m_sha1_context);
// Include the exception lists in the SHA-1 calculation (but not in the compression...)
mbedtls_sha1_update_ret(&m_sha1_context, in.data.data(), *in_bytes_read);
m_started = true;
}
while (!m_done && in.bytes_written != *in_bytes_read &&
(m_segment_bytes_written < sizeof(m_segment) || out->data.size() != out->bytes_written))
{
if (m_segment_bytes_written == 0 && *in_bytes_read == in.data.size() - sizeof(SHA1))
{
const size_t zeroes_to_write = std::min<size_t>(m_decompressed_size - m_out_bytes_written,
out->data.size() - out->bytes_written);
std::memset(out->data.data() + out->bytes_written, 0, zeroes_to_write);
out->bytes_written += zeroes_to_write;
m_out_bytes_written += zeroes_to_write;
if (m_out_bytes_written == m_decompressed_size && in.bytes_written == in.data.size())
{
SHA1 actual_hash;
mbedtls_sha1_finish_ret(&m_sha1_context, actual_hash.data());
SHA1 expected_hash;
std::memcpy(expected_hash.data(), in.data.data() + *in_bytes_read, expected_hash.size());
*in_bytes_read += expected_hash.size();
m_done = true;
if (actual_hash != expected_hash)
return false;
}
return true;
}
if (m_segment_bytes_written < sizeof(m_segment))
{
const size_t bytes_to_copy =
std::min(in.bytes_written - *in_bytes_read, sizeof(m_segment) - m_segment_bytes_written);
std::memcpy(reinterpret_cast<u8*>(&m_segment) + m_segment_bytes_written,
in.data.data() + *in_bytes_read, bytes_to_copy);
mbedtls_sha1_update_ret(&m_sha1_context, in.data.data() + *in_bytes_read, bytes_to_copy);
*in_bytes_read += bytes_to_copy;
m_bytes_read += bytes_to_copy;
m_segment_bytes_written += bytes_to_copy;
}
if (m_segment_bytes_written < sizeof(m_segment))
return true;
const size_t offset = Common::swap32(m_segment.offset);
const size_t size = Common::swap32(m_segment.size);
if (m_out_bytes_written < offset)
{
const size_t zeroes_to_write =
std::min(offset - m_out_bytes_written, out->data.size() - out->bytes_written);
std::memset(out->data.data() + out->bytes_written, 0, zeroes_to_write);
out->bytes_written += zeroes_to_write;
m_out_bytes_written += zeroes_to_write;
}
if (m_out_bytes_written >= offset && m_out_bytes_written < offset + size)
{
const size_t bytes_to_copy = std::min(
std::min(offset + size - m_out_bytes_written, out->data.size() - out->bytes_written),
in.bytes_written - *in_bytes_read);
std::memcpy(out->data.data() + out->bytes_written, in.data.data() + *in_bytes_read,
bytes_to_copy);
mbedtls_sha1_update_ret(&m_sha1_context, in.data.data() + *in_bytes_read, bytes_to_copy);
*in_bytes_read += bytes_to_copy;
m_bytes_read += bytes_to_copy;
out->bytes_written += bytes_to_copy;
m_out_bytes_written += bytes_to_copy;
}
if (m_out_bytes_written >= offset + size)
m_segment_bytes_written = 0;
}
return true;
}
Bzip2Decompressor::~Bzip2Decompressor()
{
if (m_started)
BZ2_bzDecompressEnd(&m_stream);
}
bool Bzip2Decompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read)
{
if (!m_started)
{
if (BZ2_bzDecompressInit(&m_stream, 0, 0) != BZ_OK)
return false;
m_started = true;
}
constexpr auto clamped_cast = [](size_t x) {
return static_cast<unsigned int>(
std::min<size_t>(std::numeric_limits<unsigned int>().max(), x));
};
char* const in_ptr = reinterpret_cast<char*>(const_cast<u8*>(in.data.data() + *in_bytes_read));
m_stream.next_in = in_ptr;
m_stream.avail_in = clamped_cast(in.bytes_written - *in_bytes_read);
char* const out_ptr = reinterpret_cast<char*>(out->data.data() + out->bytes_written);
m_stream.next_out = out_ptr;
m_stream.avail_out = clamped_cast(out->data.size() - out->bytes_written);
const int result = BZ2_bzDecompress(&m_stream);
*in_bytes_read += m_stream.next_in - in_ptr;
out->bytes_written += m_stream.next_out - out_ptr;
m_done = result == BZ_STREAM_END;
return result == BZ_OK || result == BZ_STREAM_END;
}
LZMADecompressor::LZMADecompressor(bool lzma2, const u8* filter_options, size_t filter_options_size)
{
m_options.preset_dict = nullptr;
if (!lzma2 && filter_options_size == 5)
{
// The dictionary size is stored as a 32-bit little endian unsigned integer
static_assert(sizeof(m_options.dict_size) == sizeof(u32));
std::memcpy(&m_options.dict_size, filter_options + 1, sizeof(u32));
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 : LZMA2DictionarySize(d);
}
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;
}
LZMADecompressor::~LZMADecompressor()
{
if (m_started)
lzma_end(&m_stream);
}
bool LZMADecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read)
{
if (!m_started)
{
if (m_error_occurred || lzma_raw_decoder(&m_stream, m_filters) != LZMA_OK)
return false;
m_started = true;
}
const u8* const in_ptr = in.data.data() + *in_bytes_read;
m_stream.next_in = in_ptr;
m_stream.avail_in = in.bytes_written - *in_bytes_read;
u8* const out_ptr = out->data.data() + out->bytes_written;
m_stream.next_out = out_ptr;
m_stream.avail_out = out->data.size() - out->bytes_written;
const lzma_ret result = lzma_code(&m_stream, LZMA_RUN);
*in_bytes_read += m_stream.next_in - in_ptr;
out->bytes_written += m_stream.next_out - out_ptr;
m_done = result == LZMA_STREAM_END;
return result == LZMA_OK || result == LZMA_STREAM_END;
}
ZstdDecompressor::ZstdDecompressor()
{
m_stream = ZSTD_createDStream();
}
ZstdDecompressor::~ZstdDecompressor()
{
ZSTD_freeDStream(m_stream);
}
bool ZstdDecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read)
{
if (!m_stream)
return false;
ZSTD_inBuffer in_buffer{in.data.data(), in.bytes_written, *in_bytes_read};
ZSTD_outBuffer out_buffer{out->data.data(), out->data.size(), out->bytes_written};
const size_t result = ZSTD_decompressStream(m_stream, &out_buffer, &in_buffer);
*in_bytes_read = in_buffer.pos;
out->bytes_written = out_buffer.pos;
m_done = result == 0;
return !ZSTD_isError(result);
}
RVZPackDecompressor::RVZPackDecompressor(std::unique_ptr<Decompressor> decompressor,
DecompressionBuffer decompressed, u64 data_offset)
: m_decompressor(std::move(decompressor)), m_decompressed(std::move(decompressed)),
m_data_offset(data_offset)
{
}
std::optional<bool> RVZPackDecompressor::ReadToDecompressed(const DecompressionBuffer& in,
size_t* in_bytes_read,
size_t decompressed_bytes_read,
size_t bytes_to_read)
{
if (m_decompressed.data.size() < decompressed_bytes_read + bytes_to_read)
m_decompressed.data.resize(decompressed_bytes_read + bytes_to_read);
if (m_decompressed.bytes_written < decompressed_bytes_read + bytes_to_read)
{
if (!m_decompressor->Decompress(in, &m_decompressed, in_bytes_read))
return false;
if (m_decompressed.bytes_written < decompressed_bytes_read + bytes_to_read)
return true;
}
return std::nullopt;
}
bool RVZPackDecompressor::Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read)
{
while (out->data.size() != out->bytes_written && !Done())
{
if (m_size == 0)
{
if (m_decompressed.bytes_written == m_decompressed_bytes_read)
{
m_decompressed.data.resize(sizeof(u32));
m_decompressed.bytes_written = 0;
m_decompressed_bytes_read = 0;
}
std::optional<bool> result =
ReadToDecompressed(in, in_bytes_read, m_decompressed_bytes_read, sizeof(u32));
if (result)
return *result;
m_size = Common::swap32(m_decompressed.data.data() + m_decompressed_bytes_read);
m_junk = m_size & 0x80000000;
if (m_junk)
{
m_size &= 0x7FFFFFFF;
constexpr size_t SEED_SIZE = LaggedFibonacciGenerator::SEED_SIZE * sizeof(u32);
constexpr size_t BLOCK_SIZE = 0x8000;
result = ReadToDecompressed(in, in_bytes_read, m_decompressed_bytes_read + sizeof(u32),
SEED_SIZE);
if (result)
return *result;
m_lfg.SetSeed(m_decompressed.data.data() + m_decompressed_bytes_read + sizeof(u32));
m_lfg.Forward(m_data_offset % BLOCK_SIZE);
m_decompressed_bytes_read += SEED_SIZE;
}
m_decompressed_bytes_read += sizeof(u32);
}
size_t bytes_to_write = std::min<size_t>(m_size, out->data.size() - out->bytes_written);
if (m_junk)
{
m_lfg.GetBytes(bytes_to_write, out->data.data() + out->bytes_written);
out->bytes_written += bytes_to_write;
}
else
{
if (m_decompressed.bytes_written != m_decompressed_bytes_read)
{
bytes_to_write =
std::min(bytes_to_write, m_decompressed.bytes_written - m_decompressed_bytes_read);
std::memcpy(out->data.data() + out->bytes_written,
m_decompressed.data.data() + m_decompressed_bytes_read, bytes_to_write);
m_decompressed_bytes_read += bytes_to_write;
out->bytes_written += bytes_to_write;
}
else
{
const size_t prev_out_bytes_written = out->bytes_written;
const size_t old_out_size = out->data.size();
const size_t new_out_size = out->bytes_written + bytes_to_write;
if (new_out_size < old_out_size)
out->data.resize(new_out_size);
if (!m_decompressor->Decompress(in, out, in_bytes_read))
return false;
out->data.resize(old_out_size);
bytes_to_write = out->bytes_written - prev_out_bytes_written;
if (bytes_to_write == 0)
return true;
}
}
m_data_offset += bytes_to_write;
m_size -= static_cast<u32>(bytes_to_write);
}
// If out is full but not all data has been read from in, give the decompressor a chance to read
// from in anyway. This is needed for the case where zstd has read everything except the checksum.
if (out->data.size() == out->bytes_written && in.bytes_written != *in_bytes_read)
{
if (!m_decompressor->Decompress(in, out, in_bytes_read))
return false;
}
return true;
}
bool RVZPackDecompressor::Done() const
{
return m_size == 0 && m_decompressed.bytes_written == m_decompressed_bytes_read &&
m_decompressor->Done();
}
Compressor::~Compressor() = default;
PurgeCompressor::PurgeCompressor()
{
mbedtls_sha1_init(&m_sha1_context);
}
PurgeCompressor::~PurgeCompressor() = default;
bool PurgeCompressor::Start()
{
m_buffer.clear();
m_bytes_written = 0;
mbedtls_sha1_starts_ret(&m_sha1_context);
return true;
}
bool PurgeCompressor::AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size)
{
mbedtls_sha1_update_ret(&m_sha1_context, data, size);
return true;
}
bool PurgeCompressor::Compress(const u8* data, size_t size)
{
// We could add support for calling this twice if we're fine with
// making the code more complicated, but there's no need to support it
ASSERT_MSG(DISCIO, m_bytes_written == 0,
"Calling PurgeCompressor::Compress() twice is not supported");
m_buffer.resize(size + sizeof(PurgeSegment) + sizeof(SHA1));
size_t bytes_read = 0;
while (true)
{
const auto first_non_zero =
std::find_if(data + bytes_read, data + size, [](u8 x) { return x != 0; });
const u32 non_zero_data_start = static_cast<u32>(first_non_zero - data);
if (non_zero_data_start == size)
break;
size_t non_zero_data_end = non_zero_data_start;
size_t sequence_length = 0;
for (size_t i = non_zero_data_start; i < size; ++i)
{
if (data[i] == 0)
{
++sequence_length;
}
else
{
sequence_length = 0;
non_zero_data_end = i + 1;
}
// To avoid wasting space, only count runs of zeroes that are of a certain length
// (unless there is nothing after the run of zeroes, then we might as well always count it)
if (sequence_length > sizeof(PurgeSegment))
break;
}
const u32 non_zero_data_length = static_cast<u32>(non_zero_data_end - non_zero_data_start);
const PurgeSegment segment{Common::swap32(non_zero_data_start),
Common::swap32(non_zero_data_length)};
std::memcpy(m_buffer.data() + m_bytes_written, &segment, sizeof(segment));
m_bytes_written += sizeof(segment);
std::memcpy(m_buffer.data() + m_bytes_written, data + non_zero_data_start,
non_zero_data_length);
m_bytes_written += non_zero_data_length;
bytes_read = non_zero_data_end;
}
return true;
}
bool PurgeCompressor::End()
{
mbedtls_sha1_update_ret(&m_sha1_context, m_buffer.data(), m_bytes_written);
mbedtls_sha1_finish_ret(&m_sha1_context, m_buffer.data() + m_bytes_written);
m_bytes_written += sizeof(SHA1);
ASSERT(m_bytes_written <= m_buffer.size());
return true;
}
const u8* PurgeCompressor::GetData() const
{
return m_buffer.data();
}
size_t PurgeCompressor::GetSize() const
{
return m_bytes_written;
}
Bzip2Compressor::Bzip2Compressor(int compression_level) : m_compression_level(compression_level)
{
}
Bzip2Compressor::~Bzip2Compressor()
{
BZ2_bzCompressEnd(&m_stream);
}
bool Bzip2Compressor::Start()
{
ASSERT_MSG(DISCIO, m_stream.state == nullptr,
"Called Bzip2Compressor::Start() twice without calling Bzip2Compressor::End()");
m_buffer.clear();
m_stream.next_out = reinterpret_cast<char*>(m_buffer.data());
return BZ2_bzCompressInit(&m_stream, m_compression_level, 0, 0) == BZ_OK;
}
bool Bzip2Compressor::Compress(const u8* data, size_t size)
{
m_stream.next_in = reinterpret_cast<char*>(const_cast<u8*>(data));
m_stream.avail_in = static_cast<unsigned int>(size);
ExpandBuffer(size);
while (m_stream.avail_in != 0)
{
if (m_stream.avail_out == 0)
ExpandBuffer(0x100);
if (BZ2_bzCompress(&m_stream, BZ_RUN) != BZ_RUN_OK)
return false;
}
return true;
}
bool Bzip2Compressor::End()
{
bool success = true;
while (true)
{
if (m_stream.avail_out == 0)
ExpandBuffer(0x100);
const int result = BZ2_bzCompress(&m_stream, BZ_FINISH);
if (result != BZ_FINISH_OK && result != BZ_STREAM_END)
success = false;
if (result != BZ_FINISH_OK)
break;
}
if (BZ2_bzCompressEnd(&m_stream) != BZ_OK)
success = false;
return success;
}
void Bzip2Compressor::ExpandBuffer(size_t bytes_to_add)
{
const size_t bytes_written = GetSize();
m_buffer.resize(m_buffer.size() + bytes_to_add);
m_stream.next_out = reinterpret_cast<char*>(m_buffer.data()) + bytes_written;
m_stream.avail_out = static_cast<unsigned int>(m_buffer.size() - bytes_written);
}
const u8* Bzip2Compressor::GetData() const
{
return m_buffer.data();
}
size_t Bzip2Compressor::GetSize() const
{
return static_cast<size_t>(reinterpret_cast<u8*>(m_stream.next_out) - m_buffer.data());
}
LZMACompressor::LZMACompressor(bool lzma2, int compression_level, u8 compressor_data_out[7],
u8* compressor_data_size_out)
{
// lzma_lzma_preset returns false on success for some reason
if (lzma_lzma_preset(&m_options, static_cast<uint32_t>(compression_level)))
{
m_initialization_failed = true;
return;
}
if (!lzma2)
{
if (compressor_data_size_out)
*compressor_data_size_out = 5;
if (compressor_data_out)
{
ASSERT(m_options.lc < 9);
ASSERT(m_options.lp < 5);
ASSERT(m_options.pb < 5);
compressor_data_out[0] =
static_cast<u8>((m_options.pb * 5 + m_options.lp) * 9 + m_options.lc);
// The dictionary size is stored as a 32-bit little endian unsigned integer
static_assert(sizeof(m_options.dict_size) == sizeof(u32));
std::memcpy(compressor_data_out + 1, &m_options.dict_size, sizeof(u32));
}
}
else
{
if (compressor_data_size_out)
*compressor_data_size_out = 1;
if (compressor_data_out)
{
u8 encoded_dict_size = 0;
while (encoded_dict_size < 40 && m_options.dict_size > LZMA2DictionarySize(encoded_dict_size))
++encoded_dict_size;
compressor_data_out[0] = encoded_dict_size;
}
}
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;
}
LZMACompressor::~LZMACompressor()
{
lzma_end(&m_stream);
}
bool LZMACompressor::Start()
{
if (m_initialization_failed)
return false;
m_buffer.clear();
m_stream.next_out = m_buffer.data();
return lzma_raw_encoder(&m_stream, m_filters) == LZMA_OK;
}
bool LZMACompressor::Compress(const u8* data, size_t size)
{
m_stream.next_in = data;
m_stream.avail_in = size;
ExpandBuffer(size);
while (m_stream.avail_in != 0)
{
if (m_stream.avail_out == 0)
ExpandBuffer(0x100);
if (lzma_code(&m_stream, LZMA_RUN) != LZMA_OK)
return false;
}
return true;
}
bool LZMACompressor::End()
{
while (true)
{
if (m_stream.avail_out == 0)
ExpandBuffer(0x100);
switch (lzma_code(&m_stream, LZMA_FINISH))
{
case LZMA_OK:
break;
case LZMA_STREAM_END:
return true;
default:
return false;
}
}
}
void LZMACompressor::ExpandBuffer(size_t bytes_to_add)
{
const size_t bytes_written = GetSize();
m_buffer.resize(m_buffer.size() + bytes_to_add);
m_stream.next_out = m_buffer.data() + bytes_written;
m_stream.avail_out = m_buffer.size() - bytes_written;
}
const u8* LZMACompressor::GetData() const
{
return m_buffer.data();
}
size_t LZMACompressor::GetSize() const
{
return static_cast<size_t>(m_stream.next_out - m_buffer.data());
}
ZstdCompressor::ZstdCompressor(int compression_level)
{
m_stream = ZSTD_createCStream();
if (ZSTD_isError(ZSTD_CCtx_setParameter(m_stream, ZSTD_c_compressionLevel, compression_level)))
m_stream = nullptr;
}
ZstdCompressor::~ZstdCompressor()
{
ZSTD_freeCStream(m_stream);
}
bool ZstdCompressor::Start()
{
if (!m_stream)
return false;
m_buffer.clear();
m_out_buffer = {};
return !ZSTD_isError(ZSTD_CCtx_reset(m_stream, ZSTD_reset_session_only));
}
bool ZstdCompressor::Compress(const u8* data, size_t size)
{
ZSTD_inBuffer in_buffer{data, size, 0};
ExpandBuffer(size);
while (in_buffer.size != in_buffer.pos)
{
if (m_out_buffer.size == m_out_buffer.pos)
ExpandBuffer(0x100);
if (ZSTD_isError(ZSTD_compressStream(m_stream, &m_out_buffer, &in_buffer)))
return false;
}
return true;
}
bool ZstdCompressor::End()
{
while (true)
{
if (m_out_buffer.size == m_out_buffer.pos)
ExpandBuffer(0x100);
const size_t result = ZSTD_endStream(m_stream, &m_out_buffer);
if (ZSTD_isError(result))
return false;
if (result == 0)
return true;
}
}
void ZstdCompressor::ExpandBuffer(size_t bytes_to_add)
{
m_buffer.resize(m_buffer.size() + bytes_to_add);
m_out_buffer.dst = m_buffer.data();
m_out_buffer.size = m_buffer.size();
}
} // namespace DiscIO

View file

@ -0,0 +1,249 @@
// Copyright 2020 Dolphin Emulator Project
// Licensed under GPLv2+
// Refer to the license.txt file included.
#pragma once
#include <cstddef>
#include <memory>
#include <optional>
#include <vector>
#include <bzlib.h>
#include <lzma.h>
#include <mbedtls/sha1.h>
#include <zstd.h>
#include "Common/CommonTypes.h"
#include "DiscIO/LaggedFibonacciGenerator.h"
namespace DiscIO
{
struct DecompressionBuffer
{
std::vector<u8> data;
size_t bytes_written = 0;
};
using SHA1 = std::array<u8, 20>;
struct PurgeSegment
{
u32 offset;
u32 size;
};
static_assert(sizeof(PurgeSegment) == 0x08, "Wrong size for WIA purge segment");
class Decompressor
{
public:
virtual ~Decompressor();
virtual bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) = 0;
virtual bool Done() const { return m_done; };
protected:
bool m_done = false;
};
class NoneDecompressor final : public Decompressor
{
public:
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
};
// This class assumes that more bytes won't be added to in once in.bytes_written == in.data.size()
// and that *in_bytes_read initially will be equal to the size of the exception lists
class PurgeDecompressor final : public Decompressor
{
public:
PurgeDecompressor(u64 decompressed_size);
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
private:
const u64 m_decompressed_size;
PurgeSegment m_segment = {};
size_t m_bytes_read = 0;
size_t m_segment_bytes_written = 0;
size_t m_out_bytes_written = 0;
bool m_started = false;
mbedtls_sha1_context m_sha1_context;
};
class Bzip2Decompressor final : public Decompressor
{
public:
~Bzip2Decompressor();
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
private:
bz_stream m_stream = {};
bool m_started = false;
};
class LZMADecompressor final : public Decompressor
{
public:
LZMADecompressor(bool lzma2, const u8* filter_options, size_t filter_options_size);
~LZMADecompressor();
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) 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_error_occurred = false;
};
class ZstdDecompressor final : public Decompressor
{
public:
ZstdDecompressor();
~ZstdDecompressor();
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
private:
ZSTD_DStream* m_stream;
};
class RVZPackDecompressor final : public Decompressor
{
public:
RVZPackDecompressor(std::unique_ptr<Decompressor> decompressor, DecompressionBuffer decompressed,
u64 data_offset);
bool Decompress(const DecompressionBuffer& in, DecompressionBuffer* out,
size_t* in_bytes_read) override;
bool Done() const override;
private:
std::optional<bool> ReadToDecompressed(const DecompressionBuffer& in, size_t* in_bytes_read,
size_t decompressed_bytes_read, size_t bytes_to_read);
std::unique_ptr<Decompressor> m_decompressor;
DecompressionBuffer m_decompressed;
size_t m_decompressed_bytes_read = 0;
u64 m_data_offset;
u32 m_size = 0;
bool m_junk;
LaggedFibonacciGenerator m_lfg;
};
class Compressor
{
public:
virtual ~Compressor();
// First call Start, then AddDataOnlyForPurgeHashing/Compress any number of times,
// then End, then GetData/GetSize any number of times.
virtual bool Start() = 0;
virtual bool AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) { return true; }
virtual bool Compress(const u8* data, size_t size) = 0;
virtual bool End() = 0;
virtual const u8* GetData() const = 0;
virtual size_t GetSize() const = 0;
};
class PurgeCompressor final : public Compressor
{
public:
PurgeCompressor();
~PurgeCompressor();
bool Start() override;
bool AddPrecedingDataOnlyForPurgeHashing(const u8* data, size_t size) override;
bool Compress(const u8* data, size_t size) override;
bool End() override;
const u8* GetData() const override;
size_t GetSize() const override;
private:
std::vector<u8> m_buffer;
size_t m_bytes_written;
mbedtls_sha1_context m_sha1_context;
};
class Bzip2Compressor final : public Compressor
{
public:
Bzip2Compressor(int compression_level);
~Bzip2Compressor();
bool Start() override;
bool Compress(const u8* data, size_t size) override;
bool End() override;
const u8* GetData() const override;
size_t GetSize() const override;
private:
void ExpandBuffer(size_t bytes_to_add);
bz_stream m_stream = {};
std::vector<u8> m_buffer;
int m_compression_level;
};
class LZMACompressor final : public Compressor
{
public:
LZMACompressor(bool lzma2, int compression_level, u8 compressor_data_out[7],
u8* compressor_data_size_out);
~LZMACompressor();
bool Start() override;
bool Compress(const u8* data, size_t size) override;
bool End() override;
const u8* GetData() const override;
size_t GetSize() const override;
private:
void ExpandBuffer(size_t bytes_to_add);
lzma_stream m_stream = LZMA_STREAM_INIT;
lzma_options_lzma m_options = {};
lzma_filter m_filters[2];
std::vector<u8> m_buffer;
bool m_initialization_failed = false;
};
class ZstdCompressor final : public Compressor
{
public:
ZstdCompressor(int compression_level);
~ZstdCompressor();
bool Start() override;
bool Compress(const u8* data, size_t size) override;
bool End() override;
const u8* GetData() const override { return m_buffer.data(); }
size_t GetSize() const override { return m_out_buffer.pos; }
private:
void ExpandBuffer(size_t bytes_to_add);
ZSTD_CStream* m_stream;
ZSTD_outBuffer m_out_buffer;
std::vector<u8> m_buffer;
};
} // namespace DiscIO