zstd: Create a function for use from U-Boot

The existing zstd API requires the same sequence of calls to perform its
task. Create a helper for U-Boot, to avoid code duplication, as is done
with other compression algorithms. Make use of of this from the image
code.

Note that the zstd code lacks a test in test/compression.c and this should
be added by the maintainer.

Signed-off-by: Simon Glass <sjg@chromium.org>

2 files changed, 65 insertions, 1 deletions

diff --git a/lib/zstd/Makefile b/lib/zstd/Makefile
index 33c1df48ec..1217089292 100644
--- a/lib/zstd/Makefile
+++ b/lib/zstd/Makefile
@@ -1,4 +1,4 @@
 obj-y += zstd_decompress.o
 
 zstd_decompress-y := huf_decompress.o decompress.o \
-		     entropy_common.o fse_decompress.o zstd_common.o
+		     entropy_common.o fse_decompress.o zstd_common.o zstd.o
diff --git a/lib/zstd/zstd.c b/lib/zstd/zstd.c
new file mode 100644
index 0000000000..bf9cd19cfa
--- /dev/null
+++ b/lib/zstd/zstd.c
@@ -0,0 +1,64 @@
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Copyright 2021 Google LLC
+ */
+
+#define LOG_CATEGORY	LOGC_BOOT
+
+#include <common.h>
+#include <abuf.h>
+#include <log.h>
+#include <malloc.h>
+#include <linux/zstd.h>
+
+int zstd_decompress(struct abuf *in, struct abuf *out)
+{
+	ZSTD_DStream *dstream;
+	ZSTD_inBuffer in_buf;
+	ZSTD_outBuffer out_buf;
+	void *workspace;
+	size_t wsize;
+	int ret;
+
+	wsize = ZSTD_DStreamWorkspaceBound(abuf_size(in));
+	workspace = malloc(wsize);
+	if (!workspace) {
+		debug("%s: cannot allocate workspace of size %zu\n", __func__,
+			wsize);
+		return -ENOMEM;
+	}
+
+	dstream = ZSTD_initDStream(abuf_size(in), workspace, wsize);
+	if (!dstream) {
+		log_err("%s: ZSTD_initDStream failed\n", __func__);
+		ret = -EPERM;
+		goto do_free;
+	}
+
+	in_buf.src = abuf_data(in);
+	in_buf.pos = 0;
+	in_buf.size = abuf_size(in);
+
+	out_buf.dst = abuf_data(out);
+	out_buf.pos = 0;
+	out_buf.size = abuf_size(out);
+
+	while (1) {
+		size_t res;
+
+		res = ZSTD_decompressStream(dstream, &out_buf, &in_buf);
+		if (ZSTD_isError(res)) {
+			ret = ZSTD_getErrorCode(res);
+			log_err("ZSTD_decompressStream error %d\n", ret);
+			goto do_free;
+		}
+
+		if (in_buf.pos >= abuf_size(in) || !res)
+			break;
+	}
+
+	ret = out_buf.pos;
+do_free:
+	free(workspace);
+	return ret;
+}