fs: btrfs: add zstd decompression support

This adds decompression support for Zstandard, which has been included
in Linux btrfs driver for some time.

Signed-off-by: Marek BehĂșn <marek.behun@nic.cz>
diff --git a/fs/btrfs/Kconfig b/fs/btrfs/Kconfig
index 22909d9..f302b1f 100644
--- a/fs/btrfs/Kconfig
+++ b/fs/btrfs/Kconfig
@@ -2,6 +2,7 @@
 	bool "Enable BTRFS filesystem support"
 	select CRC32C
 	select LZO
+	select ZSTD
 	select RBTREE
 	help
 	  This provides a single-device read-only BTRFS support. BTRFS is a
diff --git a/fs/btrfs/btrfs_tree.h b/fs/btrfs/btrfs_tree.h
index f90fbb2..aa0f3d6 100644
--- a/fs/btrfs/btrfs_tree.h
+++ b/fs/btrfs/btrfs_tree.h
@@ -647,8 +647,9 @@
 	BTRFS_COMPRESS_NONE  = 0,
 	BTRFS_COMPRESS_ZLIB  = 1,
 	BTRFS_COMPRESS_LZO   = 2,
-	BTRFS_COMPRESS_TYPES = 2,
-	BTRFS_COMPRESS_LAST  = 3,
+	BTRFS_COMPRESS_ZSTD  = 3,
+	BTRFS_COMPRESS_TYPES = 3,
+	BTRFS_COMPRESS_LAST  = 4,
 };
 
 struct btrfs_file_extent_item {
diff --git a/fs/btrfs/compression.c b/fs/btrfs/compression.c
index e5601b8..346875d 100644
--- a/fs/btrfs/compression.c
+++ b/fs/btrfs/compression.c
@@ -6,7 +6,9 @@
  */
 
 #include "btrfs.h"
+#include <malloc.h>
 #include <linux/lzo.h>
+#include <linux/zstd.h>
 #include <u-boot/zlib.h>
 #include <asm/unaligned.h>
 
@@ -108,6 +110,61 @@
 	return res;
 }
 
+#define ZSTD_BTRFS_MAX_WINDOWLOG 17
+#define ZSTD_BTRFS_MAX_INPUT (1 << ZSTD_BTRFS_MAX_WINDOWLOG)
+
+static u32 decompress_zstd(const u8 *cbuf, u32 clen, u8 *dbuf, u32 dlen)
+{
+	ZSTD_DStream *dstream;
+	ZSTD_inBuffer in_buf;
+	ZSTD_outBuffer out_buf;
+	void *workspace;
+	size_t wsize;
+	u32 res = -1;
+
+	wsize = ZSTD_DStreamWorkspaceBound(ZSTD_BTRFS_MAX_INPUT);
+	workspace = malloc(wsize);
+	if (!workspace) {
+		debug("%s: cannot allocate workspace of size %zu\n", __func__,
+		      wsize);
+		return -1;
+	}
+
+	dstream = ZSTD_initDStream(ZSTD_BTRFS_MAX_INPUT, workspace, wsize);
+	if (!dstream) {
+		printf("%s: ZSTD_initDStream failed\n", __func__);
+		goto err_free;
+	}
+
+	in_buf.src = cbuf;
+	in_buf.pos = 0;
+	in_buf.size = clen;
+
+	out_buf.dst = dbuf;
+	out_buf.pos = 0;
+	out_buf.size = dlen;
+
+	while (1) {
+		size_t ret;
+
+		ret = ZSTD_decompressStream(dstream, &out_buf, &in_buf);
+		if (ZSTD_isError(ret)) {
+			printf("%s: ZSTD_decompressStream error %d\n", __func__,
+			       ZSTD_getErrorCode(ret));
+			goto err_free;
+		}
+
+		if (in_buf.pos >= clen || !ret)
+			break;
+	}
+
+	res = out_buf.pos;
+
+err_free:
+	free(workspace);
+	return res;
+}
+
 u32 btrfs_decompress(u8 type, const char *c, u32 clen, char *d, u32 dlen)
 {
 	u32 res;
@@ -126,6 +183,8 @@
 		return decompress_zlib(cbuf, clen, dbuf, dlen);
 	case BTRFS_COMPRESS_LZO:
 		return decompress_lzo(cbuf, clen, dbuf, dlen);
+	case BTRFS_COMPRESS_ZSTD:
+		return decompress_zstd(cbuf, clen, dbuf, dlen);
 	default:
 		printf("%s: Unsupported compression in extent: %i\n", __func__,
 		       type);