loader/strat: fix/refactor zstd-zbic integration

This commit is contained in:
Alula 2026-06-24 20:50:41 +02:00
parent 7b6ee4916e
commit 38251801df
No known key found for this signature in database
GPG Key ID: 3E00485503A1D8BA
5 changed files with 66 additions and 31 deletions

View File

@ -21,11 +21,13 @@ namespace ams::util {
/* Compression utilities. */
int CompressLZ4(void *dst, size_t dst_size, const void *src, size_t src_size);
size_t CompressZstd(void *dst, size_t dst_size, const void *src, size_t src_size);
size_t CompressZbic(void *dst, size_t dst_size, const void *src, size_t src_size);
constexpr size_t ZstdDctxWorkspaceSize = 0x176E8;
/* Decompression utilities. */
int DecompressLZ4(void *dst, size_t dst_size, const void *src, size_t src_size);
size_t DecompressZstd(void *dst, size_t dst_size, const void *src, size_t src_size);
bool DecompressZstdForLoader(void* workspace, size_t workspace_size, void *dst, size_t dst_size, size_t expected_dec_size, const void *src, size_t src_size);
size_t DecompressZbic(void *dst, size_t dst_size, const void *src, size_t src_size);
bool DecompressZbicForLoader(void* workspace, size_t workspace_size, void *dst, size_t dst_size, size_t expected_dec_size, const void *src, size_t src_size);
}

View File

@ -0,0 +1,41 @@
/*
* Copyright (c) Atmosphère-NX
*
* This program is free software; you can redistribute it and/or modify it
* under the terms and conditions of the GNU General Public License,
* version 2, as published by the Free Software Foundation.
*
* This program is distributed in the hope it will be useful, but WITHOUT
* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
* FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
* more details.
*
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
#include <stratosphere.hpp>
#include "lz4.h"
namespace ams::util {
/* Compression utilities. */
int CompressLZ4(void *dst, size_t dst_size, const void *src, size_t src_size) {
/* Size checks. */
AMS_ABORT_UNLESS(dst_size <= std::numeric_limits<int>::max());
AMS_ABORT_UNLESS(src_size <= std::numeric_limits<int>::max());
/* This is just a thin wrapper around LZ4. */
return LZ4_compress_default(reinterpret_cast<const char *>(src), reinterpret_cast<char *>(dst), static_cast<int>(src_size), static_cast<int>(dst_size));
}
/* Decompression utilities. */
int DecompressLZ4(void *dst, size_t dst_size, const void *src, size_t src_size) {
/* Size checks. */
AMS_ABORT_UNLESS(dst_size <= std::numeric_limits<int>::max());
AMS_ABORT_UNLESS(src_size <= std::numeric_limits<int>::max());
/* This is just a thin wrapper around LZ4. */
return LZ4_decompress_safe(reinterpret_cast<const char *>(src), reinterpret_cast<char *>(dst), static_cast<int>(src_size), static_cast<int>(dst_size));
}
}

View File

