14963bb2bSNick Terrell // SPDX-License-Identifier: GPL-2.0
24963bb2bSNick Terrell
34963bb2bSNick Terrell /*
44963bb2bSNick Terrell * Important notes about in-place decompression
54963bb2bSNick Terrell *
64963bb2bSNick Terrell * At least on x86, the kernel is decompressed in place: the compressed data
74963bb2bSNick Terrell * is placed to the end of the output buffer, and the decompressor overwrites
84963bb2bSNick Terrell * most of the compressed data. There must be enough safety margin to
94963bb2bSNick Terrell * guarantee that the write position is always behind the read position.
104963bb2bSNick Terrell *
114963bb2bSNick Terrell * The safety margin for ZSTD with a 128 KB block size is calculated below.
124963bb2bSNick Terrell * Note that the margin with ZSTD is bigger than with GZIP or XZ!
134963bb2bSNick Terrell *
144963bb2bSNick Terrell * The worst case for in-place decompression is that the beginning of
154963bb2bSNick Terrell * the file is compressed extremely well, and the rest of the file is
164963bb2bSNick Terrell * uncompressible. Thus, we must look for worst-case expansion when the
174963bb2bSNick Terrell * compressor is encoding uncompressible data.
184963bb2bSNick Terrell *
1905911c5dSZhen Lei * The structure of the .zst file in case of a compressed kernel is as follows.
204963bb2bSNick Terrell * Maximum sizes (as bytes) of the fields are in parenthesis.
214963bb2bSNick Terrell *
224963bb2bSNick Terrell * Frame Header: (18)
234963bb2bSNick Terrell * Blocks: (N)
244963bb2bSNick Terrell * Checksum: (4)
254963bb2bSNick Terrell *
264963bb2bSNick Terrell * The frame header and checksum overhead is at most 22 bytes.
274963bb2bSNick Terrell *
284963bb2bSNick Terrell * ZSTD stores the data in blocks. Each block has a header whose size is
294963bb2bSNick Terrell * a 3 bytes. After the block header, there is up to 128 KB of payload.
304963bb2bSNick Terrell * The maximum uncompressed size of the payload is 128 KB. The minimum
314963bb2bSNick Terrell * uncompressed size of the payload is never less than the payload size
324963bb2bSNick Terrell * (excluding the block header).
334963bb2bSNick Terrell *
344963bb2bSNick Terrell * The assumption, that the uncompressed size of the payload is never
354963bb2bSNick Terrell * smaller than the payload itself, is valid only when talking about
364963bb2bSNick Terrell * the payload as a whole. It is possible that the payload has parts where
374963bb2bSNick Terrell * the decompressor consumes more input than it produces output. Calculating
384963bb2bSNick Terrell * the worst case for this would be tricky. Instead of trying to do that,
394963bb2bSNick Terrell * let's simply make sure that the decompressor never overwrites any bytes
404963bb2bSNick Terrell * of the payload which it is currently reading.
414963bb2bSNick Terrell *
424963bb2bSNick Terrell * Now we have enough information to calculate the safety margin. We need
434963bb2bSNick Terrell * - 22 bytes for the .zst file format headers;
444963bb2bSNick Terrell * - 3 bytes per every 128 KiB of uncompressed size (one block header per
454963bb2bSNick Terrell * block); and
464963bb2bSNick Terrell * - 128 KiB (biggest possible zstd block size) to make sure that the
474963bb2bSNick Terrell * decompressor never overwrites anything from the block it is currently
484963bb2bSNick Terrell * reading.
494963bb2bSNick Terrell *
504963bb2bSNick Terrell * We get the following formula:
514963bb2bSNick Terrell *
524963bb2bSNick Terrell * safety_margin = 22 + uncompressed_size * 3 / 131072 + 131072
534963bb2bSNick Terrell * <= 22 + (uncompressed_size >> 15) + 131072
544963bb2bSNick Terrell */
554963bb2bSNick Terrell
564963bb2bSNick Terrell /*
574963bb2bSNick Terrell * Preboot environments #include "path/to/decompress_unzstd.c".
584963bb2bSNick Terrell * All of the source files we depend on must be #included.
5905911c5dSZhen Lei * zstd's only source dependency is xxhash, which has no source
604963bb2bSNick Terrell * dependencies.
614963bb2bSNick Terrell *
624963bb2bSNick Terrell * When UNZSTD_PREBOOT is defined we declare __decompress(), which is
634963bb2bSNick Terrell * used for kernel decompression, instead of unzstd().
644963bb2bSNick Terrell *
654963bb2bSNick Terrell * Define __DISABLE_EXPORTS in preboot environments to prevent symbols
664963bb2bSNick Terrell * from xxhash and zstd from being exported by the EXPORT_SYMBOL macro.
674963bb2bSNick Terrell */
684963bb2bSNick Terrell #ifdef STATIC
694963bb2bSNick Terrell # define UNZSTD_PREBOOT
704963bb2bSNick Terrell # include "xxhash.c"
712479b523SNick Terrell # include "zstd/decompress_sources.h"
72*00444448SArnd Bergmann #else
73*00444448SArnd Bergmann #include <linux/decompress/unzstd.h>
744963bb2bSNick Terrell #endif
754963bb2bSNick Terrell
764963bb2bSNick Terrell #include <linux/decompress/mm.h>
774963bb2bSNick Terrell #include <linux/kernel.h>
784963bb2bSNick Terrell #include <linux/zstd.h>
794963bb2bSNick Terrell
804963bb2bSNick Terrell /* 128MB is the maximum window size supported by zstd. */
814963bb2bSNick Terrell #define ZSTD_WINDOWSIZE_MAX (1 << ZSTD_WINDOWLOG_MAX)
824963bb2bSNick Terrell /*
834963bb2bSNick Terrell * Size of the input and output buffers in multi-call mode.
844963bb2bSNick Terrell * Pick a larger size because it isn't used during kernel decompression,
854963bb2bSNick Terrell * since that is single pass, and we have to allocate a large buffer for
864963bb2bSNick Terrell * zstd's window anyway. The larger size speeds up initramfs decompression.
874963bb2bSNick Terrell */
884963bb2bSNick Terrell #define ZSTD_IOBUF_SIZE (1 << 17)
894963bb2bSNick Terrell
handle_zstd_error(size_t ret,void (* error)(char * x))904963bb2bSNick Terrell static int INIT handle_zstd_error(size_t ret, void (*error)(char *x))
914963bb2bSNick Terrell {
92cf30f6a5SNick Terrell const zstd_error_code err = zstd_get_error_code(ret);
934963bb2bSNick Terrell
94cf30f6a5SNick Terrell if (!zstd_is_error(ret))
954963bb2bSNick Terrell return 0;
964963bb2bSNick Terrell
97cf30f6a5SNick Terrell /*
98cf30f6a5SNick Terrell * zstd_get_error_name() cannot be used because error takes a char *
99cf30f6a5SNick Terrell * not a const char *
100cf30f6a5SNick Terrell */
1014963bb2bSNick Terrell switch (err) {
1024963bb2bSNick Terrell case ZSTD_error_memory_allocation:
1034963bb2bSNick Terrell error("ZSTD decompressor ran out of memory");
1044963bb2bSNick Terrell break;
1054963bb2bSNick Terrell case ZSTD_error_prefix_unknown:
1064963bb2bSNick Terrell error("Input is not in the ZSTD format (wrong magic bytes)");
1074963bb2bSNick Terrell break;
1084963bb2bSNick Terrell case ZSTD_error_dstSize_tooSmall:
1094963bb2bSNick Terrell case ZSTD_error_corruption_detected:
1104963bb2bSNick Terrell case ZSTD_error_checksum_wrong:
1114963bb2bSNick Terrell error("ZSTD-compressed data is corrupt");
1124963bb2bSNick Terrell break;
1134963bb2bSNick Terrell default:
1144963bb2bSNick Terrell error("ZSTD-compressed data is probably corrupt");
1154963bb2bSNick Terrell break;
1164963bb2bSNick Terrell }
1174963bb2bSNick Terrell return -1;
1184963bb2bSNick Terrell }
1194963bb2bSNick Terrell
1204963bb2bSNick Terrell /*
1214963bb2bSNick Terrell * Handle the case where we have the entire input and output in one segment.
1224963bb2bSNick Terrell * We can allocate less memory (no circular buffer for the sliding window),
1234963bb2bSNick Terrell * and avoid some memcpy() calls.
1244963bb2bSNick Terrell */
decompress_single(const u8 * in_buf,long in_len,u8 * out_buf,long out_len,long * in_pos,void (* error)(char * x))1254963bb2bSNick Terrell static int INIT decompress_single(const u8 *in_buf, long in_len, u8 *out_buf,
1264963bb2bSNick Terrell long out_len, long *in_pos,
1274963bb2bSNick Terrell void (*error)(char *x))
1284963bb2bSNick Terrell {
129cf30f6a5SNick Terrell const size_t wksp_size = zstd_dctx_workspace_bound();
1304963bb2bSNick Terrell void *wksp = large_malloc(wksp_size);
131cf30f6a5SNick Terrell zstd_dctx *dctx = zstd_init_dctx(wksp, wksp_size);
1324963bb2bSNick Terrell int err;
1334963bb2bSNick Terrell size_t ret;
1344963bb2bSNick Terrell
1354963bb2bSNick Terrell if (dctx == NULL) {
136cf30f6a5SNick Terrell error("Out of memory while allocating zstd_dctx");
1374963bb2bSNick Terrell err = -1;
1384963bb2bSNick Terrell goto out;
1394963bb2bSNick Terrell }
1404963bb2bSNick Terrell /*
1414963bb2bSNick Terrell * Find out how large the frame actually is, there may be junk at
142cf30f6a5SNick Terrell * the end of the frame that zstd_decompress_dctx() can't handle.
1434963bb2bSNick Terrell */
144cf30f6a5SNick Terrell ret = zstd_find_frame_compressed_size(in_buf, in_len);
1454963bb2bSNick Terrell err = handle_zstd_error(ret, error);
1464963bb2bSNick Terrell if (err)
1474963bb2bSNick Terrell goto out;
1484963bb2bSNick Terrell in_len = (long)ret;
1494963bb2bSNick Terrell
150cf30f6a5SNick Terrell ret = zstd_decompress_dctx(dctx, out_buf, out_len, in_buf, in_len);
1514963bb2bSNick Terrell err = handle_zstd_error(ret, error);
1524963bb2bSNick Terrell if (err)
1534963bb2bSNick Terrell goto out;
1544963bb2bSNick Terrell
1554963bb2bSNick Terrell if (in_pos != NULL)
1564963bb2bSNick Terrell *in_pos = in_len;
1574963bb2bSNick Terrell
1584963bb2bSNick Terrell err = 0;
1594963bb2bSNick Terrell out:
1604963bb2bSNick Terrell if (wksp != NULL)
1614963bb2bSNick Terrell large_free(wksp);
1624963bb2bSNick Terrell return err;
1634963bb2bSNick Terrell }
1644963bb2bSNick Terrell
__unzstd(unsigned char * in_buf,long in_len,long (* fill)(void *,unsigned long),long (* flush)(void *,unsigned long),unsigned char * out_buf,long out_len,long * in_pos,void (* error)(char * x))1654963bb2bSNick Terrell static int INIT __unzstd(unsigned char *in_buf, long in_len,
1664963bb2bSNick Terrell long (*fill)(void*, unsigned long),
1674963bb2bSNick Terrell long (*flush)(void*, unsigned long),
1684963bb2bSNick Terrell unsigned char *out_buf, long out_len,
1694963bb2bSNick Terrell long *in_pos,
1704963bb2bSNick Terrell void (*error)(char *x))
1714963bb2bSNick Terrell {
172cf30f6a5SNick Terrell zstd_in_buffer in;
173cf30f6a5SNick Terrell zstd_out_buffer out;
174cf30f6a5SNick Terrell zstd_frame_header header;
1754963bb2bSNick Terrell void *in_allocated = NULL;
1764963bb2bSNick Terrell void *out_allocated = NULL;
1774963bb2bSNick Terrell void *wksp = NULL;
1784963bb2bSNick Terrell size_t wksp_size;
179cf30f6a5SNick Terrell zstd_dstream *dstream;
1804963bb2bSNick Terrell int err;
1814963bb2bSNick Terrell size_t ret;
1824963bb2bSNick Terrell
1831c4dd334SPaul Cercueil /*
1841c4dd334SPaul Cercueil * ZSTD decompression code won't be happy if the buffer size is so big
1851c4dd334SPaul Cercueil * that its end address overflows. When the size is not provided, make
1861c4dd334SPaul Cercueil * it as big as possible without having the end address overflow.
1871c4dd334SPaul Cercueil */
1884963bb2bSNick Terrell if (out_len == 0)
1891c4dd334SPaul Cercueil out_len = UINTPTR_MAX - (uintptr_t)out_buf;
1904963bb2bSNick Terrell
1914963bb2bSNick Terrell if (fill == NULL && flush == NULL)
1924963bb2bSNick Terrell /*
1934963bb2bSNick Terrell * We can decompress faster and with less memory when we have a
1944963bb2bSNick Terrell * single chunk.
1954963bb2bSNick Terrell */
1964963bb2bSNick Terrell return decompress_single(in_buf, in_len, out_buf, out_len,
1974963bb2bSNick Terrell in_pos, error);
1984963bb2bSNick Terrell
1994963bb2bSNick Terrell /*
2004963bb2bSNick Terrell * If in_buf is not provided, we must be using fill(), so allocate
2014963bb2bSNick Terrell * a large enough buffer. If it is provided, it must be at least
2024963bb2bSNick Terrell * ZSTD_IOBUF_SIZE large.
2034963bb2bSNick Terrell */
2044963bb2bSNick Terrell if (in_buf == NULL) {
2054963bb2bSNick Terrell in_allocated = large_malloc(ZSTD_IOBUF_SIZE);
2064963bb2bSNick Terrell if (in_allocated == NULL) {
2074963bb2bSNick Terrell error("Out of memory while allocating input buffer");
2084963bb2bSNick Terrell err = -1;
2094963bb2bSNick Terrell goto out;
2104963bb2bSNick Terrell }
2114963bb2bSNick Terrell in_buf = in_allocated;
2124963bb2bSNick Terrell in_len = 0;
2134963bb2bSNick Terrell }
2144963bb2bSNick Terrell /* Read the first chunk, since we need to decode the frame header. */
2154963bb2bSNick Terrell if (fill != NULL)
2164963bb2bSNick Terrell in_len = fill(in_buf, ZSTD_IOBUF_SIZE);
2174963bb2bSNick Terrell if (in_len < 0) {
2184963bb2bSNick Terrell error("ZSTD-compressed data is truncated");
2194963bb2bSNick Terrell err = -1;
2204963bb2bSNick Terrell goto out;
2214963bb2bSNick Terrell }
2224963bb2bSNick Terrell /* Set the first non-empty input buffer. */
2234963bb2bSNick Terrell in.src = in_buf;
2244963bb2bSNick Terrell in.pos = 0;
2254963bb2bSNick Terrell in.size = in_len;
2264963bb2bSNick Terrell /* Allocate the output buffer if we are using flush(). */
2274963bb2bSNick Terrell if (flush != NULL) {
2284963bb2bSNick Terrell out_allocated = large_malloc(ZSTD_IOBUF_SIZE);
2294963bb2bSNick Terrell if (out_allocated == NULL) {
2304963bb2bSNick Terrell error("Out of memory while allocating output buffer");
2314963bb2bSNick Terrell err = -1;
2324963bb2bSNick Terrell goto out;
2334963bb2bSNick Terrell }
2344963bb2bSNick Terrell out_buf = out_allocated;
2354963bb2bSNick Terrell out_len = ZSTD_IOBUF_SIZE;
2364963bb2bSNick Terrell }
2374963bb2bSNick Terrell /* Set the output buffer. */
2384963bb2bSNick Terrell out.dst = out_buf;
2394963bb2bSNick Terrell out.pos = 0;
2404963bb2bSNick Terrell out.size = out_len;
2414963bb2bSNick Terrell
2424963bb2bSNick Terrell /*
243cf30f6a5SNick Terrell * We need to know the window size to allocate the zstd_dstream.
2444963bb2bSNick Terrell * Since we are streaming, we need to allocate a buffer for the sliding
2454963bb2bSNick Terrell * window. The window size varies from 1 KB to ZSTD_WINDOWSIZE_MAX
2464963bb2bSNick Terrell * (8 MB), so it is important to use the actual value so as not to
2474963bb2bSNick Terrell * waste memory when it is smaller.
2484963bb2bSNick Terrell */
249cf30f6a5SNick Terrell ret = zstd_get_frame_header(&header, in.src, in.size);
2504963bb2bSNick Terrell err = handle_zstd_error(ret, error);
2514963bb2bSNick Terrell if (err)
2524963bb2bSNick Terrell goto out;
2534963bb2bSNick Terrell if (ret != 0) {
2544963bb2bSNick Terrell error("ZSTD-compressed data has an incomplete frame header");
2554963bb2bSNick Terrell err = -1;
2564963bb2bSNick Terrell goto out;
2574963bb2bSNick Terrell }
258cf30f6a5SNick Terrell if (header.windowSize > ZSTD_WINDOWSIZE_MAX) {
2594963bb2bSNick Terrell error("ZSTD-compressed data has too large a window size");
2604963bb2bSNick Terrell err = -1;
2614963bb2bSNick Terrell goto out;
2624963bb2bSNick Terrell }
2634963bb2bSNick Terrell
2644963bb2bSNick Terrell /*
265cf30f6a5SNick Terrell * Allocate the zstd_dstream now that we know how much memory is
2664963bb2bSNick Terrell * required.
2674963bb2bSNick Terrell */
268cf30f6a5SNick Terrell wksp_size = zstd_dstream_workspace_bound(header.windowSize);
2694963bb2bSNick Terrell wksp = large_malloc(wksp_size);
270cf30f6a5SNick Terrell dstream = zstd_init_dstream(header.windowSize, wksp, wksp_size);
2714963bb2bSNick Terrell if (dstream == NULL) {
2724963bb2bSNick Terrell error("Out of memory while allocating ZSTD_DStream");
2734963bb2bSNick Terrell err = -1;
2744963bb2bSNick Terrell goto out;
2754963bb2bSNick Terrell }
2764963bb2bSNick Terrell
2774963bb2bSNick Terrell /*
2784963bb2bSNick Terrell * Decompression loop:
2794963bb2bSNick Terrell * Read more data if necessary (error if no more data can be read).
2804963bb2bSNick Terrell * Call the decompression function, which returns 0 when finished.
2814963bb2bSNick Terrell * Flush any data produced if using flush().
2824963bb2bSNick Terrell */
2834963bb2bSNick Terrell if (in_pos != NULL)
2844963bb2bSNick Terrell *in_pos = 0;
2854963bb2bSNick Terrell do {
2864963bb2bSNick Terrell /*
2874963bb2bSNick Terrell * If we need to reload data, either we have fill() and can
2884963bb2bSNick Terrell * try to get more data, or we don't and the input is truncated.
2894963bb2bSNick Terrell */
2904963bb2bSNick Terrell if (in.pos == in.size) {
2914963bb2bSNick Terrell if (in_pos != NULL)
2924963bb2bSNick Terrell *in_pos += in.pos;
2934963bb2bSNick Terrell in_len = fill ? fill(in_buf, ZSTD_IOBUF_SIZE) : -1;
2944963bb2bSNick Terrell if (in_len < 0) {
2954963bb2bSNick Terrell error("ZSTD-compressed data is truncated");
2964963bb2bSNick Terrell err = -1;
2974963bb2bSNick Terrell goto out;
2984963bb2bSNick Terrell }
2994963bb2bSNick Terrell in.pos = 0;
3004963bb2bSNick Terrell in.size = in_len;
3014963bb2bSNick Terrell }
3024963bb2bSNick Terrell /* Returns zero when the frame is complete. */
303cf30f6a5SNick Terrell ret = zstd_decompress_stream(dstream, &out, &in);
3044963bb2bSNick Terrell err = handle_zstd_error(ret, error);
3054963bb2bSNick Terrell if (err)
3064963bb2bSNick Terrell goto out;
3074963bb2bSNick Terrell /* Flush all of the data produced if using flush(). */
3084963bb2bSNick Terrell if (flush != NULL && out.pos > 0) {
3094963bb2bSNick Terrell if (out.pos != flush(out.dst, out.pos)) {
3104963bb2bSNick Terrell error("Failed to flush()");
3114963bb2bSNick Terrell err = -1;
3124963bb2bSNick Terrell goto out;
3134963bb2bSNick Terrell }
3144963bb2bSNick Terrell out.pos = 0;
3154963bb2bSNick Terrell }
3164963bb2bSNick Terrell } while (ret != 0);
3174963bb2bSNick Terrell
3184963bb2bSNick Terrell if (in_pos != NULL)
3194963bb2bSNick Terrell *in_pos += in.pos;
3204963bb2bSNick Terrell
3214963bb2bSNick Terrell err = 0;
3224963bb2bSNick Terrell out:
3234963bb2bSNick Terrell if (in_allocated != NULL)
3244963bb2bSNick Terrell large_free(in_allocated);
3254963bb2bSNick Terrell if (out_allocated != NULL)
3264963bb2bSNick Terrell large_free(out_allocated);
3274963bb2bSNick Terrell if (wksp != NULL)
3284963bb2bSNick Terrell large_free(wksp);
3294963bb2bSNick Terrell return err;
3304963bb2bSNick Terrell }
3314963bb2bSNick Terrell
3324963bb2bSNick Terrell #ifndef UNZSTD_PREBOOT
unzstd(unsigned char * buf,long len,long (* fill)(void *,unsigned long),long (* flush)(void *,unsigned long),unsigned char * out_buf,long * pos,void (* error)(char * x))3334963bb2bSNick Terrell STATIC int INIT unzstd(unsigned char *buf, long len,
3344963bb2bSNick Terrell long (*fill)(void*, unsigned long),
3354963bb2bSNick Terrell long (*flush)(void*, unsigned long),
3364963bb2bSNick Terrell unsigned char *out_buf,
3374963bb2bSNick Terrell long *pos,
3384963bb2bSNick Terrell void (*error)(char *x))
3394963bb2bSNick Terrell {
3404963bb2bSNick Terrell return __unzstd(buf, len, fill, flush, out_buf, 0, pos, error);
3414963bb2bSNick Terrell }
3424963bb2bSNick Terrell #else
__decompress(unsigned char * buf,long len,long (* fill)(void *,unsigned long),long (* flush)(void *,unsigned long),unsigned char * out_buf,long out_len,long * pos,void (* error)(char * x))3434963bb2bSNick Terrell STATIC int INIT __decompress(unsigned char *buf, long len,
3444963bb2bSNick Terrell long (*fill)(void*, unsigned long),
3454963bb2bSNick Terrell long (*flush)(void*, unsigned long),
3464963bb2bSNick Terrell unsigned char *out_buf, long out_len,
3474963bb2bSNick Terrell long *pos,
3484963bb2bSNick Terrell void (*error)(char *x))
3494963bb2bSNick Terrell {
3504963bb2bSNick Terrell return __unzstd(buf, len, fill, flush, out_buf, out_len, pos, error);
3514963bb2bSNick Terrell }
3524963bb2bSNick Terrell #endif
353