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 181 if (out_len == 0) 182 out_len = LONG_MAX; /* no limit */ 183 184 if (fill == NULL && flush == NULL) 185 /* 186 * We can decompress faster and with less memory when we have a 187 * single chunk. 188 */ 189 return decompress_single(in_buf, in_len, out_buf, out_len, 190 in_pos, error); 191 192 /* 193 * If in_buf is not provided, we must be using fill(), so allocate 194 * a large enough buffer. If it is provided, it must be at least 195 * ZSTD_IOBUF_SIZE large. 196 */ 197 if (in_buf == NULL) { 198 in_allocated = large_malloc(ZSTD_IOBUF_SIZE); 199 if (in_allocated == NULL) { 200 error("Out of memory while allocating input buffer"); 201 err = -1; 202 goto out; 203 } 204 in_buf = in_allocated; 205 in_len = 0; 206 } 207 /* Read the first chunk, since we need to decode the frame header. */ 208 if (fill != NULL) 209 in_len = fill(in_buf, ZSTD_IOBUF_SIZE); 210 if (in_len < 0) { 211 error("ZSTD-compressed data is truncated"); 212 err = -1; 213 goto out; 214 } 215 /* Set the first non-empty input buffer. */ 216 in.src = in_buf; 217 in.pos = 0; 218 in.size = in_len; 219 /* Allocate the output buffer if we are using flush(). */ 220 if (flush != NULL) { 221 out_allocated = large_malloc(ZSTD_IOBUF_SIZE); 222 if (out_allocated == NULL) { 223 error("Out of memory while allocating output buffer"); 224 err = -1; 225 goto out; 226 } 227 out_buf = out_allocated; 228 out_len = ZSTD_IOBUF_SIZE; 229 } 230 /* Set the output buffer. */ 231 out.dst = out_buf; 232 out.pos = 0; 233 out.size = out_len; 234 235 /* 236 * We need to know the window size to allocate the ZSTD_DStream. 237 * Since we are streaming, we need to allocate a buffer for the sliding 238 * window. The window size varies from 1 KB to ZSTD_WINDOWSIZE_MAX 239 * (8 MB), so it is important to use the actual value so as not to 240 * waste memory when it is smaller. 241 */ 242 ret = ZSTD_getFrameParams(¶ms, in.src, in.size); 243 err = handle_zstd_error(ret, error); 244 if (err) 245 goto out; 246 if (ret != 0) { 247 error("ZSTD-compressed data has an incomplete frame header"); 248 err = -1; 249 goto out; 250 } 251 if (params.windowSize > ZSTD_WINDOWSIZE_MAX) { 252 error("ZSTD-compressed data has too large a window size"); 253 err = -1; 254 goto out; 255 } 256 257 /* 258 * Allocate the ZSTD_DStream now that we know how much memory is 259 * required. 260 */ 261 wksp_size = ZSTD_DStreamWorkspaceBound(params.windowSize); 262 wksp = large_malloc(wksp_size); 263 dstream = ZSTD_initDStream(params.windowSize, wksp, wksp_size); 264 if (dstream == NULL) { 265 error("Out of memory while allocating ZSTD_DStream"); 266 err = -1; 267 goto out; 268 } 269 270 /* 271 * Decompression loop: 272 * Read more data if necessary (error if no more data can be read). 273 * Call the decompression function, which returns 0 when finished. 274 * Flush any data produced if using flush(). 275 */ 276 if (in_pos != NULL) 277 *in_pos = 0; 278 do { 279 /* 280 * If we need to reload data, either we have fill() and can 281 * try to get more data, or we don't and the input is truncated. 282 */ 283 if (in.pos == in.size) { 284 if (in_pos != NULL) 285 *in_pos += in.pos; 286 in_len = fill ? fill(in_buf, ZSTD_IOBUF_SIZE) : -1; 287 if (in_len < 0) { 288 error("ZSTD-compressed data is truncated"); 289 err = -1; 290 goto out; 291 } 292 in.pos = 0; 293 in.size = in_len; 294 } 295 /* Returns zero when the frame is complete. */ 296 ret = ZSTD_decompressStream(dstream, &out, &in); 297 err = handle_zstd_error(ret, error); 298 if (err) 299 goto out; 300 /* Flush all of the data produced if using flush(). */ 301 if (flush != NULL && out.pos > 0) { 302 if (out.pos != flush(out.dst, out.pos)) { 303 error("Failed to flush()"); 304 err = -1; 305 goto out; 306 } 307 out.pos = 0; 308 } 309 } while (ret != 0); 310 311 if (in_pos != NULL) 312 *in_pos += in.pos; 313 314 err = 0; 315 out: 316 if (in_allocated != NULL) 317 large_free(in_allocated); 318 if (out_allocated != NULL) 319 large_free(out_allocated); 320 if (wksp != NULL) 321 large_free(wksp); 322 return err; 323 } 324 325 #ifndef UNZSTD_PREBOOT 326 STATIC int INIT unzstd(unsigned char *buf, long len, 327 long (*fill)(void*, unsigned long), 328 long (*flush)(void*, unsigned long), 329 unsigned char *out_buf, 330 long *pos, 331 void (*error)(char *x)) 332 { 333 return __unzstd(buf, len, fill, flush, out_buf, 0, pos, error); 334 } 335 #else 336 STATIC int INIT __decompress(unsigned char *buf, long len, 337 long (*fill)(void*, unsigned long), 338 long (*flush)(void*, unsigned long), 339 unsigned char *out_buf, long out_len, 340 long *pos, 341 void (*error)(char *x)) 342 { 343 return __unzstd(buf, len, fill, flush, out_buf, out_len, pos, error); 344 } 345 #endif 346