1 files changed, 59 insertions, 0 deletions

diff --git a/fs/btrfs/compression.c b/fs/btrfs/compression.c
index e5601b8f2bf..346875d45a1 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 @@ static u32 decompress_zlib(const u8 *_cbuf, u32 clen, u8 *dbuf, u32 dlen)
 	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 @@ u32 btrfs_decompress(u8 type, const char *c, u32 clen, char *d, u32 dlen)
 		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);