Multi-client server, protocol versioning, streaming decompression fixes

src/shared/compression.c

@@ -4,15 +4,22 @@
 #include "stdlib.h"
 #include "zstd.h"
 
+#define INITIAL_DECOMPRESS_BUF_SIZE (1024 * 1024)
+
 Data *data_compress(Data *data_to_compress, int compression_level) {
   log_message(LOG_LEVEL_DEBUG, "Starting to compress data");
   size_t dst_size = ZSTD_compressBound(data_to_compress->size);
   Data *compressed_data = data_create_empty(dst_size);
+  if (!compressed_data) {
+    log_message(LOG_LEVEL_ERROR, "Failed to allocate compression buffer");
+    return NULL;
+  }
 
   ZSTD_CCtx *cctx = ZSTD_createCCtx();
   if (!cctx) {
     log_message(LOG_LEVEL_ERROR, "Failed to create ZSTD compression context");
-    exit(EXIT_FAILURE);
+    data_destroy(compressed_data);
+    return NULL;
   }
 
   ZSTD_inBuffer input = {data_to_compress->data, data_to_compress->size, 0};
@@ -24,7 +31,9 @@ Data *data_compress(Data *data_to_compress, int compression_level) {
     if (ZSTD_isError(ret)) {
       log_message(LOG_LEVEL_ERROR, "Compression failed: %s",
                   ZSTD_getErrorName(ret));
-      exit(EXIT_FAILURE);
+      ZSTD_freeCCtx(cctx);
+      data_destroy(compressed_data);
+      return NULL;
     }
   } while (ret > 0);
 
@@ -40,23 +49,26 @@ Data *data_decompress(Data *compressed_data) {
   log_message(LOG_LEVEL_DEBUG, "Start to decompress data");
   unsigned long long dst_size = ZSTD_getFrameContentSize(
       compressed_data->data, compressed_data->size);
-  if (ZSTD_isError(dst_size)) {
-    log_message(LOG_LEVEL_ERROR, "Failed to get decompressed size: %s",
-                ZSTD_getErrorName(dst_size));
-    exit(EXIT_FAILURE);
-  }
-
-  Data *uncompressed_data = data_create_empty((size_t)dst_size);
 
   ZSTD_DCtx *dctx = ZSTD_createDCtx();
   if (!dctx) {
     log_message(LOG_LEVEL_ERROR,
                 "Failed to create ZSTD decompression context");
-    exit(EXIT_FAILURE);
+    return NULL;
+  }
+
+  size_t buf_size = (!ZSTD_isError(dst_size) && dst_size > 0)
+                        ? (size_t)dst_size
+                        : INITIAL_DECOMPRESS_BUF_SIZE;
+  Data *uncompressed_data = data_create_empty(buf_size);
+  if (!uncompressed_data) {
+    log_message(LOG_LEVEL_ERROR, "Failed to allocate decompression buffer");
+    ZSTD_freeDCtx(dctx);
+    return NULL;
   }
 
   ZSTD_inBuffer input = {compressed_data->data, compressed_data->size, 0};
-  ZSTD_outBuffer output = {uncompressed_data->data, (size_t)dst_size, 0};
+  ZSTD_outBuffer output = {uncompressed_data->data, buf_size, 0};
 
   size_t ret;
   do {
@@ -64,7 +76,22 @@ Data *data_decompress(Data *compressed_data) {
     if (ZSTD_isError(ret)) {
       log_message(LOG_LEVEL_ERROR, "Decompression failed: %s",
                   ZSTD_getErrorName(ret));
-      exit(EXIT_FAILURE);
+      ZSTD_freeDCtx(dctx);
+      data_destroy(uncompressed_data);
+      return NULL;
+    }
+    if (ret > 0 && output.pos == output.size) {
+      buf_size *= 2;
+      void *new_data = realloc(uncompressed_data->data, buf_size);
+      if (!new_data) {
+        log_message(LOG_LEVEL_ERROR, "Failed to grow decompression buffer");
+        ZSTD_freeDCtx(dctx);
+        data_destroy(uncompressed_data);
+        return NULL;
+      }
+      uncompressed_data->data = new_data;
+      output.dst = new_data;
+      output.size = buf_size;
     }
   } while (ret > 0);