loader/util: support regular zstd alongside zbic

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

@@ -19,15 +19,17 @@
 
 namespace ams::util {
     
+    constexpr size_t DecompressZstdWithBicWorkBufferSizeDefault = 0x176E8;
+
     /* Compression utilities. */
     int CompressLZ4(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);
+    size_t CompressZstd(void *dst, size_t dst_size, const void *src, size_t src_size);
-
+    size_t CompressZstdWithBic(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 DecompressZbic(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 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);
+    size_t DecompressZstdWithBic(void *dst, size_t dst_size, const void *src, size_t src_size);
+    bool DecompressZstdWithBic(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/source/util/util_compression_lz4.cpp

@@ -18,7 +18,6 @@
 
 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());
@@ -28,7 +27,6 @@ namespace ams::util {
         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());

libraries/libstratosphere/source/util/util_compression_zstd.cpp

@@ -0,0 +1,60 @@
+/*
+ * 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>
+#define ZSTD_STATIC_LINKING_ONLY
+
+#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 {
+
+    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);
+    }
+    
+    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);
+    }
+
+}

libraries/libstratosphere/source/util/util_compression_zstd_bic.cpp

@@ -14,7 +14,6 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 #include <stratosphere.hpp>
-#include "lz4.h"
 #define ZSTD_STATIC_LINKING_ONLY
 #define ZSTD_ZBIC_SUPPORT 1
 
@@ -27,9 +26,9 @@
 #pragma GCC diagnostic pop
 
 namespace ams::util {
-    static_assert(sizeof(ZSTD_DCtx) <= ZstdDctxWorkspaceSize);
+    static_assert(sizeof(ZSTD_DCtx) <= DecompressZstdWithBicWorkBufferSizeDefault);
 
-    size_t CompressZbic(void *dst, size_t dst_size, const void *src, size_t src_size) {
+    size_t CompressZstdWithBic(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());
@@ -42,11 +41,11 @@ namespace ams::util {
         /* Use Zstd default level. */
         int compressionLevel = 3;
         
-        /* This is just a wrapper around Zstd. */
+        /* This is just a wrapper around Zstd with BIC support enabled. */
         return ZSTD_compress(dst, dst_size, src, src_size, compressionLevel);
     }
     
-    size_t DecompressZbic(void *dst, size_t dst_size, const void *src, size_t src_size) {
+    size_t DecompressZstdWithBic(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());
@@ -56,16 +55,22 @@ namespace ams::util {
         AMS_ABORT_UNLESS(!ZSTD_isError(bound));
         AMS_ABORT_UNLESS(dst_size >= bound);
         
-        /* This is just a wrapper around Zstd. */
+        /* This is just a wrapper around Zstd with BIC support enabled. */
         return ZSTD_decompress(dst, dst_size, src, 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) {
+    bool DecompressZstdWithBic(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);
+            auto error_code = ZSTD_getErrorCode(margin);
-            AMS_LOG("[ldr] can't determine decompression margin: %u (%s)\n", ec, ZSTD_getErrorString(ec));
+            auto error_string = ZSTD_getErrorString(error_code);
+            
+            AMS_LOG("[util] Can't determine decompression margin: %u (%s)\n", error_code, error_string);
+            
+            /* TODO: Unused variables in non-debug build. */
+            AMS_UNUSED(error_code, error_string);
+            
             return false;
         }
         
@@ -76,7 +81,7 @@ 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);
+            AMS_LOG("[util] Not enough space for decompression %lu + %lu > %lu\n", margin, expected_dec_size, dst_size);
             return false;
         }
         
@@ -88,11 +93,18 @@ namespace ams::util {
         size_t dec_size = ZSTD_decompressDCtx(dctx, dst, dst_size, src, src_size);
 
         if (ZSTD_isError(dec_size)) {
-            auto ec = ZSTD_getErrorCode(dec_size);
+            auto error_code = ZSTD_getErrorCode(margin);
-            AMS_LOG("[ldr] decompression failed: %u (%s)\n", ec, ZSTD_getErrorString(ec));
+            auto error_string = ZSTD_getErrorString(error_code);
+            
+            AMS_LOG("[util] Decompression failed: %u (%s)\n", error_code, error_string);
+            
+            /* TODO: Unused variables in non-debug build. */
+            AMS_UNUSED(error_code, error_string);
+            
             return false;
         }
         
+        /* Make sure we match the expected size. */
         if (dec_size != expected_dec_size) {
             return false;
         }

stratosphere/loader/source/ldr_process_creation.cpp

@@ -121,7 +121,7 @@ namespace ams::ldr {
         NsoHeader g_nso_headers[Nso_Count];
 
         /* Global Zstd decompression context. */
-        alignas(8) u8 g_zstd_dctx_workspace[util::ZstdDctxWorkspaceSize];
+        alignas(8) u8 g_zstd_dctx_workspace[util::DecompressZstdWithBicWorkBufferSizeDefault];
 
         Result ValidateProgramVersion(ncm::ProgramId program_id, u32 version) {
             /* No version verification is done before 8.1.0. */
@@ -660,7 +660,7 @@ namespace ams::ldr {
             auto compressed_data_buf = reinterpret_cast<const void *>(load_address);
 
             if (is_zbic) {
-                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);
+                bool decompressed = util::DecompressZstdWithBic(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_LOG(decompressed, ldr::ResultInvalidNso(), "[ldr] Failed to decompress segment with zbic!\n");
             } else {
                 bool decompressed = (util::DecompressLZ4(reinterpret_cast<void *>(map_base), segment_size, compressed_data_buf, file_size) == static_cast<int>(segment_size));