353 lines · c
1// SPDX-License-Identifier: GPL-2.02 3/*4 * Important notes about in-place decompression5 *6 * At least on x86, the kernel is decompressed in place: the compressed data7 * is placed to the end of the output buffer, and the decompressor overwrites8 * most of the compressed data. There must be enough safety margin to9 * 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 of15 * the file is compressed extremely well, and the rest of the file is16 * uncompressible. Thus, we must look for worst-case expansion when the17 * 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 is29 * 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 minimum31 * uncompressed size of the payload is never less than the payload size32 * (excluding the block header).33 *34 * The assumption, that the uncompressed size of the payload is never35 * smaller than the payload itself, is valid only when talking about36 * the payload as a whole. It is possible that the payload has parts where37 * the decompressor consumes more input than it produces output. Calculating38 * 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 bytes40 * of the payload which it is currently reading.41 *42 * Now we have enough information to calculate the safety margin. We need43 * - 22 bytes for the .zst file format headers;44 * - 3 bytes per every 128 KiB of uncompressed size (one block header per45 * block); and46 * - 128 KiB (biggest possible zstd block size) to make sure that the47 * decompressor never overwrites anything from the block it is currently48 * reading.49 *50 * We get the following formula:51 *52 * safety_margin = 22 + uncompressed_size * 3 / 131072 + 13107253 * <= 22 + (uncompressed_size >> 15) + 13107254 */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 source60 * dependencies.61 *62 * When UNZSTD_PREBOOT is defined we declare __decompress(), which is63 * used for kernel decompression, instead of unzstd().64 *65 * Define __DISABLE_EXPORTS in preboot environments to prevent symbols66 * from xxhash and zstd from being exported by the EXPORT_SYMBOL macro.67 */68#ifdef STATIC69# define UNZSTD_PREBOOT70# include "xxhash.c"71# include "zstd/decompress_sources.h"72#else73#include <linux/decompress/unzstd.h>74#endif75 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 for86 * 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 at142 * 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 big185 * that its end address overflows. When the size is not provided, make186 * 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 a194 * 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 allocate201 * a large enough buffer. If it is provided, it must be at least202 * 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 sliding245 * window. The window size varies from 1 KB to ZSTD_WINDOWSIZE_MAX246 * (8 MB), so it is important to use the actual value so as not to247 * 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 is266 * 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 can288 * 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_PREBOOT333STATIC 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#else343STATIC 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#endif353