From 67f778670c085615470b67eb3c54885b8c2d482e Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Fri, 3 May 2024 22:23:02 +0200
Subject: [PATCH] libutil: add makeDecompressionSource

Change-Id: Iac7f24d79e24417436b9b5cbefd6af051aeea0a6
---
 src/libutil/compression.cc        | 21 ++++++++++++++++-----
 src/libutil/compression.hh        |  1 +
 tests/unit/libutil/compression.cc | 11 ++++++-----
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/libutil/compression.cc b/src/libutil/compression.cc
index e78d76500..d93a1b1d6 100644
--- a/src/libutil/compression.cc
+++ b/src/libutil/compression.cc
@@ -192,11 +192,9 @@ struct BrotliDecompressionSource : Source
 
 std::string decompress(const std::string & method, std::string_view in)
 {
-    StringSink ssink;
-    auto sink = makeDecompressionSink(method, ssink);
-    (*sink)(in);
-    sink->finish();
-    return std::move(ssink.s);
+    StringSource src{in};
+    auto filter = makeDecompressionSource(method, src);
+    return filter->drain();
 }
 
 std::unique_ptr<FinishSink> makeDecompressionSink(const std::string & method, Sink & nextSink)
@@ -224,6 +222,19 @@ std::unique_ptr<FinishSink> makeDecompressionSink(const std::string & method, Si
         });
 }
 
+std::unique_ptr<Source> makeDecompressionSource(const std::string & method, Source & inner)
+{
+    if (method == "none" || method == "") {
+        return std::make_unique<LambdaSource>([&](char * data, size_t len) {
+            return inner.read(data, len);
+        });
+    } else if (method == "br") {
+        return std::make_unique<BrotliDecompressionSource>(inner);
+    } else {
+        return std::make_unique<ArchiveDecompressionSource>(inner);
+    }
+}
+
 struct BrotliCompressionSink : ChunkedCompressionSink
 {
     Sink & nextSink;
diff --git a/src/libutil/compression.hh b/src/libutil/compression.hh
index 4e53a7b3c..8affdddd6 100644
--- a/src/libutil/compression.hh
+++ b/src/libutil/compression.hh
@@ -19,6 +19,7 @@ struct CompressionSink : BufferedSink, FinishSink
 std::string decompress(const std::string & method, std::string_view in);
 
 std::unique_ptr<FinishSink> makeDecompressionSink(const std::string & method, Sink & nextSink);
+std::unique_ptr<Source> makeDecompressionSource(const std::string & method, Source & inner);
 
 std::string compress(const std::string & method, std::string_view in, const bool parallel = false, int level = -1);
 
diff --git a/tests/unit/libutil/compression.cc b/tests/unit/libutil/compression.cc
index bbbf3500f..0542e7d33 100644
--- a/tests/unit/libutil/compression.cc
+++ b/tests/unit/libutil/compression.cc
@@ -81,16 +81,17 @@ namespace nix {
     }
 
     TEST(makeCompressionSink, compressAndDecompress) {
-        StringSink strSink;
         auto inputString = "slfja;sljfklsa;jfklsjfkl;sdjfkl;sadjfkl;sdjf;lsdfjsadlf";
-        auto decompressionSink = makeDecompressionSink("bzip2", strSink);
-        auto sink = makeCompressionSink("bzip2", *decompressionSink);
 
+        StringSink strSink;
+        auto sink = makeCompressionSink("bzip2", strSink);
         (*sink)(inputString);
         sink->finish();
-        decompressionSink->finish();
 
-        ASSERT_STREQ(strSink.s.c_str(), inputString);
+        StringSource strSource{strSink.s};
+        auto decompressionSource = makeDecompressionSource("bzip2", strSource);
+
+        ASSERT_STREQ(decompressionSource->drain().c_str(), inputString);
     }
 
 }