From 7f8790eff2313af5a3014bbedf9d2c758b74c9da Mon Sep 17 00:00:00 2001
From: edef <edef@edef.eu>
Date: Mon, 23 Oct 2023 17:36:13 +0000
Subject: [PATCH] libutil: add ZstdDecompressionSink

---
 Makefile.config.in         |  1 +
 configure.ac               |  2 +
 nix.spec.in                |  1 +
 release-common.nix         |  2 +-
 src/libutil/compression.cc | 76 ++++++++++++++++++++++++++++++++++++++
 src/libutil/local.mk       |  2 +-
 6 files changed, 82 insertions(+), 2 deletions(-)

diff --git a/Makefile.config.in b/Makefile.config.in
index 588669f0ac2..5f024417692 100644
--- a/Makefile.config.in
+++ b/Makefile.config.in
@@ -19,6 +19,7 @@ SODIUM_LIBS = @SODIUM_LIBS@
 LIBLZMA_LIBS = @LIBLZMA_LIBS@
 SQLITE3_LIBS = @SQLITE3_LIBS@
 LIBBROTLI_LIBS = @LIBBROTLI_LIBS@
+LIBZSTD_LIBS = @LIBZSTD_LIBS@
 EDITLINE_LIBS = @EDITLINE_LIBS@
 bash = @bash@
 bindir = @bindir@
diff --git a/configure.ac b/configure.ac
index 824a838a04d..b0b620a0863 100644
--- a/configure.ac
+++ b/configure.ac
@@ -223,6 +223,8 @@ AC_CHECK_LIB([lzma], [lzma_stream_encoder_mt],
 # Look for libbrotli{enc,dec}.
 PKG_CHECK_MODULES([LIBBROTLI], [libbrotlienc libbrotlidec], [CXXFLAGS="$LIBBROTLI_CFLAGS $CXXFLAGS"])
 
+# Look for libzstd.
+PKG_CHECK_MODULES([LIBZSTD], [libzstd], [CXXFLAGS="$LIBZSTD_CFLAGS $CXXFLAGS"])
 
 # Look for libseccomp, required for Linux sandboxing.
 case "$host_os" in
diff --git a/nix.spec.in b/nix.spec.in
index 6b9e3763738..1aeafa1746c 100644
--- a/nix.spec.in
+++ b/nix.spec.in
@@ -28,6 +28,7 @@ Requires: curl
 Requires: bzip2
 Requires: gzip
 Requires: xz
+Requires: zstd
 BuildRequires: bison
 BuildRequires: boost-devel >= 1.60
 BuildRequires: bzip2-devel
diff --git a/release-common.nix b/release-common.nix
index a33b71f9311..26355ba0b76 100644
--- a/release-common.nix
+++ b/release-common.nix
@@ -49,7 +49,7 @@ rec {
 
   buildDeps =
     [ curl
-      bzip2 xz brotli editline
+      bzip2 xz brotli zstd editline
       openssl pkgconfig sqlite
       boost
 
diff --git a/src/libutil/compression.cc b/src/libutil/compression.cc
index 0dd84e32034..c168e8ac5a5 100644
--- a/src/libutil/compression.cc
+++ b/src/libutil/compression.cc
@@ -5,8 +5,10 @@
 
 #include <lzma.h>
 #include <bzlib.h>
+#include <zstd.h>
 #include <cstdio>
 #include <cstring>
+#include <vector>
 
 #include <brotli/decode.h>
 #include <brotli/encode.h>
@@ -198,6 +200,78 @@ struct BrotliDecompressionSink : ChunkedCompressionSink
     }
 };
 
+struct ZstdDecompressionSink : CompressionSink
+{
+    Sink & nextSink;
+    ZSTD_DStream *strm;
+
+    std::vector<uint8_t> inbuf;
+    size_t outbuf_size = ZSTD_DStreamOutSize();
+    uint8_t *outbuf = new uint8_t[outbuf_size];
+
+    ZstdDecompressionSink(Sink & nextSink) : nextSink(nextSink)
+    {
+        strm = ZSTD_createDStream();
+        if (!strm)
+            throw CompressionError("unable to initialise zstd decoder");
+
+         ZSTD_initDStream(strm);
+    }
+
+    ~ZstdDecompressionSink()
+    {
+        delete[] outbuf;
+        ZSTD_freeDStream(strm);
+    }
+
+    void finish() override
+    {
+        // this call doesn't make any sense, but it's here for consistency with the other compression sinks
+        // CompressionSink inherits from BufferedSink, but none of the subclasses appear to ever make use of the buffer
+        flush();
+
+        // if we still have undecoded data in the input buffer, we can't signal EOF to libzstd
+        // if we don't, then we're done here anyway
+        if (inbuf.size())
+            throw CompressionError("received unexpected EOF while decompressing zstd file");
+
+        nextSink(nullptr, 0);
+    }
+
+    void write(const unsigned char * data, size_t len) override
+    {
+        inbuf.insert(inbuf.end(), data, data + len);
+
+        ZSTD_inBuffer in = {
+            .src = inbuf.data(),
+            .size = inbuf.size(),
+            .pos = 0
+        };
+
+        ZSTD_outBuffer out = {
+            .dst = outbuf,
+            .size = outbuf_size,
+            .pos = 0
+        };
+
+        while (in.pos < in.size) {
+            out.pos = 0;
+
+            size_t ret = ZSTD_decompressStream(strm, &out, &in);
+            if (ZSTD_isError(ret))
+                throw CompressionError("error %s while decompressing zstd file", ZSTD_getErrorName(ret));
+
+            if (out.pos)
+                nextSink(outbuf, out.pos);
+            else
+                break;
+        }
+
+        // drop consumed input
+        inbuf.erase(inbuf.begin(), inbuf.begin() + in.pos);
+    }
+};
+
 ref<std::string> decompress(const std::string & method, const std::string & in)
 {
     StringSink ssink;
@@ -217,6 +291,8 @@ ref<CompressionSink> makeDecompressionSink(const std::string & method, Sink & ne
         return make_ref<BzipDecompressionSink>(nextSink);
     else if (method == "br")
         return make_ref<BrotliDecompressionSink>(nextSink);
+    else if (method == "zstd")
+        return make_ref<ZstdDecompressionSink>(nextSink);
     else
         throw UnknownCompressionMethod("unknown compression method '%s'", method);
 }
diff --git a/src/libutil/local.mk b/src/libutil/local.mk
index e41a67d1f9e..ebcfddad2b7 100644
--- a/src/libutil/local.mk
+++ b/src/libutil/local.mk
@@ -6,4 +6,4 @@ libutil_DIR := $(d)
 
 libutil_SOURCES := $(wildcard $(d)/*.cc)
 
-libutil_LDFLAGS = $(LIBLZMA_LIBS) -lbz2 -pthread $(OPENSSL_LIBS) $(LIBBROTLI_LIBS) $(BOOST_LDFLAGS) -lboost_context
+libutil_LDFLAGS = $(LIBLZMA_LIBS) -lbz2 -pthread $(OPENSSL_LIBS) $(LIBBROTLI_LIBS) $(LIBZSTD_LIBS) $(BOOST_LDFLAGS) -lboost_context