[PATCH] zstd compressed firmware loading

[Rene Rebe <rene@exactcode.com>]

Now with zstd compressed kernel & initrd upstream, we would rather
compress everything with one type of compressor, so I added support
for zstd compressed firmware loading, too.

Tested on x86-64, sparc64 and mips64.

Signed-off-by: René Rebe <rene@exactcode.de>

diff --git a/drivers/base/firmware_loader/Kconfig b/drivers/base/firmware_loader/Kconfig
index 5b24f3959255..30d440bec257 100644
--- a/drivers/base/firmware_loader/Kconfig
+++ b/drivers/base/firmware_loader/Kconfig
@@ -169,6 +169,16 @@ config FW_LOADER_COMPRESS
 	  be compressed with either none or crc32 integrity check type (pass
 	  "-C crc32" option to xz command).
 
+config FW_LOADER_COMPRESS_ZSTD
+	bool "Enable Zstd compressed firmware support"
+	select FW_LOADER_PAGED_BUF
+	select ZSTD_DECOMPRESS
+	help
+	  This option enables the support for loading Zstd compressed firmware
+	  files. The caller of firmware API receives the decompressed file
+	  content. The compressed file is loaded as a fallback, only after
+	  loading the raw file failed at first.
+
 config FW_CACHE
 	bool "Enable firmware caching during suspend"
 	depends on PM_SLEEP
diff --git a/drivers/base/firmware_loader/main.c b/drivers/base/firmware_loader/main.c
index 9da0c9d5f538..5deb73dbe81e 100644
--- a/drivers/base/firmware_loader/main.c
+++ b/drivers/base/firmware_loader/main.c
@@ -33,7 +33,9 @@
 #include <linux/syscore_ops.h>
 #include <linux/reboot.h>
 #include <linux/security.h>
+#include <linux/decompress/mm.h>
 #include <linux/xz.h>
+#include <linux/zstd.h>
 
 #include <generated/utsrelease.h>
 
@@ -435,6 +437,140 @@ static int fw_decompress_xz(struct device *dev, struct fw_priv *fw_priv,
 }
 #endif /* CONFIG_FW_LOADER_COMPRESS */
 
+/*
+ * Zstd-compressed firmware support
+ */
+#ifdef CONFIG_FW_LOADER_COMPRESS_ZSTD
+/* show an error and return the standard error code */
+static int handle_zstd_error(size_t ret)
+{
+	const int err = ZSTD_getErrorCode(ret);
+
+	if (!ZSTD_isError(ret))
+		return 0;
+
+	switch (err) {
+	case ZSTD_error_memory_allocation:
+		printk("ZSTD decompressor ran out of memory");
+		break;
+	case ZSTD_error_prefix_unknown:
+		printk("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:
+		printk("ZSTD-compressed data is corrupt");
+		break;
+	default:
+		printk("ZSTD-compressed data is probably corrupt");
+		break;
+	}
+	return -1;
+}
+
+/* single-shot decompression onto the pre-allocated buffer */
+static int fw_decompress_zstd_single(struct device *dev, struct fw_priv *fw_priv,
+				   size_t in_size, const void *in_buffer)
+{
+	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;
+
+	printk("zstd_single\n");
+
+	if (dctx == NULL) {
+	        dev_warn(dev, "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_buffer, in_size);
+	err = handle_zstd_error(ret);
+	if (err)
+		goto out;
+	in_size = (long)ret;
+
+	ret = ZSTD_decompressDCtx(dctx, fw_priv->data, fw_priv->allocated_size, in_buffer, in_size);
+	err = handle_zstd_error(ret);
+	if (err)
+		goto out;
+
+	fw_priv->size = ret;
+
+out:
+	if (wksp != NULL)
+		large_free(wksp);
+	return err;
+}
+
+/* decompression on paged buffer and map it */
+static int fw_decompress_zstd_pages(struct device *dev, struct fw_priv *fw_priv,
+				  size_t in_size, const void *in_buffer)
+{
+#if 0
+	struct zstd_dec *zstd_dec;
+	struct zstd_buf zstd_buf;
+	enum zstd_ret zstd_ret;
+	struct page *page;
+	int err = 0;
+
+	printk("zstd_pages\n");
+
+	zstd_dec = zstd_dec_init(ZSTD_DYNALLOC, (u32)-1);
+	if (!zstd_dec)
+		return -ENOMEM;
+
+	zstd_buf.in_size = in_size;
+	zstd_buf.in = in_buffer;
+	zstd_buf.in_pos = 0;
+
+	fw_priv->is_paged_buf = true;
+	fw_priv->size = 0;
+	do {
+		if (fw_grow_paged_buf(fw_priv, fw_priv->nr_pages + 1)) {
+			err = -ENOMEM;
+			goto out;
+		}
+
+		/* decompress onto the new allocated page */
+		page = fw_priv->pages[fw_priv->nr_pages - 1];
+		zstd_buf.out = kmap(page);
+		zstd_buf.out_pos = 0;
+		zstd_buf.out_size = PAGE_SIZE;
+		zstd_ret = zstd_dec_run(zstd_dec, &zstd_buf);
+		kunmap(page);
+		fw_priv->size += zstd_buf.out_pos;
+		/* partial decompression means either end or error */
+		if (zstd_buf.out_pos != PAGE_SIZE)
+			break;
+	} while (zstd_ret == ZSTD_OK);
+
+	err = fw_decompress_zstd_error(dev, zstd_ret);
+	if (!err)
+		err = fw_map_paged_buf(fw_priv);
+
+ out:
+	zstd_dec_end(zstd_dec);
+	return err;
+#endif
+	return -ENOMEM;
+}
+
+static int fw_decompress_zstd(struct device *dev, struct fw_priv *fw_priv,
+			    size_t in_size, const void *in_buffer)
+{
+	/* if the buffer is pre-allocated, we can perform in single-shot mode */
+	if (fw_priv->data)
+		return fw_decompress_zstd_single(dev, fw_priv, in_size, in_buffer);
+	else
+		return fw_decompress_zstd_pages(dev, fw_priv, in_size, in_buffer);
+}
+#endif /* CONFIG_FW_LOADER_COMPRESS_ZSTD */
+
 /* direct firmware loading support */
 static char fw_path_para[256];
 static const char * const fw_path[] = {
@@ -773,6 +909,11 @@ _request_firmware(const struct firmware **firmware_p, const char *name,
 		ret = fw_get_filesystem_firmware(device, fw->priv, ".xz",
 						 fw_decompress_xz);
 #endif
+#ifdef CONFIG_FW_LOADER_COMPRESS_ZSTD
+	if (ret == -ENOENT)
+		ret = fw_get_filesystem_firmware(device, fw->priv, ".zst",
+						 fw_decompress_zstd);
+#endif
 
 	if (ret == -ENOENT)
 		ret = firmware_fallback_platform(fw->priv, opt_flags);

-- 
  René Rebe, ExactCODE GmbH, Lietzenburger Str. 42, DE-10789 Berlin
  https://exactcode.com | https://t2sde.org | https://rene.rebe.de