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 90static 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 */ 125static 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; 159out: 160 if (wksp != NULL) 161 large_free(wksp); 162 return err; 163} 164 165static 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; 322out: 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 333STATIC 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 343STATIC 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