xref: /openbmc/linux/lib/decompress_unzstd.c (revision e9b7b8b3)
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 compressed 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 dependency 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/decompress_sources.h"
72 #else
73 #include <linux/decompress/unzstd.h>
74 #endif
75 
76 #include <linux/decompress/mm.h>
77 #include <linux/kernel.h>
78 #include <linux/zstd.h>
79 
80 /* 128MB is the maximum window size supported by zstd. */
81 #define ZSTD_WINDOWSIZE_MAX	(1 << ZSTD_WINDOWLOG_MAX)
82 /*
83  * Size of the input and output buffers in multi-call mode.
84  * Pick a larger size because it isn't used during kernel decompression,
85  * since that is single pass, and we have to allocate a large buffer for
86  * zstd's window anyway. The larger size speeds up initramfs decompression.
87  */
88 #define ZSTD_IOBUF_SIZE		(1 << 17)
89 
90 static int INIT handle_zstd_error(size_t ret, void (*error)(char *x))
91 {
92 	const zstd_error_code err = zstd_get_error_code(ret);
93 
94 	if (!zstd_is_error(ret))
95 		return 0;
96 
97 	/*
98 	 * zstd_get_error_name() cannot be used because error takes a char *
99 	 * not a const char *
100 	 */
101 	switch (err) {
102 	case ZSTD_error_memory_allocation:
103 		error("ZSTD decompressor ran out of memory");
104 		break;
105 	case ZSTD_error_prefix_unknown:
106 		error("Input is not in the ZSTD format (wrong magic bytes)");
107 		break;
108 	case ZSTD_error_dstSize_tooSmall:
109 	case ZSTD_error_corruption_detected:
110 	case ZSTD_error_checksum_wrong:
111 		error("ZSTD-compressed data is corrupt");
112 		break;
113 	default:
114 		error("ZSTD-compressed data is probably corrupt");
115 		break;
116 	}
117 	return -1;
118 }
119 
120 /*
121  * Handle the case where we have the entire input and output in one segment.
122  * We can allocate less memory (no circular buffer for the sliding window),
123  * and avoid some memcpy() calls.
124  */
125 static int INIT decompress_single(const u8 *in_buf, long in_len, u8 *out_buf,
126 				  long out_len, long *in_pos,
127 				  void (*error)(char *x))
128 {
129 	const size_t wksp_size = zstd_dctx_workspace_bound();
130 	void *wksp = large_malloc(wksp_size);
131 	zstd_dctx *dctx = zstd_init_dctx(wksp, wksp_size);
132 	int err;
133 	size_t ret;
134 
135 	if (dctx == NULL) {
136 		error("Out of memory while allocating zstd_dctx");
137 		err = -1;
138 		goto out;
139 	}
140 	/*
141 	 * Find out how large the frame actually is, there may be junk at
142 	 * the end of the frame that zstd_decompress_dctx() can't handle.
143 	 */
144 	ret = zstd_find_frame_compressed_size(in_buf, in_len);
145 	err = handle_zstd_error(ret, error);
146 	if (err)
147 		goto out;
148 	in_len = (long)ret;
149 
150 	ret = zstd_decompress_dctx(dctx, out_buf, out_len, in_buf, in_len);
151 	err = handle_zstd_error(ret, error);
152 	if (err)
153 		goto out;
154 
155 	if (in_pos != NULL)
156 		*in_pos = in_len;
157 
158 	err = 0;
159 out:
160 	if (wksp != NULL)
161 		large_free(wksp);
162 	return err;
163 }
164 
165 static int INIT __unzstd(unsigned char *in_buf, long in_len,
166 			 long (*fill)(void*, unsigned long),
167 			 long (*flush)(void*, unsigned long),
168 			 unsigned char *out_buf, long out_len,
169 			 long *in_pos,
170 			 void (*error)(char *x))
171 {
172 	zstd_in_buffer in;
173 	zstd_out_buffer out;
174 	zstd_frame_header header;
175 	void *in_allocated = NULL;
176 	void *out_allocated = NULL;
177 	void *wksp = NULL;
178 	size_t wksp_size;
179 	zstd_dstream *dstream;
180 	int err;
181 	size_t ret;
182 
183 	/*
184 	 * ZSTD decompression code won't be happy if the buffer size is so big
185 	 * that its end address overflows. When the size is not provided, make
186 	 * it as big as possible without having the end address overflow.
187 	 */
188 	if (out_len == 0)
189 		out_len = UINTPTR_MAX - (uintptr_t)out_buf;
190 
191 	if (fill == NULL && flush == NULL)
192 		/*
193 		 * We can decompress faster and with less memory when we have a
194 		 * single chunk.
195 		 */
196 		return decompress_single(in_buf, in_len, out_buf, out_len,
197 					 in_pos, error);
198 
199 	/*
200 	 * If in_buf is not provided, we must be using fill(), so allocate
201 	 * a large enough buffer. If it is provided, it must be at least
202 	 * ZSTD_IOBUF_SIZE large.
203 	 */
204 	if (in_buf == NULL) {
205 		in_allocated = large_malloc(ZSTD_IOBUF_SIZE);
206 		if (in_allocated == NULL) {
207 			error("Out of memory while allocating input buffer");
208 			err = -1;
209 			goto out;
210 		}
211 		in_buf = in_allocated;
212 		in_len = 0;
213 	}
214 	/* Read the first chunk, since we need to decode the frame header. */
215 	if (fill != NULL)
216 		in_len = fill(in_buf, ZSTD_IOBUF_SIZE);
217 	if (in_len < 0) {
218 		error("ZSTD-compressed data is truncated");
219 		err = -1;
220 		goto out;
221 	}
222 	/* Set the first non-empty input buffer. */
223 	in.src = in_buf;
224 	in.pos = 0;
225 	in.size = in_len;
226 	/* Allocate the output buffer if we are using flush(). */
227 	if (flush != NULL) {
228 		out_allocated = large_malloc(ZSTD_IOBUF_SIZE);
229 		if (out_allocated == NULL) {
230 			error("Out of memory while allocating output buffer");
231 			err = -1;
232 			goto out;
233 		}
234 		out_buf = out_allocated;
235 		out_len = ZSTD_IOBUF_SIZE;
236 	}
237 	/* Set the output buffer. */
238 	out.dst = out_buf;
239 	out.pos = 0;
240 	out.size = out_len;
241 
242 	/*
243 	 * We need to know the window size to allocate the zstd_dstream.
244 	 * Since we are streaming, we need to allocate a buffer for the sliding
245 	 * window. The window size varies from 1 KB to ZSTD_WINDOWSIZE_MAX
246 	 * (8 MB), so it is important to use the actual value so as not to
247 	 * waste memory when it is smaller.
248 	 */
249 	ret = zstd_get_frame_header(&header, in.src, in.size);
250 	err = handle_zstd_error(ret, error);
251 	if (err)
252 		goto out;
253 	if (ret != 0) {
254 		error("ZSTD-compressed data has an incomplete frame header");
255 		err = -1;
256 		goto out;
257 	}
258 	if (header.windowSize > ZSTD_WINDOWSIZE_MAX) {
259 		error("ZSTD-compressed data has too large a window size");
260 		err = -1;
261 		goto out;
262 	}
263 
264 	/*
265 	 * Allocate the zstd_dstream now that we know how much memory is
266 	 * required.
267 	 */
268 	wksp_size = zstd_dstream_workspace_bound(header.windowSize);
269 	wksp = large_malloc(wksp_size);
270 	dstream = zstd_init_dstream(header.windowSize, wksp, wksp_size);
271 	if (dstream == NULL) {
272 		error("Out of memory while allocating ZSTD_DStream");
273 		err = -1;
274 		goto out;
275 	}
276 
277 	/*
278 	 * Decompression loop:
279 	 * Read more data if necessary (error if no more data can be read).
280 	 * Call the decompression function, which returns 0 when finished.
281 	 * Flush any data produced if using flush().
282 	 */
283 	if (in_pos != NULL)
284 		*in_pos = 0;
285 	do {
286 		/*
287 		 * If we need to reload data, either we have fill() and can
288 		 * try to get more data, or we don't and the input is truncated.
289 		 */
290 		if (in.pos == in.size) {
291 			if (in_pos != NULL)
292 				*in_pos += in.pos;
293 			in_len = fill ? fill(in_buf, ZSTD_IOBUF_SIZE) : -1;
294 			if (in_len < 0) {
295 				error("ZSTD-compressed data is truncated");
296 				err = -1;
297 				goto out;
298 			}
299 			in.pos = 0;
300 			in.size = in_len;
301 		}
302 		/* Returns zero when the frame is complete. */
303 		ret = zstd_decompress_stream(dstream, &out, &in);
304 		err = handle_zstd_error(ret, error);
305 		if (err)
306 			goto out;
307 		/* Flush all of the data produced if using flush(). */
308 		if (flush != NULL && out.pos > 0) {
309 			if (out.pos != flush(out.dst, out.pos)) {
310 				error("Failed to flush()");
311 				err = -1;
312 				goto out;
313 			}
314 			out.pos = 0;
315 		}
316 	} while (ret != 0);
317 
318 	if (in_pos != NULL)
319 		*in_pos += in.pos;
320 
321 	err = 0;
322 out:
323 	if (in_allocated != NULL)
324 		large_free(in_allocated);
325 	if (out_allocated != NULL)
326 		large_free(out_allocated);
327 	if (wksp != NULL)
328 		large_free(wksp);
329 	return err;
330 }
331 
332 #ifndef UNZSTD_PREBOOT
333 STATIC int INIT unzstd(unsigned char *buf, long len,
334 		       long (*fill)(void*, unsigned long),
335 		       long (*flush)(void*, unsigned long),
336 		       unsigned char *out_buf,
337 		       long *pos,
338 		       void (*error)(char *x))
339 {
340 	return __unzstd(buf, len, fill, flush, out_buf, 0, pos, error);
341 }
342 #else
343 STATIC int INIT __decompress(unsigned char *buf, long len,
344 			     long (*fill)(void*, unsigned long),
345 			     long (*flush)(void*, unsigned long),
346 			     unsigned char *out_buf, long out_len,
347 			     long *pos,
348 			     void (*error)(char *x))
349 {
350 	return __unzstd(buf, len, fill, flush, out_buf, out_len, pos, error);
351 }
352 #endif
353