/* SPDX-License-Identifier: GPL-2.0 */

#ifndef LINUX_DECOMPRESS_UNZSTD_H

#define LINUX_DECOMPRESS_UNZSTD_H

int unzstd(unsigned char *inbuf, long len,

	   long (*fill)(void*, unsigned long),

	   long (*flush)(void*, unsigned long),

	   unsigned char *output,

	   long *pos,

	   void (*error_fn)(char *x));

#endif

config DECOMPRESS_ZSTD

	select ZSTD_DECOMPRESS

	tristate

lib-$(CONFIG_DECOMPRESS_ZSTD) += decompress_unzstd.o

#include <linux/decompress/unzstd.h>

#ifndef CONFIG_DECOMPRESS_ZSTD

# define unzstd NULL

#endif

	{ {0x28, 0xb5}, "zstd", unzstd },

// SPDX-License-Identifier: GPL-2.0

/*

 * Important notes about in-place decompression

 *

 * At least on x86, the kernel is decompressed in place: the compressed data

 * is placed to the end of the output buffer, and the decompressor overwrites

 * most of the compressed data. There must be enough safety margin to

 * guarantee that the write position is always behind the read position.

 *

 * The safety margin for ZSTD with a 128 KB block size is calculated below.

 * Note that the margin with ZSTD is bigger than with GZIP or XZ!

 *

 * The worst case for in-place decompression is that the beginning of

 * the file is compressed extremely well, and the rest of the file is

 * uncompressible. Thus, we must look for worst-case expansion when the

 * compressor is encoding uncompressible data.

 *

 * The structure of the .zst file in case of a compresed kernel is as follows.

 * Maximum sizes (as bytes) of the fields are in parenthesis.

 *

 *    Frame Header: (18)

 *    Blocks: (N)

 *    Checksum: (4)

 *

 * The frame header and checksum overhead is at most 22 bytes.

 *

 * ZSTD stores the data in blocks. Each block has a header whose size is

 * a 3 bytes. After the block header, there is up to 128 KB of payload.

 * The maximum uncompressed size of the payload is 128 KB. The minimum

 * uncompressed size of the payload is never less than the payload size

 * (excluding the block header).

 *

 * The assumption, that the uncompressed size of the payload is never

 * smaller than the payload itself, is valid only when talking about

 * the payload as a whole. It is possible that the payload has parts where

 * the decompressor consumes more input than it produces output. Calculating

 * the worst case for this would be tricky. Instead of trying to do that,

 * let's simply make sure that the decompressor never overwrites any bytes

 * of the payload which it is currently reading.

 *

 * Now we have enough information to calculate the safety margin. We need

 *   - 22 bytes for the .zst file format headers;

 *   - 3 bytes per every 128 KiB of uncompressed size (one block header per

 *     block); and

 *   - 128 KiB (biggest possible zstd block size) to make sure that the

 *     decompressor never overwrites anything from the block it is currently

 *     reading.

 *

 * We get the following formula:

 *

 *    safety_margin = 22 + uncompressed_size * 3 / 131072 + 131072

 *                 <= 22 + (uncompressed_size >> 15) + 131072

 */

/*

 * Preboot environments #include "path/to/decompress_unzstd.c".

 * All of the source files we depend on must be #included.

 * zstd's only source dependeny is xxhash, which has no source

 * dependencies.

 *

 * zstd and xxhash avoid declaring themselves as modules

 * when ZSTD_PREBOOT and XXH_PREBOOT are defined.

 */

#ifdef STATIC

# define ZSTD_PREBOOT

# define XXH_PREBOOT

# include "xxhash.c"

# include "zstd/entropy_common.c"

# include "zstd/fse_decompress.c"

# include "zstd/huf_decompress.c"

# include "zstd/zstd_common.c"

# include "zstd/decompress.c"

#endif

#include <linux/decompress/mm.h>

#include <linux/kernel.h>

#include <linux/zstd.h>

