Add Zstandard (zstd) compression support to the compress plugin

Author: JakeChampion

CMakeLists.txt

@@ -338,6 +338,11 @@ set(TS_USE_MALLOC_ALLOCATOR ${ENABLE_MALLOC_ALLOCATOR})
 set(TS_USE_ALLOCATOR_METRICS ${ENABLE_ALLOCATOR_METRICS})
 find_package(ZLIB REQUIRED)
 
+find_package(zstd)
+if(zstd_FOUND)
+  set(HAVE_ZSTD_H TRUE)
+endif()
+
 # ncurses is used in traffic_top
 find_package(Curses)
 set(HAVE_CURSES_H ${CURSES_HAVE_CURSES_H})

plugins/compress/CMakeLists.txt

@@ -20,5 +20,10 @@ target_link_libraries(compress PRIVATE libswoc::libswoc)
 if(HAVE_BROTLI_ENCODE_H)
   target_link_libraries(compress PRIVATE brotli::brotlienc)
 endif()
+
+if(HAVE_ZSTD_H)
+  target_link_libraries(compress PRIVATE zstd)
+endif()
+
 verify_global_plugin(compress)
 verify_remap_plugin(compress)

plugins/compress/compress.cc

@@ -1,6 +1,6 @@
 /** @file
 
-  Transforms content using gzip, deflate or brotli
+  Transforms content using gzip, deflate, brotli or zstd
 
   @section license License
 
@@ -23,6 +23,9 @@
 
 #include <cstring>
 #include <zlib.h>
+#if HAVE_ZSTD_F
+#include <zstd.h>
+#endif
 
 #include "ts/apidefs.h"
 #include "tscore/ink_config.h"
@@ -71,6 +74,10 @@ const int BROTLI_COMPRESSION_LEVEL = 6;
 const int BROTLI_LGW               = 16;
 #endif
 
+#if HAVE_ZSTD_H
+const int ZSTD_COMPRESSION_LEVEL = 6;
+#endif
+
 static const char *global_hidden_header_name = nullptr;
 
 static TSMutex compress_config_mutex = nullptr;
@@ -191,6 +198,16 @@ data_alloc(int compression_type, int compression_algorithms)
     data->bstrm.avail_out = 0;
     data->bstrm.total_out = 0;
   }
+#endif
+#if HAVE_ZSTD_H
+  data->zstd_ctx = nullptr;
+  if (compression_type & COMPRESSION_TYPE_ZSTD) {
+    debug("zstd compression. Create Zstd Compression Context.");
+    data->zstd_ctx = ZSTD_createCCtx();
+    if (!data->zstd_ctx) {
+      fatal("Zstd Compression Context Creation Failed");
+    }
+  }
 #endif
   return data;
 }
@@ -212,6 +229,11 @@ data_destroy(Data *data)
 #if HAVE_BROTLI_ENCODE_H
   BrotliEncoderDestroyInstance(data->bstrm.br);
 #endif
+#if HAVE_ZSTD_H
+  if (data->zstd_ctx) {
+    ZSTD_freeCCtx(data->zstd_ctx);
+  }
+#endif
 
   TSfree(data);
 }
@@ -233,6 +255,9 @@ content_encoding_header(TSMBuffer bufp, TSMLoc hdr_loc, const int compression_ty
   } else if (compression_type & COMPRESSION_TYPE_DEFLATE && (algorithm & ALGORITHM_DEFLATE)) {
     value     = TS_HTTP_VALUE_DEFLATE;
     value_len = TS_HTTP_LEN_DEFLATE;
+  } else if (compression_type & COMPRESSION_TYPE_ZSTD && (algorithm & ALGORITHM_ZSTD)) {
+    value     = "zstd";
+    value_len = strlen("zstd");
   }
 
   if (value_len == 0) {
@@ -357,6 +382,16 @@ compress_transform_init(TSCont contp, Data *data)
     data->downstream_vio    = TSVConnWrite(downstream_conn, contp, data->downstream_reader, INT64_MAX);
   }
 
+#if HAVE_ZSTD_H
+  if (data->compression_type & COMPRESSION_TYPE_ZSTD) {
+    zstd_compress_init(data);
+    if (!data->zstd_cctx) {
+      TSError("Failed to create Zstandard compression context");
+      return;
+    }
+  }
+#endif
+
   TSHandleMLocRelease(bufp, TS_NULL_MLOC, hdr_loc);
 }
 
@@ -462,6 +497,39 @@ brotli_transform_one(Data *data, const char *upstream_buffer, int64_t upstream_l
 }
 #endif
 
+#if HAVE_ZSTD_H
+static void
+zstd_compress_init(Data *data)
+{
+  data->zstd_cctx = ZSTD_createCCtx();
+  if (!data->zstd_cctx) {
+    error("Failed to initialize Zstd compression context");
+  }
+}
+
+static void
+zstd_compress_finish(Data *data)
+{
+  if (data->zstd_cctx) {
+    ZSTD_freeCCtx(data->zstd_cctx);
+    data->zstd_cctx = nullptr;
+  }
+}
+
+static void
+zstd_compress_one(Data *data, const char *upstream_buffer, int64_t upstream_length)
+{
+  char output_buffer[ZSTD_CStreamOutSize()];
+  size_t compressed_size = ZSTD_compressCCtx(data->zstd_cctx, output_buffer, sizeof(output_buffer), upstream_buffer, upstream_length, ZSTD_COMPRESSION_LEVEL);
+  if (ZSTD_isError(compressed_size)) {
+    error("Zstd compression failed: %s", ZSTD_getErrorName(compressed_size));
+    return;
+  }
+  TSIOBufferWrite(data->downstream_buffer, output_buffer, compressed_size);
+  data->downstream_length += compressed_size;
+}
+#endif
+
 static void
 compress_transform_one(Data *data, TSIOBufferReader upstream_reader, int amount)
 {
@@ -488,6 +556,11 @@ compress_transform_one(Data *data, TSIOBufferReader upstream_reader, int amount)
     if (data->compression_type & COMPRESSION_TYPE_BROTLI && (data->compression_algorithms & ALGORITHM_BROTLI)) {
       brotli_transform_one(data, upstream_buffer, upstream_length);
     } else
+#endif
+#if HAVE_ZSTD_H
+    if (data->compression_type & COMPRESSION_TYPE_ZSTD && (data->compression_algorithms & ALGORITHM_ZSTD)) {
+      zstd_compress_one(data, upstream_buffer, upstream_length);
+    } else
 #endif
       if ((data->compression_type & (COMPRESSION_TYPE_GZIP | COMPRESSION_TYPE_DEFLATE)) &&
           (data->compression_algorithms & (ALGORITHM_GZIP | ALGORITHM_DEFLATE))) {
@@ -576,6 +649,12 @@ compress_transform_finish(Data *data)
     brotli_transform_finish(data);
     debug("compress_transform_finish: brotli compression finish");
   } else
+#endif
+#if HAVE_ZSTD_H
+  if (data->compression_type & COMPRESSION_TYPE_ZSTD && data->compression_algorithms & ALGORITHM_ZSTD) {
+    zstd_compress_finish(data);
+    debug("compress_transform_finish: zstd compression finish");
+  } else
 #endif
     if ((data->compression_type & (COMPRESSION_TYPE_GZIP | COMPRESSION_TYPE_DEFLATE)) &&
         (data->compression_algorithms & (ALGORITHM_GZIP | ALGORITHM_DEFLATE))) {
@@ -771,6 +850,11 @@ transformable(TSHttpTxn txnp, bool server, HostConfiguration *host_configuration
           compression_acceptable = 1;
         }
         *compress_type |= COMPRESSION_TYPE_GZIP;
+      } else if (strncasecmp(value, "zstd", sizeof("zstd") - 1) == 0) {
+        if (*algorithms & ALGORITHM_ZSTD) {
+          compression_acceptable = 1;
+        }
+        *compress_type |= COMPRESSION_TYPE_ZSTD;
       }
     }
 

plugins/compress/configuration.cc

@@ -229,6 +229,12 @@ HostConfiguration::add_compression_algorithms(string &line)
     string token = extractFirstToken(line, isCommaOrSpace);
     if (token.empty()) {
       break;
+    } else if (token == "zstd") {
+#ifdef HAVE_ZSTD_H
+      compression_algorithms_ |= ALGORITHM_ZSTD;
+#else
+      error("supported-algorithms: zstd support not compiled in.");
+#endif
     } else if (token == "br") {
 #ifdef HAVE_BROTLI_ENCODE_H
       compression_algorithms_ |= ALGORITHM_BROTLI;
@@ -334,6 +340,10 @@ Configuration::Parse(const char *path)
       continue;
     }
 
+    if (strstr(line.c_str(), "zstd")) {
+      current_host_configuration->add_compression_algorithms(line);
+    }
+
     for (;;) {
       string token = extractFirstToken(line, isspace);
 

plugins/compress/configuration.h

@@ -38,7 +38,8 @@ enum CompressionAlgorithm {
   ALGORITHM_DEFAULT = 0,
   ALGORITHM_DEFLATE = 1,
   ALGORITHM_GZIP    = 2,
-  ALGORITHM_BROTLI  = 4 // For bit manipulations
+  ALGORITHM_BROTLI  = 4, // For bit manipulations
+  ALGORITHM_ZSTD    = 8
 };
 
 enum class RangeRequestCtrl : int {

plugins/compress/misc.cc

@@ -84,6 +84,7 @@ normalize_accept_encoding(TSHttpTxn /* txnp ATS_UNUSED */, TSMBuffer reqp, TSMLo
   bool   deflate = false;
   bool   gzip    = false;
   bool   br      = false;
