From a5d431a9119c2560608768d9ec70fa37d5a6a985 Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Wed, 15 May 2024 17:13:45 +0200
Subject: [PATCH] libstore: turn copyNAR into a generator

Change-Id: Id452f6a03faa1037ff13af0f63e32883966ff40d
---
 src/libstore/daemon.cc           | 3 +--
 src/libstore/export-import.cc    | 2 +-
 src/libstore/legacy-ssh-store.cc | 6 +++---
 src/libstore/remote-store.cc     | 4 ++--
 src/libutil/archive.cc           | 8 +++-----
 src/libutil/archive.hh           | 4 ++--
 6 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/src/libstore/daemon.cc b/src/libstore/daemon.cc
index 52e8ab32e..f7b6a38a1 100644
--- a/src/libstore/daemon.cc
+++ b/src/libstore/daemon.cc
@@ -463,7 +463,7 @@ static void performOp(TunnelLogger * logger, ref<Store> store,
                        command. (We don't trust `addToStoreFromDump` to not
                        eagerly consume the entire stream it's given, past the
                        length of the Nar. */
-                    copyNAR(from, saved);
+                    saved << copyNAR(from);
                 } else {
                     /* Incrementally parse the NAR file, stripping the
                        metadata, and streaming the sole file we expect into
@@ -884,7 +884,6 @@ static void performOp(TunnelLogger * logger, ref<Store> store,
 
         else {
             std::unique_ptr<Source> source;
-            StringSink saved;
             source = std::make_unique<TunnelSource>(from, to);
 
             logger->startWork();
diff --git a/src/libstore/export-import.cc b/src/libstore/export-import.cc
index e0f3eb6c9..a05a9f23b 100644
--- a/src/libstore/export-import.cc
+++ b/src/libstore/export-import.cc
@@ -64,7 +64,7 @@ StorePaths Store::importPaths(Source & source, CheckSigsFlag checkSigs)
 
         /* Extract the NAR from the source. */
         StringSink saved;
-        copyNAR(source, saved);
+        saved << copyNAR(source);
 
         uint32_t magic = readInt(source);
         if (magic != exportMagic)
diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index 584254afe..b4f3854d5 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -193,7 +193,7 @@ struct LegacySSHStore : public virtual LegacySSHStoreConfig, public virtual Stor
                 << info.sigs
                 << renderContentAddress(info.ca);
             try {
-                copyNAR(source, conn->to);
+                conn->to << copyNAR(source);
             } catch (...) {
                 conn->good = false;
                 throw;
@@ -206,7 +206,7 @@ struct LegacySSHStore : public virtual LegacySSHStoreConfig, public virtual Stor
                 << ServeProto::Command::ImportPaths
                 << 1;
             try {
-                copyNAR(source, conn->to);
+                conn->to << copyNAR(source);
             } catch (...) {
                 conn->good = false;
                 throw;
@@ -233,7 +233,7 @@ struct LegacySSHStore : public virtual LegacySSHStoreConfig, public virtual Stor
 
         conn->to << ServeProto::Command::DumpStorePath << printStorePath(path);
         conn->to.flush();
-        copyNAR(conn->from, sink);
+        sink << copyNAR(conn->from);
     }
 
     std::optional<StorePath> queryPathFromHashPart(const std::string & hashPart) override
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index d955f8449..9cebdbe82 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -469,7 +469,7 @@ void RemoteStore::addToStore(const ValidPathInfo & info, Source & source,
 
     if (GET_PROTOCOL_MINOR(conn->daemonVersion) >= 23) {
         conn.withFramedSink([&](Sink & sink) {
-            copyNAR(source, sink);
+            sink << copyNAR(source);
         });
     } else {
         conn.processStderr(0, &source);
@@ -853,7 +853,7 @@ void RemoteStore::narFromPath(const StorePath & path, Sink & sink)
     auto conn(connections->get());
     conn->to << WorkerProto::Op::NarFromPath << printStorePath(path);
     conn->processStderr();
-    copyNAR(conn->from, sink);
+    sink << copyNAR(conn->from);
 }
 
 ref<FSAccessor> RemoteStore::getFSAccessor()
diff --git a/src/libutil/archive.cc b/src/libutil/archive.cc
index 1b8038649..1c82c3f78 100644
--- a/src/libutil/archive.cc
+++ b/src/libutil/archive.cc
@@ -413,16 +413,14 @@ void restorePath(const Path & path, Source & source)
 }
 
 
-void copyNAR(Source & source, Sink & sink)
+WireFormatGenerator copyNAR(Source & source)
 {
     // FIXME: if 'source' is the output of dumpPath() followed by EOF,
     // we should just forward all data directly without parsing.
 
-    ParseSink parseSink; /* null sink; just parse the NAR */
+    static ParseSink parseSink; /* null sink; just parse the NAR */
 
-    TeeSource wrapper { source, sink };
-
-    parseDump(parseSink, wrapper);
+    return parseAndCopyDump(parseSink, source);
 }
 
 
diff --git a/src/libutil/archive.hh b/src/libutil/archive.hh
index 9b4ad9faf..97d99f2f4 100644
--- a/src/libutil/archive.hh
+++ b/src/libutil/archive.hh
@@ -122,9 +122,9 @@ void parseDump(ParseSink & sink, Source & source);
 void restorePath(const Path & path, Source & source);
 
 /**
- * Read a NAR from 'source' and write it to 'sink'.
+ * Read a NAR from 'source' and return it as a generator.
  */
-void copyNAR(Source & source, Sink & sink);
+WireFormatGenerator copyNAR(Source & source);
 
 
 inline constexpr std::string_view narVersionMagic1 = "nix-archive-1";