From 27d5209f4dcb2da435cd13fd5aba0caa46b9ef3d Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Tue, 18 Mar 2025 14:51:04 +0100
Subject: [PATCH] libutil: fix copyNAR not reading the whole nar when dropped
 early

if a copyNAR generator was not drained to completion it would not read
the full nar data from its source. this could happen if the copier was
passed to parseAndDump wrapped as a source because copyNAR would yield
nar metadata *before* it had read it, and GeneratorSource will drain a
generator fully *only* if the source is allowed to throw EndOfFile. in
the parseAndDump case this never happened because parseAndDump expects
to be given an unterminated stream, and thus the combination left some
nar metadata in the input Source, breaking the remote store protocols.

fixes #732

Change-Id: Ia59a53375992bfcdb7bc6b37764ca779622bc8f7
---
 lix/libutil/archive.cc        | 23 ++++++++-----
 tests/unit/libutil/archive.cc | 63 ++++++++++++++++++++++++++++++++---
 2 files changed, 74 insertions(+), 12 deletions(-)

diff --git a/lix/libutil/archive.cc b/lix/libutil/archive.cc
index 7f620c87a..bc81109c2 100644
--- a/lix/libutil/archive.cc
+++ b/lix/libutil/archive.cc
@@ -56,7 +56,6 @@ static WireFormatGenerator dumpContents(Path path, off_t size)
 
 static WireFormatGenerator dumpSingle(nar::File f)
 {
-    co_yield "(";
     co_yield "type";
     co_yield "regular";
     if (f.executable) {
@@ -67,22 +66,18 @@ static WireFormatGenerator dumpSingle(nar::File f)
     co_yield f.size;
     co_yield std::move(f.contents);
     co_yield SerializingTransform::padding(f.size);
-    co_yield ")";
 }
 
 static WireFormatGenerator dumpSingle(nar::Symlink s)
 {
-    co_yield "(";
     co_yield "type";
     co_yield "symlink";
     co_yield "target";
     co_yield s.target;
-    co_yield ")";
 }
 
 static WireFormatGenerator dumpSingle(nar::Directory d)
 {
-    co_yield "(";
     co_yield "type";
     co_yield "directory";
     while (auto e = d.contents.next()) {
@@ -94,22 +89,25 @@ static WireFormatGenerator dumpSingle(nar::Directory d)
                     co_yield "name";
                     co_yield name;
                     co_yield "node";
+                    co_yield "(";
                     co_yield dumpSingle(std::move(i));
                     co_yield ")";
+                    co_yield ")";
                 }(e->first, i);
             },
             e->second
         );
     }
-    co_yield ")";
 }
 
 WireFormatGenerator nar::dump(nar::Entry nar)
 {
     co_yield narVersionMagic1;
+    co_yield "(";
     co_yield std::visit(
         [](auto i) -> WireFormatGenerator { return dumpSingle(std::move(i)); }, std::move(nar)
     );
+    co_yield ")";
 }
 
 // list the given path under the given filter and return the oldest mtime.
@@ -1077,8 +1075,17 @@ WireFormatGenerator copyNAR(Source & source)
     // we should just forward all data directly without parsing.
 
     auto items = nar::parse(source);
-    co_yield dump(*items.next());
-    assert(!items.next().has_value());
+
+    // we can't use dump() here because we must read the entire nar *before*
+    // returning the final `)` tag, otherwise the source will not be emptied
+    // before the returned generator is exhausted. that in turn confuses the
+    // remote store protocols that expect copyNAR to not finish any earlier.
+    co_yield narVersionMagic1;
+    co_yield "(";
+    for (auto && item : items) {
+        co_yield std::visit([](auto i) { return dumpSingle(std::move(i)); }, std::move(item));
+    }
+    co_yield ")";
 }
 
 box_ptr<AsyncInputStream> copyNAR(AsyncInputStream & source)
diff --git a/tests/unit/libutil/archive.cc b/tests/unit/libutil/archive.cc
index 44b5675a3..51189d092 100644
--- a/tests/unit/libutil/archive.cc
+++ b/tests/unit/libutil/archive.cc
@@ -243,8 +243,7 @@ TEST_P(NarTest, parse)
     }
 }
 
-TEST_P(NarTest, parseAsync)
-{
+namespace parseAsync {
     struct File
     {
         bool executable;
@@ -316,10 +315,15 @@ TEST_P(NarTest, parseAsync)
             parent.emplace(name, Symlink{target});
         }
     };
+}
+
+TEST_P(NarTest, parseAsync)
+{
+    using namespace parseAsync;
 
     AsyncGeneratorInputStream source(rawStream());
 
-    std::map<std::string, Entry> contents;
+    std::map<std::string, parseAsync::Entry> contents;
     ReconstructVisitor rv{contents};
 
     kj::EventLoop el;
@@ -327,7 +331,7 @@ TEST_P(NarTest, parseAsync)
 
     auto entries = entriesFn()();
     parseDump(rv, source).wait(ws).value();
-    auto parsed = Directory::toNar(contents.at(""));
+    auto parsed = parseAsync::Directory::toNar(contents.at(""));
     while (true) {
         auto e = entries.next();
         auto p = parsed.next();
@@ -357,6 +361,57 @@ TEST_P(NarTest, copyAsync)
     ASSERT_EQ(raw(), copied);
 }
 
+TEST_P(NarTest, parseCopied)
+{
+    GeneratorSource input(rawStream());
+    GeneratorSource source(copyNAR(input));
+
+    auto entries = entriesFn()();
+    auto parsed = parse(source);
+    while (true) {
+        auto e = entries.next();
+        auto p = parsed.next();
+        ASSERT_EQ(e.has_value(), p.has_value());
+        if (!e) {
+            break;
+        }
+        assert_eq(*e, *p);
+    }
+
+    char buf;
+    ASSERT_THROW(input(&buf, 1), EndOfFile);
+}
+
+TEST_P(NarTest, parseCopiedAsync)
+{
+    using namespace parseAsync;
+
+    AsyncGeneratorInputStream input(rawStream());
+    auto source = copyNAR(input);
+
+    std::map<std::string, parseAsync::Entry> contents;
+    ReconstructVisitor rv{contents};
+
+    kj::EventLoop el;
+    kj::WaitScope ws{el};
+
+    auto entries = entriesFn()();
+    parseDump(rv, *source).wait(ws).value();
+    auto parsed = parseAsync::Directory::toNar(contents.at(""));
+    while (true) {
+        auto e = entries.next();
+        auto p = parsed.next();
+        ASSERT_EQ(e.has_value(), p.has_value());
+        if (!e) {
+            break;
+        }
+        assert_eq(*e, *p);
+    }
+
+    char buf;
+    ASSERT_EQ(input.read(&buf, 1).wait(ws).value(), 0);
+}
+
 TEST_P(NarTest, index)
 {
     GeneratorSource source(rawStream());