Merge pull request #2753 from alula/22_support

Basic ("the console boots and functions as expected") 22.0.0 support - Part 2

libraries/libstratosphere/include/stratosphere/ldr/ldr_types.hpp

@@ -94,13 +94,14 @@ namespace ams::ldr {
         };
 
         enum Flag : u32 {
-            Flag_CompressedText   = (1 << 0),
+            Flag_CompressedText     = (1 << 0),
-            Flag_CompressedRo     = (1 << 1),
+            Flag_CompressedRo       = (1 << 1),
-            Flag_CompressedRw     = (1 << 2),
+            Flag_CompressedRw       = (1 << 2),
-            Flag_CheckHashText    = (1 << 3),
+            Flag_CheckHashText      = (1 << 3),
-            Flag_CheckHashRo      = (1 << 4),
+            Flag_CheckHashRo        = (1 << 4),
-            Flag_CheckHashRw      = (1 << 5),
+            Flag_CheckHashRw        = (1 << 5),
-            Flag_PreventCodeReads = (1 << 6),
+            Flag_PreventCodeReads   = (1 << 6),
+            Flag_UseZbicCompression = (1 << 7),
         };
 
         struct SegmentInfo {

libraries/libstratosphere/include/stratosphere/util/util_zbic_for_loader.hpp

@@ -0,0 +1,59 @@
+/*
+ * 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/>.
+ */
+
+#pragma once
+#include <vapours.hpp>
+#include <stratosphere/diag.hpp>
+
+#define ZSTD_ZBIC 1
+#define ZSTD_STATIC_LINKING_ONLY
+
+#include "zstd.h"
+
+namespace ams::util {
+    
+    constexpr size_t DCtxWorkspaceSize = 0x176E8;
+
+    inline bool DecompressZbicForLoader(void *workspace, void *map_base, size_t map_size, size_t segment_size, size_t compressed_size, const void *compressed_data_buf) {
+        /* TODO: how to assert that workspace >= DCtxWorkspaceSize? */
+
+        /* Check decompression margin */
+        auto margin           = ZSTD_decompressionMargin(compressed_data_buf, compressed_size);
+
+        if(ZSTD_isError(margin))                               return false;
+        if(!util::CanAddWithoutOverflow(margin, segment_size)) return false;
+        if(margin + segment_size > map_size)                   return false;
+
+        AMS_ABORT_UNLESS(ZSTD_estimateDCtxSize() == DCtxWorkspaceSize); /* Why is this a runtime assert in N's code? */
+
+        auto ctx        = ZSTD_initStaticDCtx(workspace, DCtxWorkspaceSize);
+        size_t dec_size = ZSTD_decompressDCtx(ctx, map_base, map_size, compressed_data_buf, compressed_size);
+
+        if (ZSTD_isError(dec_size)) {
+            AMS_LOG("[ldr] ZSTD_decompressDCtx failed: %ld\n", dec_size);
+            return false;
+        }
+        if (dec_size != segment_size) return false;
+
+        return true;
+    }
+}
+
+#ifdef AMS_ZSTD_IMPLEMENTATION
+#include "zbicdeclib.inc"
+
+static_assert(sizeof(ZSTD_DCtx) == ams::util::DCtxWorkspaceSize);
+#endif

libraries/libstratosphere/include/stratosphere/util/zstd_errors.h

@@ -0,0 +1,107 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under both the BSD-style license (found in the
+ * LICENSE file in the root directory of this source tree) and the GPLv2 (found
+ * in the COPYING file in the root directory of this source tree).
+ * You may select, at your option, one of the above-listed licenses.
+ */
+
+#ifndef ZSTD_ERRORS_H_398273423
+#define ZSTD_ERRORS_H_398273423
+
+#if defined (__cplusplus)
+extern "C" {
+#endif
+
+/* =====   ZSTDERRORLIB_API : control library symbols visibility   ===== */
+#ifndef ZSTDERRORLIB_VISIBLE
+   /* Backwards compatibility with old macro name */
+#  ifdef ZSTDERRORLIB_VISIBILITY
+#    define ZSTDERRORLIB_VISIBLE ZSTDERRORLIB_VISIBILITY
+#  elif defined(__GNUC__) && (__GNUC__ >= 4) && !defined(__MINGW32__)
+#    define ZSTDERRORLIB_VISIBLE __attribute__ ((visibility ("default")))
+#  else
+#    define ZSTDERRORLIB_VISIBLE
+#  endif
+#endif
+
+#ifndef ZSTDERRORLIB_HIDDEN
+#  if defined(__GNUC__) && (__GNUC__ >= 4) && !defined(__MINGW32__)
+#    define ZSTDERRORLIB_HIDDEN __attribute__ ((visibility ("hidden")))
+#  else
+#    define ZSTDERRORLIB_HIDDEN
+#  endif
+#endif
+
+#if defined(ZSTD_DLL_EXPORT) && (ZSTD_DLL_EXPORT==1)
+#  define ZSTDERRORLIB_API __declspec(dllexport) ZSTDERRORLIB_VISIBLE
+#elif defined(ZSTD_DLL_IMPORT) && (ZSTD_DLL_IMPORT==1)
+#  define ZSTDERRORLIB_API __declspec(dllimport) ZSTDERRORLIB_VISIBLE /* It isn't required but allows to generate better code, saving a function pointer load from the IAT and an indirect jump.*/
+#else
+#  define ZSTDERRORLIB_API ZSTDERRORLIB_VISIBLE
+#endif
+
+/*-*********************************************
+ *  Error codes list
+ *-*********************************************
+ *  Error codes _values_ are pinned down since v1.3.1 only.
+ *  Therefore, don't rely on values if you may link to any version < v1.3.1.
+ *
+ *  Only values < 100 are considered stable.
+ *
+ *  note 1 : this API shall be used with static linking only.
+ *           dynamic linking is not yet officially supported.
+ *  note 2 : Prefer relying on the enum than on its value whenever possible
+ *           This is the only supported way to use the error list < v1.3.1
+ *  note 3 : ZSTD_isError() is always correct, whatever the library version.
+ **********************************************/
+typedef enum {
+  ZSTD_error_no_error = 0,
+  ZSTD_error_GENERIC  = 1,
+  ZSTD_error_prefix_unknown                = 10,
+  ZSTD_error_version_unsupported           = 12,
+  ZSTD_error_frameParameter_unsupported    = 14,
+  ZSTD_error_frameParameter_windowTooLarge = 16,
+  ZSTD_error_corruption_detected = 20,
+  ZSTD_error_checksum_wrong      = 22,
+  ZSTD_error_literals_headerWrong = 24,
+  ZSTD_error_dictionary_corrupted      = 30,
+  ZSTD_error_dictionary_wrong          = 32,
+  ZSTD_error_dictionaryCreation_failed = 34,
+  ZSTD_error_parameter_unsupported   = 40,
+  ZSTD_error_parameter_combination_unsupported = 41,
+  ZSTD_error_parameter_outOfBound    = 42,
+  ZSTD_error_tableLog_tooLarge       = 44,
+  ZSTD_error_maxSymbolValue_tooLarge = 46,
+  ZSTD_error_maxSymbolValue_tooSmall = 48,
+  ZSTD_error_cannotProduce_uncompressedBlock = 49,
+  ZSTD_error_stabilityCondition_notRespected = 50,
+  ZSTD_error_stage_wrong       = 60,
+  ZSTD_error_init_missing      = 62,
+  ZSTD_error_memory_allocation = 64,
+  ZSTD_error_workSpace_tooSmall= 66,
+  ZSTD_error_dstSize_tooSmall = 70,
+  ZSTD_error_srcSize_wrong    = 72,
+  ZSTD_error_dstBuffer_null   = 74,
+  ZSTD_error_noForwardProgress_destFull = 80,
+  ZSTD_error_noForwardProgress_inputEmpty = 82,
+  /* following error codes are __NOT STABLE__, they can be removed or changed in future versions */
+  ZSTD_error_frameIndex_tooLarge = 100,
+  ZSTD_error_seekableIO          = 102,
+  ZSTD_error_dstBuffer_wrong     = 104,
+  ZSTD_error_srcBuffer_wrong     = 105,
+  ZSTD_error_sequenceProducer_failed = 106,
+  ZSTD_error_externalSequences_invalid = 107,
+  ZSTD_error_maxCode = 120  /* never EVER use this value directly, it can change in future versions! Use ZSTD_isError() instead */
+} ZSTD_ErrorCode;
+
+ZSTDERRORLIB_API const char* ZSTD_getErrorString(ZSTD_ErrorCode code);   /**< Same as ZSTD_getErrorName, but using a `ZSTD_ErrorCode` enum argument */
+
+
+#if defined (__cplusplus)
+}
+#endif
+
+#endif /* ZSTD_ERRORS_H_398273423 */

stratosphere/loader/source/ldr_process_creation.cpp

@@ -22,6 +22,7 @@
 #include "ldr_patcher.hpp"
 #include "ldr_process_creation.hpp"
 #include "ldr_ro_manager.hpp"
+#include <stratosphere/util/util_zbic_for_loader.hpp>
 
 namespace ams::ldr {
 
@@ -92,6 +93,8 @@ namespace ams::ldr {
 
         struct ProcessInfo {
             os::NativeHandle process_handle;
+            uintptr_t code_address;
+            size_t    total_size;
             uintptr_t args_address;
             size_t    args_size;
             uintptr_t nso_address[Nso_Count];
@@ -103,12 +106,24 @@ namespace ams::ldr {
             bool has_main;
             bool has_sdk;
             bool has_subsdk;
-            bool has_nso[Nso_Count];
+            s8   nso_indices[Nso_Count];
+        };
+
+        struct AutoLoadModuleContext {
+            NsoHeader *headers;
+            int        nso_count;
+            int        rtld_idx;
+            int        main_nso_idx;
+            int        sdk_nso_idx;
+            AutoLoadModuleInfo ali;
         };
 
         /* Global NSO header cache. */
         NsoHeader g_nso_headers[Nso_Count];
 
+        /* Global zstd decompression context */
+        alignas(8) u8 g_zstd_dctx_workspace[0x176E8];
+
         Result ValidateProgramVersion(ncm::ProgramId program_id, u32 version) {
             /* No version verification is done before 8.1.0. */
             R_SUCCEED_IF(hos::GetVersion() < hos::Version_8_1_0);
@@ -169,10 +184,15 @@ namespace ams::ldr {
             return static_cast<Acid::PoolPartition>((meta->acid->flags & Acid::AcidFlag_PoolPartitionMask) >> Acid::AcidFlag_PoolPartitionShift);
         }
 
-        Result LoadAutoLoadHeaders(NsoHeader *nso_headers, AutoLoadModuleInfo *ali, u32 acid_flags) {
+        Result LoadAutoLoadHeaders(AutoLoadModuleContext &ctx, u32 acid_flags) {
             /* Clear NSOs. */
-            std::memset(nso_headers, 0, sizeof(*nso_headers) * Nso_Count);
+            std::memset(g_nso_headers, 0, sizeof(g_nso_headers));
-            *ali = {};
+            ctx.headers      = g_nso_headers;
+            ctx.nso_count    = 0;
+            ctx.rtld_idx     = -1;
+            ctx.main_nso_idx = -1;
+            ctx.sdk_nso_idx  = -1;
+            ctx.ali          = {};
 
             for (size_t i = 0; i < Nso_Count; i++) {
                 /* Only load browser DLLs if acid flags say to do so. */
@@ -199,16 +219,18 @@ namespace ams::ldr {
 
                     /* Read NSO header. */
                     size_t read_size;
-                    R_TRY(fs::ReadFile(std::addressof(read_size), file, 0, nso_headers + i, sizeof(*nso_headers)));
+                    R_TRY(fs::ReadFile(std::addressof(read_size), file, 0, g_nso_headers + ctx.nso_count, sizeof(NsoHeader)));
-                    R_UNLESS(read_size == sizeof(*nso_headers), ldr::ResultInvalidNso());
+                    R_UNLESS(read_size == sizeof(NsoHeader), ldr::ResultInvalidNso());
 
                     /* Note nso is present. */
                     switch (i) {
                         case Nso_Rtld:
-                            ali->has_rtld = true;
+                            ctx.rtld_idx = ctx.nso_count;
+                            ctx.ali.has_rtld = true;
                             break;
                         case Nso_Main:
-                            ali->has_main = true;
+                            ctx.main_nso_idx = ctx.nso_count;
+                            ctx.ali.has_main = true;
                             break;
                         case Nso_SubSdk0:
                         case Nso_SubSdk1:
@@ -220,64 +242,65 @@ namespace ams::ldr {
                         case Nso_SubSdk7:
                         case Nso_SubSdk8:
                         case Nso_SubSdk9:
-                            ali->has_subsdk = true;
+                            ctx.ali.has_subsdk = true;
                             break;
                         case Nso_Sdk:
-                            ali->has_sdk = true;
+                            ctx.sdk_nso_idx = ctx.nso_count;
+                            ctx.ali.has_sdk = true;
                             break;
                     }
-                    ali->has_nso[i] = true;
+                    ctx.ali.nso_indices[ctx.nso_count] = static_cast<s8>(i);
+                    ctx.nso_count++;
                 }
             }
 
             R_SUCCEED();
         }
 
-        Result CheckAutoLoad(const NsoHeader *nso_headers, const AutoLoadModuleInfo *ali, u32 acid_flags) {
+        Result CheckAutoLoad(const AutoLoadModuleContext &ctx, u32 acid_flags) {
             /* We must always have a main. */
-            R_UNLESS(ali->has_main, ldr::ResultInvalidNso());
+            R_UNLESS(ctx.ali.has_main, ldr::ResultInvalidNso());
 
-            /* All NSOs must not be --X. */
+            /* Validate flags and extents for all present NSOs. */
-            /* This is "probably" not checked on Ounce? */
+            for (int i = 0; i < ctx.nso_count; ++i) {
-            for (size_t i = 0; i < Nso_Count; ++i) {
+                const auto &hdr = ctx.headers[i];
-                R_UNLESS((nso_headers[i].flags & NsoHeader::Flag_PreventCodeReads) == 0, ldr::ResultInvalidNso());
-            }
 
-            /* If we don't have an RTLD, we must only have a main. */
+                /* All NSOs must not be --X. */
-            const bool has_browser_dll = (acid_flags & Acid::AcidFlag_LoadBrowserCoreDll) != 0;
+                /* This is "probably" not checked on Ounce? */
-            if (!ali->has_rtld) {
+                R_UNLESS((hdr.flags & NsoHeader::Flag_PreventCodeReads) == 0, ldr::ResultInvalidNso());
-                /* If don't have rtld, we must also not have sdk. */
-                R_UNLESS(!ali->has_sdk, ldr::ResultInvalidNso());
 
-                /* We must also not have both subsdk and browser dll. */
+                /* Zstd compression only allowed on main, and only when both rtld+sdk are present. */
-                R_UNLESS(!(ali->has_subsdk && has_browser_dll), ldr::ResultInvalidNso());
+                if (i != ctx.main_nso_idx || ctx.rtld_idx < 0 || ctx.sdk_nso_idx < 0) {
-            } else {
+                    R_UNLESS((hdr.flags & NsoHeader::Flag_UseZbicCompression) == 0, ldr::ResultInvalidNso());
-                /* If we have rtld, we must not have browser core dll. */
-                R_UNLESS(!has_browser_dll, ldr::ResultInvalidNso());
-            }
-
-            /* Check NSO extents. */
-            for (size_t i = 0; i < Nso_Count; i++) {
-                /* Only validate the nsos we have. */
-                if (!ali->has_nso[i]) {
-                    continue;
                 }
 
                 /* NSOs must have page-aligned segments. */
-                R_UNLESS(util::IsAligned(nso_headers[i].text_dst_offset, os::MemoryPageSize), ldr::ResultInvalidNso());
+                R_UNLESS(util::IsAligned(hdr.text_dst_offset, os::MemoryPageSize), ldr::ResultInvalidNso());
-                R_UNLESS(util::IsAligned(nso_headers[i].ro_dst_offset,   os::MemoryPageSize), ldr::ResultInvalidNso());
+                R_UNLESS(util::IsAligned(hdr.ro_dst_offset,   os::MemoryPageSize), ldr::ResultInvalidNso());
-                R_UNLESS(util::IsAligned(nso_headers[i].rw_dst_offset,   os::MemoryPageSize), ldr::ResultInvalidNso());
+                R_UNLESS(util::IsAligned(hdr.rw_dst_offset,   os::MemoryPageSize), ldr::ResultInvalidNso());
 
                 /* NSOs must have zero text offset. */
-                R_UNLESS(nso_headers[i].text_dst_offset == 0, ldr::ResultInvalidNso());
+                R_UNLESS(hdr.text_dst_offset == 0, ldr::ResultInvalidNso());
 
                 /* NSO .text must precede .rodata. */
-                const size_t text_end = static_cast<size_t>(nso_headers[i].text_dst_offset) + static_cast<size_t>(nso_headers[i].text_size);
+                const size_t text_end = static_cast<size_t>(hdr.text_dst_offset) + static_cast<size_t>(hdr.text_size);
-                R_UNLESS(text_end <= static_cast<size_t>(nso_headers[i].ro_dst_offset), ldr::ResultInvalidNso());
+                R_UNLESS(text_end <= static_cast<size_t>(hdr.ro_dst_offset), ldr::ResultInvalidNso());
 
                 /* NSO .rodata must precede .rwdata. */
-                const size_t ro_end   = static_cast<size_t>(nso_headers[i].ro_dst_offset)   + static_cast<size_t>(nso_headers[i].ro_size);
+                const size_t ro_end = static_cast<size_t>(hdr.ro_dst_offset) + static_cast<size_t>(hdr.ro_size);
-                R_UNLESS(ro_end <= static_cast<size_t>(nso_headers[i].rw_dst_offset), ldr::ResultInvalidNso());
+                R_UNLESS(ro_end <= static_cast<size_t>(hdr.rw_dst_offset), ldr::ResultInvalidNso());
+            }
+
+            const bool has_browser_dll = (acid_flags & Acid::AcidFlag_LoadBrowserCoreDll) != 0;
+            if (ctx.ali.has_rtld || ctx.ali.has_sdk) {
+                /* If we have sdk we must have rtld. */
+                R_UNLESS(ctx.ali.has_rtld, ldr::ResultInvalidNso());
+
+                /* If we have rtld, we must not have browser core dll. */
+                R_UNLESS(!has_browser_dll, ldr::ResultInvalidNso());
+            } else {
+                /* We must not have both subsdk and browser dll. */
+                R_UNLESS(!(ctx.ali.has_subsdk && has_browser_dll), ldr::ResultInvalidNso());
             }
 
             R_SUCCEED();
@@ -293,8 +316,8 @@ namespace ams::ldr {
             { 0x010049900F546001 }, /* Super Mario 3D All-Stars: Super Mario 64 */
             { 0x010057D00ECE4000 }, /* Nintendo Switch Online (Nintendo 64) [for Japan] */
             { 0x01006F8002326000 }, /* Animal Crossing: New Horizons */
-            { 0x01006FB00F50E000 }, /* [???] */
+            { 0x01006FB00F50E000 }, /* 宝可梦 走吧！伊布 [Pokemon: Let's Go, Eevee! for China] */
-            { 0x010070300F50C000 }, /* [???] */
+            { 0x010070300F50C000 }, /* 宝可梦 走吧！皮卡丘 [Pokemon: Let's Go, Pikachu! for China] */
             { 0x010075100E8EC000 }, /* 马力欧卡丁车8 豪华版 [Mario Kart 8 Deluxe for China] */
             { 0x01008DB008C2C000 }, /* Pokemon Shield */
             { 0x01009AD008C4C000 }, /* Pokemon: Let's Go, Pikachu! [Kiosk] */
@@ -514,7 +537,7 @@ namespace ams::ldr {
             return rand % (max + 1);
         }
 
-        Result DecideAddressSpaceLayout(ProcessInfo *out, svc::CreateProcessParameter *out_param, const NsoHeader *nso_headers, const AutoLoadModuleInfo *ali, const ArgumentStore::Entry *argument) {
+        Result DecideAddressSpaceLayout(ProcessInfo *out, svc::CreateProcessParameter *out_param, const AutoLoadModuleContext &ctx, const ArgumentStore::Entry *argument) {
             /* Clear output. */
             out->args_address = 0;
             out->args_size = 0;
@@ -525,35 +548,33 @@ namespace ams::ldr {
             bool argument_allocated = false;
 
             /* Calculate base offsets. */
-            for (size_t i = 0; i < Nso_Count; i++) {
+            for (int i = 0; i < ctx.nso_count; i++) {
-                if (ali->has_nso[i]) {
+                out->nso_address[i] = total_size;
-                    out->nso_address[i] = total_size;
+                const size_t text_end = static_cast<size_t>(ctx.headers[i].text_dst_offset) + static_cast<size_t>(ctx.headers[i].text_size);
-                    const size_t text_end = static_cast<size_t>(nso_headers[i].text_dst_offset) + static_cast<size_t>(nso_headers[i].text_size);
+                const size_t ro_end   = static_cast<size_t>(ctx.headers[i].ro_dst_offset)   + static_cast<size_t>(ctx.headers[i].ro_size);
-                    const size_t ro_end   = static_cast<size_t>(nso_headers[i].ro_dst_offset)   + static_cast<size_t>(nso_headers[i].ro_size);
+                const size_t rw_end   = static_cast<size_t>(ctx.headers[i].rw_dst_offset)   + static_cast<size_t>(ctx.headers[i].rw_size);
-                    const size_t rw_end   = static_cast<size_t>(nso_headers[i].rw_dst_offset)   + static_cast<size_t>(nso_headers[i].rw_size);
+                out->nso_size[i] = text_end;
-                    out->nso_size[i] = text_end;
+                out->nso_size[i] = std::max(out->nso_size[i], ro_end);
-                    out->nso_size[i] = std::max(out->nso_size[i], ro_end);
+                out->nso_size[i] = std::max(out->nso_size[i], rw_end);
-                    out->nso_size[i] = std::max(out->nso_size[i], rw_end);
+                out->nso_size[i] += static_cast<size_t>(ctx.headers[i].bss_size);
-                    out->nso_size[i] += static_cast<size_t>(nso_headers[i].bss_size);
 
-                    const size_t aligned_up_size = util::AlignUp(out->nso_size[i], os::MemoryPageSize) & (AutoLoadModuleSizeMax - 1);
+                const size_t aligned_up_size = util::AlignUp(out->nso_size[i], os::MemoryPageSize) & (AutoLoadModuleSizeMax - 1);
-                    R_UNLESS(out->nso_size[i] <= aligned_up_size, ldr::ResultInvalidNso());
+                R_UNLESS(out->nso_size[i] <= aligned_up_size, ldr::ResultInvalidNso());
-                    R_UNLESS(aligned_up_size > 0,                 ldr::ResultInvalidNso());
+                R_UNLESS(aligned_up_size > 0,                 ldr::ResultInvalidNso());
 
-                    out->nso_size[i] = aligned_up_size;
+                out->nso_size[i] = aligned_up_size;
 
-                    R_UNLESS(util::CanAddWithoutOverflow(total_size, out->nso_size[i]), ldr::ResultInvalidNso());
+                R_UNLESS(util::CanAddWithoutOverflow(total_size, out->nso_size[i]), ldr::ResultInvalidNso());
-                    total_size += out->nso_size[i];
+                total_size += out->nso_size[i];
 
-                    if (!argument_allocated && argument != nullptr) {
+                if (!argument_allocated && argument != nullptr) {
-                        out->args_address = total_size;
+                    out->args_address = total_size;
-                        out->args_size    = util::AlignUp(2 * sizeof(u32) + argument->argument_size * 2 + ArgumentStore::ArgumentBufferSize, os::MemoryPageSize);
+                    out->args_size    = util::AlignUp(2 * sizeof(u32) + argument->argument_size * 2 + ArgumentStore::ArgumentBufferSize, os::MemoryPageSize);
 
-                        R_UNLESS(util::CanAddWithoutOverflow(total_size, out->args_size), ldr::ResultInvalidNso());
+                    R_UNLESS(util::CanAddWithoutOverflow(total_size, out->args_size), ldr::ResultInvalidNso());
-                        total_size += out->args_size;
+                    total_size += out->args_size;
 
-                        argument_allocated = true;
+                    argument_allocated = true;
-                    }
                 }
             }
 
@@ -598,11 +619,9 @@ namespace ams::ldr {
 
             /* Set out. */
             aslr_start += aslr_slide;
-            for (size_t i = 0; i < Nso_Count; i++) {
+            for (int i = 0; i < ctx.nso_count; i++) {
-                if (ali->has_nso[i]) {
+                R_UNLESS(util::CanAddWithoutOverflow(out->nso_address[i], aslr_start), ldr::ResultInvalidNso());
-                    R_UNLESS(util::CanAddWithoutOverflow(out->nso_address[i], aslr_start), ldr::ResultInvalidNso());
+                out->nso_address[i] += aslr_start;
-                    out->nso_address[i] += aslr_start;
-                }
             }
             if (out->args_address) {
                 R_UNLESS(util::CanAddWithoutOverflow(out->args_address, aslr_start), ldr::ResultInvalidNso());
@@ -611,69 +630,90 @@ namespace ams::ldr {
 
             out_param->code_address   = aslr_start;
             out_param->code_num_pages = total_size >> 12;
+            out->total_size = total_size;
 
             R_SUCCEED();
         }
 
-        Result LoadAutoLoadModuleSegment(fs::FileHandle file, const NsoHeader::SegmentInfo *segment, size_t file_size, const u8 *file_hash, bool is_compressed, bool check_hash, uintptr_t map_base, uintptr_t map_end) {
+        Result LoadAutoLoadModuleSegment(fs::FileHandle file, size_t file_offset, size_t compressed_size, size_t segment_size, bool is_compressed, bool is_zstd, uintptr_t map_base, uintptr_t map_end) {
             /* Select read size based on compression. */
-            if (!is_compressed) {
+            size_t file_size = is_compressed ? compressed_size : segment_size;
-                file_size = segment->size;
-            }
 
             /* Validate size. */
-            R_UNLESS(file_size <= segment->size,                       ldr::ResultInvalidNso());
+            R_UNLESS(file_size <= segment_size,                       ldr::ResultInvalidNso());
-            R_UNLESS(segment->size <= std::numeric_limits<s32>::max(), ldr::ResultInvalidNso());
+            R_UNLESS(file_size    <= std::numeric_limits<s32>::max(), ldr::ResultInvalidNso());
+            R_UNLESS(segment_size <= std::numeric_limits<s32>::max(), ldr::ResultInvalidNso());
 
             /* Load data from file. */
-            uintptr_t load_address = is_compressed ? map_end - file_size : map_base;
+            uintptr_t load_address = is_compressed ? map_end - compressed_size : map_base;
             size_t read_size;
-            R_TRY(fs::ReadFile(std::addressof(read_size), file, segment->file_offset, reinterpret_cast<void *>(load_address), file_size));
+            R_TRY(fs::ReadFile(std::addressof(read_size), file, file_offset, reinterpret_cast<void *>(load_address), file_size));
             R_UNLESS(read_size == file_size, ldr::ResultInvalidNso());
 
             /* Uncompress if necessary. */
-            if (is_compressed) {
+            R_SUCCEED_IF(!is_compressed);
-                bool decompressed = (util::DecompressLZ4(reinterpret_cast<void *>(map_base), segment->size, reinterpret_cast<const void *>(load_address), file_size) == static_cast<int>(segment->size));
+
+            auto compressed_data_buf = reinterpret_cast<const void *>(load_address);
+
+            if (is_zstd) {
+                const size_t map_size = static_cast<size_t>(map_end - map_base);
+
+                bool decompressed = util::DecompressZbicForLoader(g_zstd_dctx_workspace, reinterpret_cast<void *>(map_base), map_size, segment_size, file_size, compressed_data_buf);
+                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));
                 R_UNLESS(decompressed, ldr::ResultInvalidNso());
-            }
-
-            /* Check hash if necessary. */
-            if (check_hash) {
-                u8 hash[crypto::Sha256Generator::HashSize];
-                crypto::GenerateSha256(hash, sizeof(hash), reinterpret_cast<void *>(map_base), segment->size);
-
-                R_UNLESS(std::memcmp(hash, file_hash, sizeof(hash)) == 0, ldr::ResultInvalidNso());
             }
 
             R_SUCCEED();
         }
 
-        Result LoadAutoLoadModule(os::NativeHandle process_handle, fs::FileHandle file, const NsoHeader *nso_header, uintptr_t nso_address, size_t nso_size) {
+        Result CheckSegmentHash(const NsoHeader *nso_header, uintptr_t map_address, NsoHeader::Segment segment) {
+            if ((nso_header->flags & (NsoHeader::Flag_CheckHashText << segment)) == 0) {
+                R_SUCCEED();
+            }
+
+            u8 hash[crypto::Sha256Generator::HashSize];
+            crypto::GenerateSha256(hash, sizeof(hash),
+                reinterpret_cast<void *>(map_address + nso_header->segments[segment].dst_offset),
+                nso_header->segments[segment].size);
+            R_UNLESS(std::memcmp(hash, nso_header->segment_hashes[segment], sizeof(hash)) == 0, ldr::ResultInvalidNso());
+            R_SUCCEED();
+        }
+
+        Result LoadAutoLoadModule(os::NativeHandle process_handle, fs::FileHandle file, const NsoHeader *nso_header, uintptr_t nso_address, size_t nso_size, size_t map_size) {
+            const bool is_zstd = (nso_header->flags & NsoHeader::Flag_UseZbicCompression) != 0;
+
             /* Map and read data from file. */
             {
                 /* Map the process memory. */
                 void *mapped_memory = nullptr;
-                R_TRY(os::MapProcessMemory(std::addressof(mapped_memory), process_handle, nso_address, nso_size, GenerateSecureRandom));
+                R_TRY(os::MapProcessMemory(std::addressof(mapped_memory), process_handle, nso_address, map_size, GenerateSecureRandom));
-                ON_SCOPE_EXIT { os::UnmapProcessMemory(mapped_memory, process_handle, nso_address, nso_size); };
+                ON_SCOPE_EXIT { os::UnmapProcessMemory(mapped_memory, process_handle, nso_address, map_size); };
 
                 const uintptr_t map_address = reinterpret_cast<uintptr_t>(mapped_memory);
+                const uintptr_t map_end     = map_address + map_size;
 
                 /* Load NSO segments. */
-                R_TRY(LoadAutoLoadModuleSegment(file, std::addressof(nso_header->segments[NsoHeader::Segment_Text]), nso_header->text_compressed_size, nso_header->text_hash, (nso_header->flags & NsoHeader::Flag_CompressedText) != 0,
+                R_TRY(LoadAutoLoadModuleSegment(file, nso_header->segments[NsoHeader::Segment_Text].file_offset, nso_header->text_compressed_size, nso_header->text_size,
-                                                      (nso_header->flags & NsoHeader::Flag_CheckHashText) != 0, map_address + nso_header->text_dst_offset, map_address + nso_size));
+                                                      (nso_header->flags & NsoHeader::Flag_CompressedText) != 0, is_zstd, map_address + nso_header->text_dst_offset, map_end));
-                R_TRY(LoadAutoLoadModuleSegment(file, std::addressof(nso_header->segments[NsoHeader::Segment_Ro]), nso_header->ro_compressed_size, nso_header->ro_hash, (nso_header->flags & NsoHeader::Flag_CompressedRo) != 0,
+                R_TRY(LoadAutoLoadModuleSegment(file, nso_header->segments[NsoHeader::Segment_Ro].file_offset, nso_header->ro_compressed_size, nso_header->ro_size,
-                                                      (nso_header->flags & NsoHeader::Flag_CheckHashRo) != 0, map_address + nso_header->ro_dst_offset, map_address + nso_size));
+                                                      (nso_header->flags & NsoHeader::Flag_CompressedRo) != 0, is_zstd, map_address + nso_header->ro_dst_offset, map_end));
-                R_TRY(LoadAutoLoadModuleSegment(file, std::addressof(nso_header->segments[NsoHeader::Segment_Rw]), nso_header->rw_compressed_size, nso_header->rw_hash, (nso_header->flags & NsoHeader::Flag_CompressedRw) != 0,
+                R_TRY(LoadAutoLoadModuleSegment(file, nso_header->segments[NsoHeader::Segment_Rw].file_offset, nso_header->rw_compressed_size, nso_header->rw_size,
-                                                      (nso_header->flags & NsoHeader::Flag_CheckHashRw) != 0, map_address + nso_header->rw_dst_offset, map_address + nso_size));
+                                                      (nso_header->flags & NsoHeader::Flag_CompressedRw) != 0, is_zstd, map_address + nso_header->rw_dst_offset, map_end));
 
                 /* Clear unused space to zero. */
                 const size_t text_end = static_cast<size_t>(nso_header->text_dst_offset) + static_cast<size_t>(nso_header->text_size);
                 const size_t ro_end   = static_cast<size_t>(nso_header->ro_dst_offset)   + static_cast<size_t>(nso_header->ro_size);
                 const size_t rw_end   = static_cast<size_t>(nso_header->rw_dst_offset)   + static_cast<size_t>(nso_header->rw_size);
-                std::memset(reinterpret_cast<void *>(map_address + 0),        0, nso_header->text_dst_offset);
                 std::memset(reinterpret_cast<void *>(map_address + text_end), 0, nso_header->ro_dst_offset - text_end);
                 std::memset(reinterpret_cast<void *>(map_address + ro_end),   0, nso_header->rw_dst_offset - ro_end);
-                std::memset(reinterpret_cast<void *>(map_address + rw_end),   0, nso_header->bss_size);
+                std::memset(reinterpret_cast<void *>(map_address + rw_end),   0, nso_size - rw_end);
+
+                /* Check segment hashes. */
+                R_TRY(CheckSegmentHash(nso_header, map_address, NsoHeader::Segment_Text));
+                R_TRY(CheckSegmentHash(nso_header, map_address, NsoHeader::Segment_Ro));
+                R_TRY(CheckSegmentHash(nso_header, map_address, NsoHeader::Segment_Rw));
 
                 /* Apply embedded patches. */
                 ApplyEmbeddedPatchesToModule(nso_header->module_id, map_address, nso_size);
@@ -691,25 +731,31 @@ namespace ams::ldr {
                 R_TRY(os::SetProcessMemoryPermission(process_handle, nso_address + nso_header->text_dst_offset, text_size, prevent_code_reads ? os::MemoryPermission_ExecuteOnly : os::MemoryPermission_ReadExecute));
             }
             if (ro_size) {
-                R_TRY(os::SetProcessMemoryPermission(process_handle, nso_address + nso_header->ro_dst_offset,   ro_size,   os::MemoryPermission_ReadOnly));
+                R_TRY(os::SetProcessMemoryPermission(process_handle, nso_address + nso_header->ro_dst_offset, ro_size, os::MemoryPermission_ReadOnly));
             }
             if (rw_size) {
-                R_TRY(os::SetProcessMemoryPermission(process_handle, nso_address + nso_header->rw_dst_offset,   rw_size,   os::MemoryPermission_ReadWrite));
+                R_TRY(os::SetProcessMemoryPermission(process_handle, nso_address + nso_header->rw_dst_offset, rw_size, os::MemoryPermission_ReadWrite));
             }
 
             R_SUCCEED();
         }
 
-        Result LoadAutoLoadModules(const ProcessInfo *process_info, const NsoHeader *nso_headers, const AutoLoadModuleInfo *ali, const ArgumentStore::Entry *argument) {
+        Result LoadAutoLoadModules(const ProcessInfo *process_info, const AutoLoadModuleContext &ctx, const ArgumentStore::Entry *argument) {
             /* Load each NSO. */
-            for (size_t i = 0; i < Nso_Count; i++) {
+            const uintptr_t total_end = process_info->code_address + process_info->total_size;
-                if (ali->has_nso[i]) {
-                    fs::FileHandle file;
-                    R_TRY(fs::OpenFile(std::addressof(file), GetNsoPath(i), fs::OpenMode_Read));
-                    ON_SCOPE_EXIT { fs::CloseFile(file); };
 
-                    R_TRY(LoadAutoLoadModule(process_info->process_handle, file, nso_headers + i, process_info->nso_address[i], process_info->nso_size[i]));
+            for (int i = 0; i < ctx.nso_count; i++) {
-                }
+                const NsoIndex nso_idx = static_cast<NsoIndex>(ctx.ali.nso_indices[i]);
+
+                fs::FileHandle file;
+                R_TRY(fs::OpenFile(std::addressof(file), GetNsoPath(nso_idx), fs::OpenMode_Read));
+                ON_SCOPE_EXIT { fs::CloseFile(file); };
+
+                const bool is_zstd    = (ctx.headers[i].flags & NsoHeader::Flag_UseZbicCompression) != 0;
+                const size_t map_size = is_zstd ? (total_end - process_info->nso_address[i]) : process_info->nso_size[i];
+
+                R_TRY(LoadAutoLoadModule(process_info->process_handle, file, ctx.headers + i,
+                      process_info->nso_address[i], process_info->nso_size[i], map_size));
             }
 
             /* Load arguments, if present. */
@@ -735,13 +781,13 @@ namespace ams::ldr {
             R_SUCCEED();
         }
 
-        Result CreateProcessAndLoadAutoLoadModules(ProcessInfo *out, const Meta *meta, const NsoHeader *nso_headers, const AutoLoadModuleInfo *ali, const ArgumentStore::Entry *argument, u32 flags, os::NativeHandle resource_limit) {
+        Result CreateProcessAndLoadAutoLoadModules(ProcessInfo *out, const Meta *meta, const AutoLoadModuleContext &ctx, const ArgumentStore::Entry *argument, u32 flags, os::NativeHandle resource_limit) {
             /* Get CreateProcessParameter. */
             svc::CreateProcessParameter param;
             R_TRY(GetCreateProcessParameter(std::addressof(param), meta, flags, resource_limit));
 
             /* Decide on an NSO layout. */
-            R_TRY(DecideAddressSpaceLayout(out, std::addressof(param), nso_headers, ali, argument));
+            R_TRY(DecideAddressSpaceLayout(out, std::addressof(param), ctx, argument));
 
             /* Actually create process. */
             svc::Handle process_handle;
@@ -749,10 +795,11 @@ namespace ams::ldr {
 
             /* Set the output handle, and ensure that if we fail after this point we clean it up. */
             out->process_handle = process_handle;
+            out->code_address   = param.code_address;
             ON_RESULT_FAILURE { svc::CloseHandle(process_handle); };
 
             /* Load all auto load modules. */
-            R_RETURN(LoadAutoLoadModules(out, nso_headers, ali, argument));
+            R_RETURN(LoadAutoLoadModules(out, ctx, argument));
         }
 
     }
@@ -793,13 +840,13 @@ namespace ams::ldr {
         }
 
         /* Load, validate NSO headers. */
-        AutoLoadModuleInfo auto_load_info = {};
+        AutoLoadModuleContext ctx;
-        R_TRY(LoadAutoLoadHeaders(g_nso_headers, std::addressof(auto_load_info), meta.acid->flags));
+        R_TRY(LoadAutoLoadHeaders(ctx, meta.acid->flags));
-        R_TRY(CheckAutoLoad(g_nso_headers, std::addressof(auto_load_info), meta.acid->flags));
+        R_TRY(CheckAutoLoad(ctx, meta.acid->flags));
 
         /* Actually create the process and load NSOs into process memory. */
         ProcessInfo info;
-        R_TRY(CreateProcessAndLoadAutoLoadModules(std::addressof(info), std::addressof(meta), g_nso_headers, std::addressof(auto_load_info), argument, flags, resource_limit));
+        R_TRY(CreateProcessAndLoadAutoLoadModules(std::addressof(info), std::addressof(meta), ctx, argument, flags, resource_limit));
 
         /* Register NSOs with the RoManager. */
         {
@@ -811,10 +858,8 @@ namespace ams::ldr {
             RoManager::GetInstance().RegisterProcess(pin_id, process_id, meta.aci->program_id, as_type == Npdm::AddressSpaceType_64Bit || as_type == Npdm::AddressSpaceType_64BitDeprecated);
 
             /* Register all NSOs. */
-            for (size_t i = 0; i < Nso_Count; i++) {
+            for (int i = 0; i < ctx.nso_count; i++) {
-                if (auto_load_info.has_nso[i]) {
+                RoManager::GetInstance().AddNso(pin_id, ctx.headers[i].module_id, info.nso_address[i], info.nso_size[i]);
-                    RoManager::GetInstance().AddNso(pin_id, g_nso_headers[i].module_id, info.nso_address[i], info.nso_size[i]);
-                }
             }
         }
 

stratosphere/loader/source/ldr_zstd.cpp

@@ -0,0 +1,18 @@
+/*
+ * 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/>.
+ */
+
+#define AMS_ZSTD_IMPLEMENTATION
+#include <stratosphere/util/util_zbic_for_loader.hpp>