/* 128MB is the maximum window size supported by zstd. */

#define ZSTD_WINDOWSIZE_MAX	(1 << ZSTD_WINDOWLOG_MAX)

/* Size of the input and output buffers in multi-call mode.

 * Pick a larger size because it isn't used during kernel decompression,

 * since that is single pass, and we have to allocate a large buffer for

 * zstd's window anyways. The larger size speeds up initramfs decompression.

 */

#define ZSTD_IOBUF_SIZE		(1 << 17)

static int INIT handle_zstd_error(size_t ret, void (*error)(char *x))

{

	const int err = ZSTD_getErrorCode(ret);

	if (!ZSTD_isError(ret))

		return 0;

	switch (err) {

	case ZSTD_error_memory_allocation:

		error("ZSTD decompressor ran out of memory");

		break;

	case ZSTD_error_prefix_unknown:

		error("Input is not in the ZSTD format (wrong magic bytes)");

		break;

	case ZSTD_error_dstSize_tooSmall:

	case ZSTD_error_corruption_detected:

	case ZSTD_error_checksum_wrong:

		error("ZSTD-compressed data is corrupt");

		break;

	default:

		error("ZSTD-compressed data is probably corrupt");

		break;

	}

	return -1;

}

/*

 * Handle the case where we have the entire input and output in one segment.

 * We can allocate less memory (no circular buffer for the sliding window),

 * and avoid some memcpy() calls.

 */

static int INIT decompress_single(const u8 *in_buf, long in_len, u8 *out_buf,

				  long out_len, long *in_pos,

				  void (*error)(char *x))

{

	const size_t wksp_size = ZSTD_DCtxWorkspaceBound();

	void *wksp = large_malloc(wksp_size);

	ZSTD_DCtx *dctx = ZSTD_initDCtx(wksp, wksp_size);

	int err;

	size_t ret;

	if (dctx == NULL) {

		error("Out of memory while allocating ZSTD_DCtx");

		err = -1;

		goto out;

	}

	/*

	 * Find out how large the frame actually is, there may be junk at

	 * the end of the frame that ZSTD_decompressDCtx() can't handle.

	 */

	ret = ZSTD_findFrameCompressedSize(in_buf, in_len);

	err = handle_zstd_error(ret, error);

	if (err)

		goto out;

	in_len = (long)ret;

	ret = ZSTD_decompressDCtx(dctx, out_buf, out_len, in_buf, in_len);

	err = handle_zstd_error(ret, error);

	if (err)

		goto out;

	if (in_pos != NULL)

		*in_pos = in_len;

	err = 0;

out:

	if (wksp != NULL)

		large_free(wksp);

	return err;

}

static int INIT __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))