+  bool   zstd    = false;
   // remove the accept encoding field(s),
   // while finding out if gzip or deflate is supported.
   while (field) {
@@ -100,6 +101,8 @@ normalize_accept_encoding(TSHttpTxn /* txnp ATS_UNUSED */, TSMBuffer reqp, TSMLo
           br = true;
         } else if (strcasecmp("deflate", next) == 0) {
           deflate = true;
+        } else if (strcasecmp("zstd", next) == 0) {
+          zstd = true;
         }
       }
     }
@@ -111,9 +114,13 @@ normalize_accept_encoding(TSHttpTxn /* txnp ATS_UNUSED */, TSMBuffer reqp, TSMLo
   }
 
   // append a new accept-encoding field in the header
-  if (deflate || gzip || br) {
+  if (deflate || gzip || br || zstd) {
     TSMimeHdrFieldCreate(reqp, hdr_loc, &field);
     TSMimeHdrFieldNameSet(reqp, hdr_loc, field, TS_MIME_FIELD_ACCEPT_ENCODING, TS_MIME_LEN_ACCEPT_ENCODING);
+    if (zstd) {
+      TSMimeHdrFieldValueStringInsert(reqp, hdr_loc, field, -1, "zstd", strlen("zstd"));
+      info("normalized accept encoding to zstd");
+    }
     if (br) {
       TSMimeHdrFieldValueStringInsert(reqp, hdr_loc, field, -1, "br", strlen("br"));
       info("normalized accept encoding to br");

plugins/compress/misc.h

@@ -44,7 +44,8 @@ enum CompressionType {
   COMPRESSION_TYPE_DEFAULT = 0,
   COMPRESSION_TYPE_DEFLATE = 1,
   COMPRESSION_TYPE_GZIP    = 2,
-  COMPRESSION_TYPE_BROTLI  = 4
+  COMPRESSION_TYPE_BROTLI  = 4,
+  COMPRESSION_TYPE_ZSTD    = 8,
 };
 
 // this one is used to rename the accept encoding header
@@ -84,13 +85,21 @@ using Data = struct {
 #if HAVE_BROTLI_ENCODE_H
   b_stream bstrm;
 #endif
+#if HAVE_ZSTD_H
+  ZSTD_CCtx *zstd_cctx;
+#endif
 };
 
 voidpf      gzip_alloc(voidpf opaque, uInt items, uInt size);
 void        gzip_free(voidpf opaque, voidpf address);
 void        normalize_accept_encoding(TSHttpTxn txnp, TSMBuffer reqp, TSMLoc hdr_loc);
 void        hide_accept_encoding(TSHttpTxn txnp, TSMBuffer reqp, TSMLoc hdr_loc, const char *hidden_header_name);
 void        restore_accept_encoding(TSHttpTxn txnp, TSMBuffer reqp, TSMLoc hdr_loc, const char *hidden_header_name);
+#if HAVE_ZSTD_H
+void        zstd_compress_init(Data *data);
+void        zstd_compress_finish(Data *data);
+void        zstd_compress_one(Data *data, const char *upstream_buffer, int64_t upstream_length);
+#endif
 const char *init_hidden_header_name();
 int         register_plugin();
 void        log_compression_ratio(int64_t in, int64_t out);

tests/gold_tests/pluginTest/compress/compress.test.py

@@ -170,6 +170,14 @@ def get_verify_command(out_path, decrompressor):
     tr.ReturnCode = 0
     tr.Processes.Default.Command = f"diff {out_path} {orig_path}"
 
+    tr = Test.AddTestRun(f'zstd: {i}')
+    tr.Processes.Default.ReturnCode = 0
+    out_path = get_out_path()
+    tr.MakeCurlCommand(curl(ts, i, "zstd", out_path))
+    tr = Test.AddTestRun(f'verify zstd: {i}')
+    tr.ReturnCode = 0
+    tr.Processes.Default.Command = get_verify_command(out_path, "zstd -d")
+
 # Test Accept-Encoding normalization.
 
 tr = Test.AddTestRun()