[linux.git] / lib / decompress_unzstd.c

// 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
Commit	Line	Data
4963bb2b NT	1	// SPDX-License-Identifier: GPL-2.0
	2
	3	/*
	4	* Important notes about in-place decompression
	5	*
	6	* At least on x86, the kernel is decompressed in place: the compressed data
	7	* is placed to the end of the output buffer, and the decompressor overwrites
	8	* most of the compressed data. There must be enough safety margin to
	9	* guarantee that the write position is always behind the read position.
	10	*
	11	* The safety margin for ZSTD with a 128 KB block size is calculated below.
	12	* Note that the margin with ZSTD is bigger than with GZIP or XZ!
	13	*
	14	* The worst case for in-place decompression is that the beginning of
	15	* the file is compressed extremely well, and the rest of the file is
	16	* uncompressible. Thus, we must look for worst-case expansion when the
	17	* compressor is encoding uncompressible data.
	18	*
	19	* The structure of the .zst file in case of a compresed kernel is as follows.
	20	* Maximum sizes (as bytes) of the fields are in parenthesis.
	21	*
	22	* Frame Header: (18)
	23	* Blocks: (N)
	24	* Checksum: (4)
	25	*
	26	* The frame header and checksum overhead is at most 22 bytes.
	27	*
	28	* ZSTD stores the data in blocks. Each block has a header whose size is
	29	* a 3 bytes. After the block header, there is up to 128 KB of payload.
	30	* The maximum uncompressed size of the payload is 128 KB. The minimum
	31	* uncompressed size of the payload is never less than the payload size
	32	* (excluding the block header).
	33	*
	34	* The assumption, that the uncompressed size of the payload is never
	35	* smaller than the payload itself, is valid only when talking about
	36	* the payload as a whole. It is possible that the payload has parts where
	37	* the decompressor consumes more input than it produces output. Calculating
	38	* the worst case for this would be tricky. Instead of trying to do that,
	39	* let's simply make sure that the decompressor never overwrites any bytes
	40	* of the payload which it is currently reading.
	41	*
	42	* Now we have enough information to calculate the safety margin. We need
	43	* - 22 bytes for the .zst file format headers;
	44	* - 3 bytes per every 128 KiB of uncompressed size (one block header per
	45	* block); and
	46	* - 128 KiB (biggest possible zstd block size) to make sure that the
	47	* decompressor never overwrites anything from the block it is currently
	48	* reading.
	49	*
	50	* We get the following formula:
	51	*
	52	* safety_margin = 22 + uncompressed_size * 3 / 131072 + 131072
	53	* <= 22 + (uncompressed_size >> 15) + 131072
	54	*/
	55
	56	/*
	57	* Preboot environments #include "path/to/decompress_unzstd.c".
	58	* All of the source files we depend on must be #included.
	59	* zstd's only source dependeny is xxhash, which has no source
	60	* dependencies.
	61	*
	62	* When UNZSTD_PREBOOT is defined we declare __decompress(), which is
	63	* used for kernel decompression, instead of unzstd().
	64	*
65	* Define __DISABLE_EXPORTS in preboot environments to prevent symbols
66	* from xxhash and zstd from being exported by the EXPORT_SYMBOL macro.
67	*/
68	#ifdef STATIC
69	# define UNZSTD_PREBOOT
70	# include "xxhash.c"
71	# include "zstd/entropy_common.c"
72	# include "zstd/fse_decompress.c"
73	# include "zstd/huf_decompress.c"
74	# include "zstd/zstd_common.c"
75	# include "zstd/decompress.c"
76	#endif
77
78	#include <linux/decompress/mm.h>
79	#include <linux/kernel.h>
80	#include <linux/zstd.h>
81
82	/* 128MB is the maximum window size supported by zstd. */
83	#define ZSTD_WINDOWSIZE_MAX (1 << ZSTD_WINDOWLOG_MAX)
84	/*
85	* Size of the input and output buffers in multi-call mode.
86	* Pick a larger size because it isn't used during kernel decompression,
87	* since that is single pass, and we have to allocate a large buffer for
88	* zstd's window anyway. The larger size speeds up initramfs decompression.
89	*/
90	#define ZSTD_IOBUF_SIZE (1 << 17)
91
92	static int INIT handle_zstd_error(size_t ret, void (error)(char x))
93	{
94	const int err = ZSTD_getErrorCode(ret);
95
96	if (!ZSTD_isError(ret))
97	return 0;
98
99	switch (err) {
100	case ZSTD_error_memory_allocation:
101	error("ZSTD decompressor ran out of memory");
102	break;
103	case ZSTD_error_prefix_unknown:
104	error("Input is not in the ZSTD format (wrong magic bytes)");
105	break;
106	case ZSTD_error_dstSize_tooSmall:
107	case ZSTD_error_corruption_detected:
108	case ZSTD_error_checksum_wrong:
109	error("ZSTD-compressed data is corrupt");
110	break;
111	default:
112	error("ZSTD-compressed data is probably corrupt");
113	break;
114	}
115	return -1;
116	}
117
118	/*
119	* Handle the case where we have the entire input and output in one segment.
120	* We can allocate less memory (no circular buffer for the sliding window),
121	* and avoid some memcpy() calls.
122	*/
123	static int INIT decompress_single(const u8 in_buf, long in_len, u8 out_buf,
124	long out_len, long *in_pos,
125	void (error)(char x))
126	{
127	const size_t wksp_size = ZSTD_DCtxWorkspaceBound();
128	void *wksp = large_malloc(wksp_size);
129	ZSTD_DCtx *dctx = ZSTD_initDCtx(wksp, wksp_size);
130	int err;
131	size_t ret;
132
133	if (dctx == NULL) {
134	error("Out of memory while allocating ZSTD_DCtx");
135	err = -1;
136	goto out;
137	}
138	/*
139	* Find out how large the frame actually is, there may be junk at
140	* the end of the frame that ZSTD_decompressDCtx() can't handle.
141	*/
142	ret = ZSTD_findFrameCompressedSize(in_buf, in_len);
143	err = handle_zstd_error(ret, error);
144	if (err)
145	goto out;
146	in_len = (long)ret;
147
148	ret = ZSTD_decompressDCtx(dctx, out_buf, out_len, in_buf, in_len);
149	err = handle_zstd_error(ret, error);
150	if (err)
151	goto out;
152
153	if (in_pos != NULL)
154	*in_pos = in_len;
155
156	err = 0;
157	out:
158	if (wksp != NULL)
159	large_free(wksp);
160	return err;
161	}
162
163	static int INIT __unzstd(unsigned char *in_buf, long in_len,
164	long (fill)(void, unsigned long),
165	long (flush)(void, unsigned long),
166	unsigned char *out_buf, long out_len,
167	long *in_pos,
168	void (error)(char x))
169	{
170	ZSTD_inBuffer in;
171	ZSTD_outBuffer out;
172	ZSTD_frameParams params;
173	void *in_allocated = NULL;
174	void *out_allocated = NULL;
175	void *wksp = NULL;
176	size_t wksp_size;
177	ZSTD_DStream *dstream;
178	int err;
179	size_t ret;
180
1c4dd334 PC	181	/*
	182	* ZSTD decompression code won't be happy if the buffer size is so big
	183	* that its end address overflows. When the size is not provided, make
	184	* it as big as possible without having the end address overflow.
	185	*/
4963bb2b	186	if (out_len == 0)
1c4dd334	187	out_len = UINTPTR_MAX - (uintptr_t)out_buf;
4963bb2b NT	188
	189	if (fill == NULL && flush == NULL)
	190	/*
	191	* We can decompress faster and with less memory when we have a
	192	* single chunk.
	193	*/
	194	return decompress_single(in_buf, in_len, out_buf, out_len,
	195	in_pos, error);
	196
	197	/*
	198	* If in_buf is not provided, we must be using fill(), so allocate
	199	* a large enough buffer. If it is provided, it must be at least
	200	* ZSTD_IOBUF_SIZE large.
	201	*/
	202	if (in_buf == NULL) {
	203	in_allocated = large_malloc(ZSTD_IOBUF_SIZE);
	204	if (in_allocated == NULL) {
	205	error("Out of memory while allocating input buffer");
	206	err = -1;
	207	goto out;
	208	}
	209	in_buf = in_allocated;
	210	in_len = 0;
	211	}
	212	/* Read the first chunk, since we need to decode the frame header. */
	213	if (fill != NULL)
	214	in_len = fill(in_buf, ZSTD_IOBUF_SIZE);
	215	if (in_len < 0) {
	216	error("ZSTD-compressed data is truncated");
	217	err = -1;
	218	goto out;
	219	}
	220	/* Set the first non-empty input buffer. */
	221	in.src = in_buf;
	222	in.pos = 0;
	223	in.size = in_len;
	224	/* Allocate the output buffer if we are using flush(). */
	225	if (flush != NULL) {
	226	out_allocated = large_malloc(ZSTD_IOBUF_SIZE);
	227	if (out_allocated == NULL) {
	228	error("Out of memory while allocating output buffer");
	229	err = -1;
	230	goto out;
	231	}
	232	out_buf = out_allocated;
	233	out_len = ZSTD_IOBUF_SIZE;
	234	}
	235	/* Set the output buffer. */
	236	out.dst = out_buf;
	237	out.pos = 0;
	238	out.size = out_len;
	239
	240	/*
	241	* We need to know the window size to allocate the ZSTD_DStream.
	242	* Since we are streaming, we need to allocate a buffer for the sliding
	243	* window. The window size varies from 1 KB to ZSTD_WINDOWSIZE_MAX
	244	* (8 MB), so it is important to use the actual value so as not to
	245	* waste memory when it is smaller.
	246	*/
	247	ret = ZSTD_getFrameParams(&params, in.src, in.size);
	248	err = handle_zstd_error(ret, error);
	249	if (err)
	250	goto out;
	251	if (ret != 0) {
252	error("ZSTD-compressed data has an incomplete frame header");
253	err = -1;
254	goto out;
255	}
256	if (params.windowSize > ZSTD_WINDOWSIZE_MAX) {
257	error("ZSTD-compressed data has too large a window size");
258	err = -1;
259	goto out;
260	}
261
262	/*
263	* Allocate the ZSTD_DStream now that we know how much memory is
264	* required.
265	*/
266	wksp_size = ZSTD_DStreamWorkspaceBound(params.windowSize);
267	wksp = large_malloc(wksp_size);
268	dstream = ZSTD_initDStream(params.windowSize, wksp, wksp_size);
269	if (dstream == NULL) {
270	error("Out of memory while allocating ZSTD_DStream");
271	err = -1;
272	goto out;
273	}
274
275	/*
276	* Decompression loop:
277	* Read more data if necessary (error if no more data can be read).
278	* Call the decompression function, which returns 0 when finished.
279	* Flush any data produced if using flush().
280	*/
281	if (in_pos != NULL)
282	*in_pos = 0;
283	do {
284	/*
285	* If we need to reload data, either we have fill() and can
286	* try to get more data, or we don't and the input is truncated.
287	*/
288	if (in.pos == in.size) {
289	if (in_pos != NULL)
290	*in_pos += in.pos;
291	in_len = fill ? fill(in_buf, ZSTD_IOBUF_SIZE) : -1;
292	if (in_len < 0) {
293	error("ZSTD-compressed data is truncated");
294	err = -1;
295	goto out;
296	}
297	in.pos = 0;
298	in.size = in_len;
299	}
300	/* Returns zero when the frame is complete. */
301	ret = ZSTD_decompressStream(dstream, &out, &in);
302	err = handle_zstd_error(ret, error);
303	if (err)
304	goto out;
305	/* Flush all of the data produced if using flush(). */
306	if (flush != NULL && out.pos > 0) {
307	if (out.pos != flush(out.dst, out.pos)) {
308	error("Failed to flush()");
309	err = -1;
310	goto out;
311	}
312	out.pos = 0;
313	}
314	} while (ret != 0);
315
316	if (in_pos != NULL)
317	*in_pos += in.pos;
318
319	err = 0;
320	out:
321	if (in_allocated != NULL)
322	large_free(in_allocated);
323	if (out_allocated != NULL)
324	large_free(out_allocated);
325	if (wksp != NULL)
326	large_free(wksp);
327	return err;
328	}
329
330	#ifndef UNZSTD_PREBOOT
331	STATIC int INIT unzstd(unsigned char *buf, long len,
332	long (fill)(void, unsigned long),
333	long (flush)(void, unsigned long),
334	unsigned char *out_buf,
335	long *pos,
336	void (error)(char x))
337	{
338	return __unzstd(buf, len, fill, flush, out_buf, 0, pos, error);
339	}
340	#else
341	STATIC int INIT __decompress(unsigned char *buf, long len,
342	long (fill)(void, unsigned long),
343	long (flush)(void, unsigned long),
344	unsigned char *out_buf, long out_len,
345	long *pos,
346	void (error)(char x))
347	{
348	return __unzstd(buf, len, fill, flush, out_buf, out_len, pos, error);
349	}
350	#endif