lib/decompress_unzstd.c

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