zstd: extend support for unknown compressed size

Signed-off-by: Eduardo Silva <[email protected]>

Author: edsiper

src/flb_zstd.c

@@ -24,6 +24,9 @@
 #include <fluent-bit/flb_compression.h>
 #include <fluent-bit/flb_zstd.h>
 
+
+#define FLB_ZSTD_DEFAULT_CHUNK 64 * 1024  /* 64 KB buffer */
+
 size_t flb_zstd_compress(void *in_data, size_t in_len, void **out_data, size_t *out_len)
 {
     void *buf;
@@ -50,8 +53,83 @@ size_t flb_zstd_compress(void *in_data, size_t in_len, void **out_data, size_t *
     return 0;
 }
 
+static int zstd_uncompress_unknown_size(void *in_data, size_t in_len, void **out_data, size_t *out_len)
+{
+    int ret = 0;
+    size_t out_size;
+    char *tmp;
+    void *buf;
+
+    ZSTD_DCtx *dctx;
+    ZSTD_inBuffer input;
+    ZSTD_outBuffer output;
+
+    /* create decompression context */
+    dctx = ZSTD_createDCtx();
+    if (!dctx) {
+        flb_error("[zstd] cannot create decompression context");
+        return -1;
+    }
+
+    /* initial output buffer */
+    out_size = FLB_ZSTD_DEFAULT_CHUNK;
+    buf = flb_malloc(out_size);
+    if (!buf) {
+        flb_errno();
+        ZSTD_freeDCtx(dctx);
+        return -1;
+    }
+
+    /* input */
+    input.src = in_data;
+    input.size = in_len;
+    input.pos = 0;
+
+    /* start the decompress loop */
+    output.dst = buf;
+    output.pos = 0;
+    output.size = out_size;
+
+    while (input.pos < input.size) {
+        ret = ZSTD_decompressStream(dctx, &output, &input);
+        if (ZSTD_isError(ret)) {
+            flb_error("[zstd] decompression failed: %s", ZSTD_getErrorName(ret));
+            flb_free(buf);
+            ZSTD_freeDCtx(dctx);
+            return -1;
+        }
+
+        /* check if we need more space */
+        if (output.pos == out_size) {
+            out_size *= 2;
+            tmp = flb_realloc(buf, out_size);
+            if (!tmp) {
+                flb_errno();
+                flb_free(buf);
+                ZSTD_freeDCtx(dctx);
+                return -1;
+            }
+            buf = tmp;
+            output.dst = buf;
+            output.size = out_size;
+        }
+
+        /* check if we have finished */
+        if (ret == 0) {
+            break;
+        }
+    }
+
+    ZSTD_freeDCtx(dctx);
+
+    *out_data = buf;
+    *out_len = output.pos;
+    return 0;
+}
+
 size_t flb_zstd_uncompress(void *in_data, size_t in_len, void **out_data, size_t *out_len)
 {
+    int ret;
     void *buf;
     size_t size;
 
@@ -60,6 +138,10 @@ size_t flb_zstd_uncompress(void *in_data, size_t in_len, void **out_data, size_t
         flb_error("[zstd] invalid content size");
         return -1;
     }
+    else if (size == ZSTD_CONTENTSIZE_UNKNOWN) {
+        ret = zstd_uncompress_unknown_size(in_data, in_len, out_data, out_len);
+        return ret;
+    }
 
     buf = flb_malloc(size);
     if (!buf) {