{

	ZSTD_inBuffer in;

	ZSTD_outBuffer out;

	ZSTD_frameParams params;

	void *in_allocated = NULL;

	void *out_allocated = NULL;

	void *wksp = NULL;

	size_t wksp_size;

	ZSTD_DStream *dstream;

	int err;

	size_t ret;

	if (out_len == 0)

		out_len = LONG_MAX; /* no limit */

	if (fill == NULL && flush == NULL)

		/*

		 * We can decompress faster and with less memory when we have a

		 * single chunk.

		 */

		return decompress_single(in_buf, in_len, out_buf, out_len,

					 in_pos, error);

	/*

	 * If in_buf is not provided, we must be using fill(), so allocate

	 * a large enough buffer. If it is provided, it must be at least

	 * ZSTD_IOBUF_SIZE large.

	 */

	if (in_buf == NULL) {

		in_allocated = large_malloc(ZSTD_IOBUF_SIZE);

		if (in_allocated == NULL) {

			error("Out of memory while allocating input buffer");

			err = -1;

			goto out;

		}

		in_buf = in_allocated;

		in_len = 0;

	}

	/* Read the first chunk, since we need to decode the frame header. */

	if (fill != NULL)

		in_len = fill(in_buf, ZSTD_IOBUF_SIZE);

	if (in_len < 0) {

		error("ZSTD-compressed data is truncated");

		err = -1;

		goto out;

	}

	/* Set the first non-empty input buffer. */

	in.src = in_buf;

	in.pos = 0;

	in.size = in_len;

	/* Allocate the output buffer if we are using flush(). */

	if (flush != NULL) {

		out_allocated = large_malloc(ZSTD_IOBUF_SIZE);

		if (out_allocated == NULL) {

			error("Out of memory while allocating output buffer");

			err = -1;

			goto out;

		}

		out_buf = out_allocated;

		out_len = ZSTD_IOBUF_SIZE;

	}

	/* Set the output buffer. */

	out.dst = out_buf;

	out.pos = 0;

	out.size = out_len;

	/*

	 * We need to know the window size to allocate the ZSTD_DStream.

	 * Since we are streaming, we need to allocate a buffer for the sliding

	 * window. The window size varies from 1 KB to ZSTD_WINDOWSIZE_MAX

	 * (8 MB), so it is important to use the actual value so as not to

	 * waste memory when it is smaller.

	 */

	ret = ZSTD_getFrameParams(&params, in.src, in.size);

	err = handle_zstd_error(ret, error);

	if (err)

		goto out;

	if (ret != 0) {

		error("ZSTD-compressed data has an incomplete frame header");

		err = -1;

		goto out;

	}

	if (params.windowSize > ZSTD_WINDOWSIZE_MAX) {

		error("ZSTD-compressed data has too large a window size");

		err = -1;

		goto out;

	}

	/*

	 * Allocate the ZSTD_DStream now that we know how much memory is

	 * required.

	 */

	wksp_size = ZSTD_DStreamWorkspaceBound(params.windowSize);

	wksp = large_malloc(wksp_size);

	dstream = ZSTD_initDStream(params.windowSize, wksp, wksp_size);

	if (dstream == NULL) {

		error("Out of memory while allocating ZSTD_DStream");

		err = -1;

		goto out;

	}

	/*

	 * Decompression loop:

	 * Read more data if necessary (error if no more data can be read).

	 * Call the decompression function, which returns 0 when finished.

	 * Flush any data produced if using flush().

	 */

	if (in_pos != NULL)

		*in_pos = 0;

	do {

		/*

		 * If we need to reload data, either we have fill() and can

		 * try to get more data, or we don't and the input is truncated.

		 */

		if (in.pos == in.size) {

			if (in_pos != NULL)

				*in_pos += in.pos;

			in_len = fill ? fill(in_buf, ZSTD_IOBUF_SIZE) : -1;

			if (in_len < 0) {

				error("ZSTD-compressed data is truncated");

				err = -1;

				goto out;

			}

			in.pos = 0;

			in.size = in_len;

		}

		/* Returns zero when the frame is complete. */

		ret = ZSTD_decompressStream(dstream, &out, &in);

		err = handle_zstd_error(ret, error);

		if (err)

			goto out;

		/* Flush all of the data produced if using flush(). */

		if (flush != NULL && out.pos > 0) {

			if (out.pos != flush(out.dst, out.pos)) {

				error("Failed to flush()");

				err = -1;

				goto out;

			}

			out.pos = 0;

		}

	} while (ret != 0);

	if (in_pos != NULL)

		*in_pos += in.pos;

	err = 0;

out:

	if (in_allocated != NULL)

		large_free(in_allocated);

	if (out_allocated != NULL)

		large_free(out_allocated);

	if (wksp != NULL)

		large_free(wksp);

	return err;

}

#ifndef ZSTD_PREBOOT

STATIC int INIT 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))

{

	return __unzstd(buf, len, fill, flush, out_buf, 0, pos, error);

}

#else

STATIC int INIT __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))

{

	return __unzstd(buf, len, fill, flush, out_buf, out_len, pos, error);

}

#endif