loader/util: fully implement zstd bic variant

Implement both compression and decompression utilities and simplify loader logic

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

@@ -21,8 +21,11 @@ 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);
 
     /* 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);
 
 }

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

@@ -1,59 +0,0 @@
-/*
- * 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/source/util/util_compression.cpp

@@ -15,6 +15,9 @@
  */
 #include <stratosphere.hpp>
 #include "lz4.h"
+#define ZSTD_STATIC_LINKING_ONLY
+#define ZSTD_ZBIC_SUPPORT 1
+#include "zstd.h"
 
 namespace ams::util {
 
@@ -27,6 +30,23 @@ namespace ams::util {
         /* 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) {
+        /* Basic size checks. */
+        AMS_ABORT_UNLESS(dst_size <= std::numeric_limits<int>::max());
+        AMS_ABORT_UNLESS(src_size <= std::numeric_limits<int>::max());
+        
+        /* Additionally, we must check the compression boundary. */ 
+        auto bound = ZSTD_compressBound(src_size);
+        AMS_ABORT_UNLESS(!ZSTD_isError(bound));
+        AMS_ABORT_UNLESS(dst_size >= bound);
+
+        /* Use Zstd default level. */
+        int compressionLevel = 3;
+        
+        /* This is just a wrapper around Zstd. */
+        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) {
@@ -37,5 +57,54 @@ namespace ams::util {
         /* 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) {
+        /* Basic size checks. */
+        AMS_ABORT_UNLESS(dst_size <= std::numeric_limits<int>::max());
+        AMS_ABORT_UNLESS(src_size <= std::numeric_limits<int>::max());
+        
+        /* Additionally, we must check the decompression boundary. */ 
+        auto bound = ZSTD_decompressBound(src, src_size);
+        AMS_ABORT_UNLESS(!ZSTD_isError(bound));
+        AMS_ABORT_UNLESS(dst_size >= bound);
+        
+        /* This is just a wrapper around Zstd. */
+        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) {
+        /* Check decompression margin. */
+        auto margin = ZSTD_decompressionMargin(src, src_size);
+        if (ZSTD_isError(margin)) {
+            return false;
+        }
+        
+        /* Don't overflow from margin. */
+        if (!util::CanAddWithoutOverflow(margin, expected_dec_size)) {
+            return false;
+        }
+        
+        /* Make sure we fit in the destination buffer. */
+        if (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);
+        
+        /* 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)) {
+            return false;
+        }
+        
+        if (dec_size != expected_dec_size) {
+            return false;
+        }
+        
+        return true;
+    }
 
 }

libraries/libstratosphere/source/util/zstd.h

@@ -134,8 +134,11 @@ ZSTDLIB_API const char* ZSTD_versionString(void);
 #  define ZSTD_CLEVEL_DEFAULT 3
 #endif
 
-#ifndef ZSTD_ZBIC
-#  define ZSTD_ZBIC 0
+/* *************************************
+ *  ZBIC support
+ ***************************************/
+#ifndef ZSTD_ZBIC_SUPPORT
+#  define ZSTD_ZBIC_SUPPORT 0
 #endif
 
 /* *************************************
@@ -143,8 +146,8 @@ ZSTDLIB_API const char* ZSTD_versionString(void);
  ***************************************/
 
 /* All magic numbers are supposed read/written to/from files/memory using little-endian convention */
-#if ZSTD_ZBIC
-#define ZSTD_MAGICNUMBER            0x4349425A    /* 0x4349425A for ZBIC, 0xFD2FB528 for zstd (valid since v0.8.0) */
+#if ZSTD_ZBIC_SUPPORT
+#define ZSTD_MAGICNUMBER            0x4349425A    /* ZBIC magicnumber */
 #else
 #define ZSTD_MAGICNUMBER            0xFD2FB528    /* valid since v0.8.0 */
 #endif

stratosphere/loader/source/ldr_process_creation.cpp

@@ -22,7 +22,6 @@
 #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 {
 
@@ -121,8 +120,9 @@ namespace ams::ldr {
         /* Global NSO header cache. */
         NsoHeader g_nso_headers[Nso_Count];
 
-        /* Global zstd decompression context */
-        alignas(8) u8 g_zstd_dctx_workspace[0x176E8];
+        /* Global Zstd decompression context. */
+        constexpr size_t ZstdDctxWorkspaceSize = 0x176E8;
+        alignas(8) u8 g_zstd_dctx_workspace[ZstdDctxWorkspaceSize];
 
         Result ValidateProgramVersion(ncm::ProgramId program_id, u32 version) {
             /* No version verification is done before 8.1.0. */
@@ -656,9 +656,7 @@ namespace ams::ldr {
             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);
+                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);
                 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));