This adds the support to bootm to decompress and boot ZSTD compressed
kernels. The decompress_unzstd.c was taken from Linux and slighlty
adapted. Also the unzstd() function is adapted since we don't support
large >2G images yet like Linux does.
Signed-off-by: Marco Felsch <m.felsch@xxxxxxxxxxxxxx>
---
common/bootm.c | 8 +
include/linux/decompress/unzstd.h | 11 ++
lib/Makefile | 1 +
lib/decompress_unzstd.c | 316 ++++++++++++++++++++++++++++++
lib/uncompress.c | 6 +
5 files changed, 342 insertions(+)
create mode 100644 include/linux/decompress/unzstd.h
create mode 100644 lib/decompress_unzstd.c
diff --git a/common/bootm.c b/common/bootm.c
index 29ea13e28cb8..4cc88eed76b5 100644
--- a/common/bootm.c
+++ b/common/bootm.c
@@ -940,6 +940,12 @@ static struct image_handler xz_bootm_handler = {
.filetype = filetype_xz_compressed,
};
+static struct image_handler zstd_bootm_handler = {
+ .name = "ZSTD compressed file",
+ .bootm = do_bootm_compressed,
+ .filetype = filetype_zstd_compressed,
+};
+
static int bootm_init(void)
{
globalvar_add_simple("bootm.image", NULL);
@@ -974,6 +980,8 @@ static int bootm_init(void)
register_image_handler(&lz4_bootm_handler);
if (IS_ENABLED(CONFIG_XZ_DECOMPRESS))
register_image_handler(&xz_bootm_handler);
+ if (IS_ENABLED(CONFIG_ZSTD_DECOMPRESS))
+ register_image_handler(&zstd_bootm_handler);
return 0;
}
diff --git a/include/linux/decompress/unzstd.h b/include/linux/decompress/unzstd.h
new file mode 100644
index 000000000000..ac078a3ba151
--- /dev/null
+++ b/include/linux/decompress/unzstd.h
@@ -0,0 +1,11 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef LINUX_DECOMPRESS_UNZSTD_H
+#define LINUX_DECOMPRESS_UNZSTD_H
+
+int unzstd(unsigned char *inbuf, int len,
+ int (*fill)(void*, unsigned int),
+ int (*flush)(void*, unsigned int),
+ unsigned char *output,
+ int *pos,
+ void (*error_fn)(char *x));
+#endif
diff --git a/lib/Makefile b/lib/Makefile
index 853d8870fe14..54866f59cc05 100644
--- a/lib/Makefile
+++ b/lib/Makefile
@@ -49,6 +49,7 @@ obj-$(CONFIG_ZSTD_DECOMPRESS) += zstd/
obj-y += show_progress.o
obj-$(CONFIG_LZO_DECOMPRESS) += decompress_unlzo.o
obj-$(CONFIG_LZ4_DECOMPRESS) += decompress_unlz4.o
+obj-$(CONFIG_ZSTD_DECOMPRESS) += decompress_unzstd.o
obj-$(CONFIG_PROCESS_ESCAPE_SEQUENCE) += process_escape_sequence.o
obj-$(CONFIG_UNCOMPRESS) += uncompress.o
obj-$(CONFIG_BCH) += bch.o
diff --git a/lib/decompress_unzstd.c b/lib/decompress_unzstd.c
new file mode 100644
index 000000000000..42166c21e255
--- /dev/null
+++ b/lib/decompress_unzstd.c
@@ -0,0 +1,316 @@
+// SPDX-License-Identifier: GPL-2.0
+
+/*
+ * Important notes about in-place decompression
+ *
+ * At least on x86, the kernel is decompressed in place: the compressed data
+ * is placed to the end of the output buffer, and the decompressor overwrites
+ * most of the compressed data. There must be enough safety margin to
+ * guarantee that the write position is always behind the read position.
+ *
+ * The safety margin for ZSTD with a 128 KB block size is calculated below.
+ * Note that the margin with ZSTD is bigger than with GZIP or XZ!
+ *
+ * The worst case for in-place decompression is that the beginning of
+ * the file is compressed extremely well, and the rest of the file is
+ * uncompressible. Thus, we must look for worst-case expansion when the
+ * compressor is encoding uncompressible data.
+ *
+ * The structure of the .zst file in case of a compressed kernel is as follows.
+ * Maximum sizes (as bytes) of the fields are in parenthesis.
+ *
+ * Frame Header: (18)
+ * Blocks: (N)
+ * Checksum: (4)
+ *
+ * The frame header and checksum overhead is at most 22 bytes.
+ *
+ * ZSTD stores the data in blocks. Each block has a header whose size is
+ * a 3 bytes. After the block header, there is up to 128 KB of payload.
+ * The maximum uncompressed size of the payload is 128 KB. The minimum
+ * uncompressed size of the payload is never less than the payload size
+ * (excluding the block header).
+ *
+ * The assumption, that the uncompressed size of the payload is never
+ * smaller than the payload itself, is valid only when talking about
+ * the payload as a whole. It is possible that the payload has parts where
+ * the decompressor consumes more input than it produces output. Calculating
+ * the worst case for this would be tricky. Instead of trying to do that,
+ * let's simply make sure that the decompressor never overwrites any bytes
+ * of the payload which it is currently reading.
+ *
+ * Now we have enough information to calculate the safety margin. We need
+ * - 22 bytes for the .zst file format headers;
+ * - 3 bytes per every 128 KiB of uncompressed size (one block header per
+ * block); and
+ * - 128 KiB (biggest possible zstd block size) to make sure that the
+ * decompressor never overwrites anything from the block it is currently
+ * reading.
+ *
+ * We get the following formula:
+ *
+ * safety_margin = 22 + uncompressed_size * 3 / 131072 + 131072
+ * <= 22 + (uncompressed_size >> 15) + 131072
+ */
+
+#include <linux/decompress/mm.h>
+#include <linux/kernel.h>
+#include <linux/zstd.h>
+
+/* 128MB is the maximum window size supported by zstd. */
+#define ZSTD_WINDOWSIZE_MAX (1 << ZSTD_WINDOWLOG_MAX)
+/*
+ * Size of the input and output buffers in multi-call mode.
+ * Pick a larger size because it isn't used during kernel decompression,
+ * since that is single pass, and we have to allocate a large buffer for
+ * zstd's window anyway. The larger size speeds up initramfs decompression.
+ */
+#define ZSTD_IOBUF_SIZE (1 << 17)
+
+static int INIT handle_zstd_error(size_t ret, void (*error)(char *x))
+{
+ const int err = ZSTD_getErrorCode(ret);
+
+ if (!ZSTD_isError(ret))
+ return 0;
+
+ switch (err) {
+ case ZSTD_error_memory_allocation:
+ error("ZSTD decompressor ran out of memory");
+ break;
+ case ZSTD_error_prefix_unknown:
+ error("Input is not in the ZSTD format (wrong magic bytes)");
+ break;
+ case ZSTD_error_dstSize_tooSmall:
+ case ZSTD_error_corruption_detected:
+ case ZSTD_error_checksum_wrong:
+ error("ZSTD-compressed data is corrupt");
+ break;
+ default:
+ error("ZSTD-compressed data is probably corrupt");
+ break;
+ }
+ return -1;
+}
+
+/*
+ * Handle the case where we have the entire input and output in one segment.
+ * We can allocate less memory (no circular buffer for the sliding window),
+ * and avoid some memcpy() calls.
+ */
+static int INIT decompress_single(const u8 *in_buf, long in_len, u8 *out_buf,
+ long out_len, long *in_pos,
+ void (*error)(char *x))
+{
+ const size_t wksp_size = ZSTD_DCtxWorkspaceBound();
+ void *wksp = large_malloc(wksp_size);
+ ZSTD_DCtx *dctx = ZSTD_initDCtx(wksp, wksp_size);
+ int err;
+ size_t ret;
+
+ if (dctx == NULL) {
+ error("Out of memory while allocating ZSTD_DCtx");
+ err = -1;
+ goto out;
+ }
+ /*
+ * Find out how large the frame actually is, there may be junk at
+ * the end of the frame that ZSTD_decompressDCtx() can't handle.
+ */
+ ret = ZSTD_findFrameCompressedSize(in_buf, in_len);
+ err = handle_zstd_error(ret, error);
+ if (err)
+ goto out;
+ in_len = (long)ret;
+
+ ret = ZSTD_decompressDCtx(dctx, out_buf, out_len, in_buf, in_len);
+ err = handle_zstd_error(ret, error);
+ if (err)
+ goto out;
+
+ if (in_pos != NULL)
+ *in_pos = in_len;
+
+ err = 0;
+out:
+ if (wksp != NULL)
+ large_free(wksp);
+ return err;
+}
+
+static int INIT __unzstd(unsigned char *in_buf, long in_len,
+ long (*fill)(void*, unsigned long),
+ long (*flush)(void*, unsigned long),
+ unsigned char *out_buf, long out_len,
+ long *in_pos,
+ void (*error)(char *x))
+{
+ ZSTD_inBuffer in;
+ ZSTD_outBuffer out;
+ ZSTD_frameParams params;
+ void *in_allocated = NULL;
+ void *out_allocated = NULL;
+ void *wksp = NULL;
+ size_t wksp_size;
+ ZSTD_DStream *dstream;
+ int err;
+ size_t ret;
+
+ /*
+ * ZSTD decompression code won't be happy if the buffer size is so big
+ * that its end address overflows. When the size is not provided, make
+ * it as big as possible without having the end address overflow.
+ */
+ if (out_len == 0)
+ out_len = UINTPTR_MAX - (uintptr_t)out_buf;
+
+ if (fill == NULL && flush == NULL)
+ /*
+ * We can decompress faster and with less memory when we have a
+ * single chunk.
+ */
+ return decompress_single(in_buf, in_len, out_buf, out_len,
+ in_pos, error);
+
+ /*
+ * If in_buf is not provided, we must be using fill(), so allocate
+ * a large enough buffer. If it is provided, it must be at least
+ * ZSTD_IOBUF_SIZE large.
+ */
+ if (in_buf == NULL) {
+ in_allocated = large_malloc(ZSTD_IOBUF_SIZE);
+ if (in_allocated == NULL) {
+ error("Out of memory while allocating input buffer");
+ err = -1;
+ goto out;
+ }
+ in_buf = in_allocated;
+ in_len = 0;
+ }
+ /* Read the first chunk, since we need to decode the frame header. */
+ if (fill != NULL)
+ in_len = fill(in_buf, ZSTD_IOBUF_SIZE);
+ if (in_len < 0) {
+ error("ZSTD-compressed data is truncated");
+ err = -1;
+ goto out;
+ }
+ /* Set the first non-empty input buffer. */
+ in.src = in_buf;
+ in.pos = 0;
+ in.size = in_len;
+ /* Allocate the output buffer if we are using flush(). */
+ if (flush != NULL) {
+ out_allocated = large_malloc(ZSTD_IOBUF_SIZE);
+ if (out_allocated == NULL) {
+ error("Out of memory while allocating output buffer");
+ err = -1;
+ goto out;
+ }
+ out_buf = out_allocated;
+ out_len = ZSTD_IOBUF_SIZE;
+ }
+ /* Set the output buffer. */
+ out.dst = out_buf;
+ out.pos = 0;
+ out.size = out_len;
+
+ /*
+ * We need to know the window size to allocate the ZSTD_DStream.
+ * Since we are streaming, we need to allocate a buffer for the sliding
+ * window. The window size varies from 1 KB to ZSTD_WINDOWSIZE_MAX
+ * (8 MB), so it is important to use the actual value so as not to
+ * waste memory when it is smaller.
+ */
+ ret = ZSTD_getFrameParams(¶ms, in.src, in.size);
+ err = handle_zstd_error(ret, error);
+ if (err)
+ goto out;
+ if (ret != 0) {
+ error("ZSTD-compressed data has an incomplete frame header");
+ err = -1;
+ goto out;
+ }
+ if (params.windowSize > ZSTD_WINDOWSIZE_MAX) {
+ error("ZSTD-compressed data has too large a window size");
+ err = -1;
+ goto out;
+ }
+
+ /*
+ * Allocate the ZSTD_DStream now that we know how much memory is
+ * required.
+ */
+ wksp_size = ZSTD_DStreamWorkspaceBound(params.windowSize);
+ wksp = large_malloc(wksp_size);
+ dstream = ZSTD_initDStream(params.windowSize, wksp, wksp_size);
+ if (dstream == NULL) {
+ error("Out of memory while allocating ZSTD_DStream");
+ err = -1;
+ goto out;
+ }
+
+ /*
+ * Decompression loop:
+ * Read more data if necessary (error if no more data can be read).
+ * Call the decompression function, which returns 0 when finished.
+ * Flush any data produced if using flush().
+ */
+ if (in_pos != NULL)
+ *in_pos = 0;
+ do {
+ /*
+ * If we need to reload data, either we have fill() and can
+ * try to get more data, or we don't and the input is truncated.
+ */
+ if (in.pos == in.size) {
+ if (in_pos != NULL)
+ *in_pos += in.pos;
+ in_len = fill ? fill(in_buf, ZSTD_IOBUF_SIZE) : -1;
+ if (in_len < 0) {
+ error("ZSTD-compressed data is truncated");
+ err = -1;
+ goto out;
+ }
+ in.pos = 0;
+ in.size = in_len;
+ }
+ /* Returns zero when the frame is complete. */
+ ret = ZSTD_decompressStream(dstream, &out, &in);
+ err = handle_zstd_error(ret, error);
+ if (err)
+ goto out;
+ /* Flush all of the data produced if using flush(). */
+ if (flush != NULL && out.pos > 0) {
+ if (out.pos != flush(out.dst, out.pos)) {
+ error("Failed to flush()");
+ err = -1;
+ goto out;
+ }
+ out.pos = 0;
+ }
+ } while (ret != 0);
+
+ if (in_pos != NULL)
+ *in_pos += in.pos;
+
+ err = 0;
+out:
+ if (in_allocated != NULL)
+ large_free(in_allocated);
+ if (out_allocated != NULL)
+ large_free(out_allocated);
+ if (wksp != NULL)
+ large_free(wksp);
+ return err;
+}
+
+STATIC int INIT unzstd(unsigned char *buf, int len,
+ int (*fill)(void*, unsigned int),
+ int (*flush)(void*, unsigned int),
+ unsigned char *out_buf,
+ int *pos,
+ void (*error)(char *x))
+{
+ return __unzstd(buf, len, fill, flush, out_buf, 0, pos, error);
+}
diff --git a/lib/uncompress.c b/lib/uncompress.c
index bfe042fcf83e..e7a3ffde823e 100644
--- a/lib/uncompress.c
+++ b/lib/uncompress.c
@@ -20,6 +20,7 @@
#include <lzo.h>
#include <linux/xz.h>
#include <linux/decompress/unlz4.h>
+#include <linux/decompress/unzstd.h>
#include <errno.h>
#include <filetype.h>
#include <malloc.h>
@@ -121,6 +122,11 @@ int uncompress(unsigned char *inbuf, int len,
case filetype_xz_compressed:
compfn = decompress_unxz;
break;
+#endif
+#ifdef CONFIG_ZSTD_DECOMPRESS
+ case filetype_zstd_compressed:
+ compfn = unzstd;
+ break;
#endif
default:
err = basprintf("cannot handle filetype %s",
--
2.39.2
