[PATCH 1/2] erofs-utils: lib: switch to on-heap fitblk_buffer for libzstd

[Gao Xiang]

 - Allocating VLAs on the stack (or using alloca()) for large sizes
   could exceed the stack limit;

 - It's easier to isolate these buffers on the heap for code sanitizers
   to detect potential bugs.

Signed-off-by: Gao Xiang <hsiangkao at linux.alibaba.com>
---
 lib/compressor_libzstd.c | 83 +++++++++++++++++++++++++++++-----------
 1 file changed, 60 insertions(+), 23 deletions(-)

diff --git a/lib/compressor_libzstd.c b/lib/compressor_libzstd.c
index e53b1a63..3233d723 100644
--- a/lib/compressor_libzstd.c
+++ b/lib/compressor_libzstd.c
@@ -4,18 +4,24 @@
 #include "erofs/config.h"
 #include <zstd.h>
 #include <zstd_errors.h>
-#include <alloca.h>
+#include <stdlib.h>
 #include "compressor.h"
 #include "erofs/atomic.h"
 
+struct erofs_libzstd_context {
+	ZSTD_CCtx *cctx;
+	u8 *fitblk_buffer;
+	unsigned int fitblk_bufsiz;
+};
+
 static int libzstd_compress(const struct erofs_compress *c,
 			    const void *src, unsigned int srcsize,
 			    void *dst, unsigned int dstcapacity)
 {
-	ZSTD_CCtx *cctx = c->private_data;
+	struct erofs_libzstd_context *ctx = c->private_data;
 	size_t csize;
 
-	csize = ZSTD_compress2(cctx, dst, dstcapacity, src, srcsize);
+	csize = ZSTD_compress2(ctx->cctx, dst, dstcapacity, src, srcsize);
 	if (ZSTD_isError(csize)) {
 		if (ZSTD_getErrorCode(csize) == ZSTD_error_dstSize_tooSmall)
 			return -ENOSPC;
@@ -29,12 +35,20 @@ static int libzstd_compress_destsize(const struct erofs_compress *c,
 				     const void *src, unsigned int *srcsize,
 				     void *dst, unsigned int dstsize)
 {
-	ZSTD_CCtx *cctx = c->private_data;
+	struct erofs_libzstd_context *ctx = c->private_data;
 	size_t l = 0;		/* largest input that fits so far */
 	size_t l_csize = 0;
 	size_t r = *srcsize + 1; /* smallest input that doesn't fit so far */
 	size_t m;
-	u8 *fitblk_buffer = alloca(dstsize + 32);
+
+	if (dstsize + 32 > ctx->fitblk_bufsiz) {
+		u8 *buf = realloc(ctx->fitblk_buffer, dstsize + 32);
+
+		if (!buf)
+			return -ENOMEM;
+		ctx->fitblk_bufsiz = dstsize + 32;
+		ctx->fitblk_buffer = buf;
+	}
 
 	m = dstsize * 4;
 	for (;;) {
@@ -43,7 +57,7 @@ static int libzstd_compress_destsize(const struct erofs_compress *c,
 		m = max(m, l + 1);
 		m = min(m, r - 1);
 
-		csize = ZSTD_compress2(cctx, fitblk_buffer,
+		csize = ZSTD_compress2(ctx->cctx, ctx->fitblk_buffer,
 				       dstsize + 32, src, m);
 		if (ZSTD_isError(csize)) {
 			if (ZSTD_getErrorCode(csize) == ZSTD_error_dstSize_tooSmall)
@@ -53,7 +67,7 @@ static int libzstd_compress_destsize(const struct erofs_compress *c,
 
 		if (csize > 0 && csize <= dstsize) {
 			/* Fits */
-			memcpy(dst, fitblk_buffer, csize);
+			memcpy(dst, ctx->fitblk_buffer, csize);
 			l = m;
 			l_csize = csize;
 			if (r <= l + 1 || csize + 1 >= dstsize)
@@ -78,9 +92,14 @@ doesnt_fit:
 
 static int compressor_libzstd_exit(struct erofs_compress *c)
 {
-	if (!c->private_data)
+	struct erofs_libzstd_context *ctx = c->private_data;
+
+	if (!ctx)
 		return -EINVAL;
-	ZSTD_freeCCtx(c->private_data);
+
+	free(ctx->fitblk_buffer);
+	ZSTD_freeCCtx(ctx->cctx);
+	free(ctx);
 	return 0;
 }
 
@@ -118,27 +137,41 @@ static int erofs_compressor_libzstd_setdictsize(struct erofs_compress *c,
 
 static int compressor_libzstd_init(struct erofs_compress *c)
 {
+	struct erofs_libzstd_context *ctx = c->private_data;
 	static erofs_atomic_bool_t __warnonce;
-	ZSTD_CCtx *cctx = c->private_data;
-	size_t err;
+	ZSTD_CCtx *cctx;
+	size_t errcode;
+	int err;
 
-	ZSTD_freeCCtx(cctx);
+	if (ctx) {
+		ZSTD_freeCCtx(ctx->cctx);
+		ctx->cctx = NULL;
+		c->private_data = NULL;
+	} else {
+		ctx = calloc(1, sizeof(*ctx));
+		if (!ctx)
+			return -ENOMEM;
+	}
 	cctx = ZSTD_createCCtx();
-	if (!cctx)
-		return -ENOMEM;
+	if (!cctx) {
+		err = -ENOMEM;
+		goto out_err;
+	}
 
-	err = ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, c->compression_level);
-	if (ZSTD_isError(err)) {
+	err = -EINVAL;
+	errcode = ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, c->compression_level);
+	if (ZSTD_isError(errcode)) {
 		erofs_err("failed to set compression level: %s",
-			  ZSTD_getErrorName(err));
-		return -EINVAL;
+			  ZSTD_getErrorName(errcode));
+		goto out_err;
 	}
-	err = ZSTD_CCtx_setParameter(cctx, ZSTD_c_windowLog, ilog2(c->dict_size));
-	if (ZSTD_isError(err)) {
-		erofs_err("failed to set window log: %s", ZSTD_getErrorName(err));
-		return -EINVAL;
+	errcode = ZSTD_CCtx_setParameter(cctx, ZSTD_c_windowLog, ilog2(c->dict_size));
+	if (ZSTD_isError(errcode)) {
+		erofs_err("failed to set window log: %s", ZSTD_getErrorName(errcode));
+		goto out_err;
 	}
-	c->private_data = cctx;
+	ctx->cctx = cctx;
+	c->private_data = ctx;
 
 	if (!erofs_atomic_test_and_set(&__warnonce)) {
 		erofs_warn("EXPERIMENTAL libzstd compressor in use. Note that `fitblk` isn't supported by upstream zstd for now.");
@@ -146,6 +179,10 @@ static int compressor_libzstd_init(struct erofs_compress *c)
 		erofs_info("You could clarify further needs in zstd repository <https://github.com/facebook/zstd/issues> for reference too.");
 	}
 	return 0;
+out_err:
+	ZSTD_freeCCtx(cctx);
+	free(ctx);
+	return err;
 }
 
 const struct erofs_compressor erofs_compressor_libzstd = {
-- 
2.43.5