4 files changed, 56 insertions, 6 deletions

diff --git a/ubifs-utils/mkfs.ubifs/compr.c b/ubifs-utils/mkfs.ubifs/compr.c
index 8eff186..6239835 100644
--- a/ubifs-utils/mkfs.ubifs/compr.c
+++ b/ubifs-utils/mkfs.ubifs/compr.c
@@ -28,6 +28,9 @@
 #include <lzo/lzo1x.h>
 #endif
 #include <linux/types.h>
+#ifndef WITHOUT_ZSTSD
+#include <zstd.h>
+#endif
 
 #define crc32 __zlib_crc32
 #include <zlib.h>
@@ -109,6 +112,25 @@ static int lzo_compress(void *in_buf, size_t in_len, void *out_buf,
 }
 #endif
 
+#ifndef WITHOUT_ZSTD
+static ZSTD_CCtx *zctx;
+
+static int zstd_compress(void *in_buf, size_t in_len, void *out_buf,
+			 size_t *out_len)
+{
+	size_t ret;
+
+	ret = ZSTD_compressCCtx(zctx, out_buf, *out_len, in_buf, in_len,
+				ZSTD_CLEVEL_DEFAULT);
+	if (ZSTD_isError(ret)) {
+		errcnt += 1;
+		return -1;
+	}
+	*out_len = ret;
+	return 0;
+}
+#endif
+
 static int no_compress(void *in_buf, size_t in_len, void *out_buf,
 		       size_t *out_len)
 {
@@ -192,6 +214,11 @@ int compress_data(void *in_buf, size_t in_len, void *out_buf, size_t *out_len,
 		case MKFS_UBIFS_COMPR_ZLIB:
 			ret = zlib_deflate(in_buf, in_len, out_buf, out_len);
 			break;
+#ifndef WITHOUT_ZSTD
+		case MKFS_UBIFS_COMPR_ZSTD:
+			ret = zstd_compress(in_buf, in_len, out_buf, out_len);
+			break;
+#endif
 		case MKFS_UBIFS_COMPR_NONE:
 			ret = 1;
 			break;
@@ -219,18 +246,29 @@ int init_compression(void)
 #endif
 
 	zlib_buf = malloc(UBIFS_BLOCK_SIZE * WORST_COMPR_FACTOR);
-	if (!zlib_buf) {
-		free(lzo_mem);
-		return -1;
-	}
+	if (!zlib_buf)
+		goto err;
+
+#ifndef WITHOUT_ZSTD
+	zctx = ZSTD_createCCtx();
+	if (!zctx)
+		goto err;
+#endif
 
 	return 0;
+err:
+	free(zlib_buf);
+	free(lzo_mem);
+	return -1;
 }
 
 void destroy_compression(void)
 {
 	free(zlib_buf);
 	free(lzo_mem);
+#ifndef WITHOUT_ZSTD
+	ZSTD_freeCCtx(zctx);
+#endif
 	if (errcnt)
 		fprintf(stderr, "%llu compression errors occurred\n", errcnt);
 }
diff --git a/ubifs-utils/mkfs.ubifs/compr.h b/ubifs-utils/mkfs.ubifs/compr.h
index e3dd95c..d58c7c7 100644
--- a/ubifs-utils/mkfs.ubifs/compr.h
+++ b/ubifs-utils/mkfs.ubifs/compr.h
@@ -36,6 +36,7 @@ enum compression_type
 	MKFS_UBIFS_COMPR_NONE,
 	MKFS_UBIFS_COMPR_LZO,
 	MKFS_UBIFS_COMPR_ZLIB,
+	MKFS_UBIFS_COMPR_ZSTD,
 };
 
 int compress_data(void *in_buf, size_t in_len, void *out_buf, size_t *out_len,
diff --git a/ubifs-utils/mkfs.ubifs/mkfs.ubifs.c b/ubifs-utils/mkfs.ubifs/mkfs.ubifs.c
index 4b31979..6d8e55d 100644
--- a/ubifs-utils/mkfs.ubifs/mkfs.ubifs.c
+++ b/ubifs-utils/mkfs.ubifs/mkfs.ubifs.c
@@ -35,6 +35,10 @@
 #include <selinux/label.h>
 #endif
 
+#ifndef WITHOUT_ZSTD
+#include <zstd.h>
+#endif
+
 #include "crypto.h"
 #include "fscrypt.h"
 
@@ -209,8 +213,8 @@ static const char *helptext =
 "-o, --output=FILE        output to FILE\n"
 "-j, --jrn-size=SIZE      journal size\n"
 "-R, --reserved=SIZE      how much space should be reserved for the super-user\n"
-"-x, --compr=TYPE         compression type - \"lzo\", \"favor_lzo\", \"zlib\" or\n"
-"                         \"none\" (default: \"lzo\")\n"
+"-x, --compr=TYPE         compression type - \"lzo\", \"favor_lzo\", \"zlib\"\n"
+"                         \"zstd\" or \"none\" (default: \"lzo\")\n"
 "-X, --favor-percent      may only be used with favor LZO compression and defines\n"
 "                         how many percent better zlib should compress to make\n"
 "                         mkfs.ubifs use zlib instead of LZO (default 20%)\n"
@@ -654,6 +658,10 @@ static int get_options(int argc, char**argv)
 				c->default_compr = UBIFS_COMPR_NONE;
 			else if (strcmp(optarg, "zlib") == 0)
 				c->default_compr = UBIFS_COMPR_ZLIB;
+#ifndef WITHOUT_ZSTD
+			else if (strcmp(optarg, "zstd") == 0)
+				c->default_compr = UBIFS_COMPR_ZSTD;
+#endif
 #ifndef WITHOUT_LZO
 			else if (strcmp(optarg, "favor_lzo") == 0) {
 				c->default_compr = UBIFS_COMPR_LZO;
diff --git a/ubifs-utils/mkfs.ubifs/mkfs.ubifs.h b/ubifs-utils/mkfs.ubifs/mkfs.ubifs.h
index 8f01860..f1425c5 100644
--- a/ubifs-utils/mkfs.ubifs/mkfs.ubifs.h
+++ b/ubifs-utils/mkfs.ubifs/mkfs.ubifs.h
@@ -77,6 +77,9 @@
 #if MKFS_UBIFS_COMPR_ZLIB != UBIFS_COMPR_ZLIB
 #error MKFS_UBIFS_COMPR_ZLIB != UBIFS_COMPR_ZLIB
 #endif
+#if MKFS_UBIFS_COMPR_ZSTD != UBIFS_COMPR_ZSTD
+#error MKFS_UBIFS_COMPR_ZSTD != UBIFS_COMPR_ZSTD
+#endif
 
 extern int verbose;
 extern int debug_level;