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(code: ret); |
93 | |
94 | if (!zstd_is_error(code: 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(workspace: wksp, workspace_size: 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(src: in_buf, src_size: 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, dst: out_buf, dst_capacity: out_len, src: in_buf, src_size: 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 ; |
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(params: &header, src: in.src, src_size: 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(max_window_size: header.windowSize); |
269 | wksp = large_malloc(wksp_size); |
270 | dstream = zstd_init_dstream(max_window_size: header.windowSize, workspace: wksp, workspace_size: 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, output: &out, input: &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(in_buf: buf, in_len: len, fill, flush, out_buf, out_len: 0, in_pos: 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 | |