lib/decompress_unzstd.c - linux - Git at Google

 // 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.
  *
  * When UNZSTD_PREBOOT is defined we declare __decompress(), which is
  * used for kernel decompression, instead of unzstd().
  *
  * Define __DISABLE_EXPORTS in preboot environments to prevent symbols
  * from xxhash and zstd from being exported by the EXPORT_SYMBOL macro.
  */
 #ifdef STATIC
 # define UNZSTD_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 anyway. 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;

 	/*
 	 * ZSTD decompression code won't be happy if the buffer size is so big
 	 * that its end address overflows. When the size is not provided, make
 	 * it as big as possible without having the end address overflow.
 	 */
 	if (out_len == 0)
 		out_len = UINTPTR_MAX - (uintptr_t)out_buf;

 	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 UNZSTD_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
	// 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.
	*
	* When UNZSTD_PREBOOT is defined we declare __decompress(), which is
	* used for kernel decompression, instead of unzstd().
	*
	* Define __DISABLE_EXPORTS in preboot environments to prevent symbols
	* from xxhash and zstd from being exported by the EXPORT_SYMBOL macro.
	*/
	#ifdef STATIC
	# define UNZSTD_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 anyway. 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;

	/*
	* ZSTD decompression code won't be happy if the buffer size is so big
	* that its end address overflows. When the size is not provided, make
	* it as big as possible without having the end address overflow.
	*/
	if (out_len == 0)
	out_len = UINTPTR_MAX - (uintptr_t)out_buf;

	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 UNZSTD_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