@ -17,21 +17,19 @@
#include "lz4.h"
#define ZSTD_STATIC_LINKING_ONLY
#define ZSTD_ZBIC_SUPPORT 1
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-function"
#define ZSTDLIB_VISIBLE static
#define ZSTDLIB_HIDDEN static
#include "zstd.h"
#include "zstd.inc"
#pragma GCC diagnostic pop
namespace ams::util {
static_assert(sizeof(ZSTD_DCtx) <= ZstdDctxWorkspaceSize);
/* Compression utilities. */
int CompressLZ4(void *dst, size_t dst_size, const void *src, size_t src_size) {
/* Size checks. */
AMS_ABORT_UNLESS(dst_size <= std::numeric_limits<int>::max());
AMS_ABORT_UNLESS(src_size <= std::numeric_limits<int>::max());
/* This is just a thin wrapper around LZ4. */
return LZ4_compress_default(reinterpret_cast<const char *>(src), reinterpret_cast<char *>(dst), static_cast<int>(src_size), static_cast<int>(dst_size));
}
size_t CompressZstd(void *dst, size_t dst_size, const void *src, size_t src_size) {
size_t CompressZbic(void *dst, size_t dst_size, const void *src, size_t src_size) {
/* Basic size checks. */
AMS_ABORT_UNLESS(dst_size <= std::numeric_limits<int>::max());
AMS_ABORT_UNLESS(src_size <= std::numeric_limits<int>::max());
@ -48,17 +46,7 @@ namespace ams::util {
return ZSTD_compress(dst, dst_size, src, src_size, compressionLevel);
}
/* Decompression utilities. */
int DecompressLZ4(void *dst, size_t dst_size, const void *src, size_t src_size) {
/* Size checks. */
AMS_ABORT_UNLESS(dst_size <= std::numeric_limits<int>::max());
AMS_ABORT_UNLESS(src_size <= std::numeric_limits<int>::max());
/* This is just a thin wrapper around LZ4. */
return LZ4_decompress_safe(reinterpret_cast<const char *>(src), reinterpret_cast<char *>(dst), static_cast<int>(src_size), static_cast<int>(dst_size));
}
size_t DecompressZstd(void *dst, size_t dst_size, const void *src, size_t src_size) {
size_t DecompressZbic(void *dst, size_t dst_size, const void *src, size_t src_size) {
/* Basic size checks. */
AMS_ABORT_UNLESS(dst_size <= std::numeric_limits<int>::max());
AMS_ABORT_UNLESS(src_size <= std::numeric_limits<int>::max());
@ -72,10 +60,12 @@ namespace ams::util {
return ZSTD_decompress(dst, dst_size, src, src_size);
}
bool DecompressZstdForLoader(void* workspace, size_t workspace_size, void *dst, size_t dst_size, size_t expected_dec_size, const void *src, size_t src_size) {
bool DecompressZbicForLoader(void* workspace, size_t workspace_size, void *dst, size_t dst_size, size_t expected_dec_size, const void *src, size_t src_size) {
/* Check decompression margin. */
auto margin = ZSTD_decompressionMargin(src, src_size);
if (ZSTD_isError(margin)) {
auto ec = ZSTD_getErrorCode(margin);
AMS_LOG("[ldr] can't determine decompression margin: %u (%s)\n", ec, ZSTD_getErrorString(ec));
return false;
}
@ -86,17 +76,20 @@ namespace ams::util {
/* Make sure we fit in the destination buffer. */
if (margin + expected_dec_size > dst_size) {
AMS_LOG("[ldr] not enough space for decompression %lu + %lu > %lu\n", margin, expected_dec_size, dst_size);
return false;
}
/* This is a runtime assert in Loader code. We replicate it here. */
AMS_ABORT_UNLESS(ZSTD_estimateDCtxSize() == workspace_size);
/* This is a runtime assert in Loader code. Note that Nintendo does == comparison here. */
AMS_ABORT_UNLESS(ZSTD_estimateDCtxSize() <= workspace_size);
/* Decompress using a static decompression context. */
auto dctx = ZSTD_initStaticDCtx(workspace, workspace_size);
size_t dec_size = ZSTD_decompressDCtx(dctx, dst, dst_size, src, src_size);
if (ZSTD_isError(dec_size)) {
auto ec = ZSTD_getErrorCode(dec_size);
AMS_LOG("[ldr] decompression failed: %u (%s)\n", ec, ZSTD_getErrorString(ec));
return false;
}

View File

@ -121,8 +121,7 @@ namespace ams::ldr {
NsoHeader g_nso_headers[Nso_Count];
/* Global Zstd decompression context. */
constexpr size_t ZstdDctxWorkspaceSize = 0x176E8;
alignas(8) u8 g_zstd_dctx_workspace[ZstdDctxWorkspaceSize];
alignas(8) u8 g_zstd_dctx_workspace[util::ZstdDctxWorkspaceSize];
Result ValidateProgramVersion(ncm::ProgramId program_id, u32 version) {
/* No version verification is done before 8.1.0. */
@ -656,7 +655,7 @@ namespace ams::ldr {
auto compressed_data_buf = reinterpret_cast<const void *>(load_address);
if (is_zstd) {
bool decompressed = util::DecompressZstdForLoader(reinterpret_cast<void *>(g_zstd_dctx_workspace), ZstdDctxWorkspaceSize, reinterpret_cast<void *>(map_base), static_cast<size_t>(map_end - map_base), segment_size, compressed_data_buf, file_size);
bool decompressed = util::DecompressZbicForLoader(reinterpret_cast<void *>(g_zstd_dctx_workspace), sizeof(g_zstd_dctx_workspace), reinterpret_cast<void *>(map_base), static_cast<size_t>(map_end - map_base), segment_size, compressed_data_buf, file_size);
R_UNLESS(decompressed, ldr::ResultInvalidNso());
} else {
bool decompressed = (util::DecompressLZ4(reinterpret_cast<void *>(map_base), segment_size, compressed_data_buf, file_size) == static_cast<int>(segment_size));