From 18a24393e8b6f0da70dafea7fafa350b9d8dca7c Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Mon, 3 Mar 2025 20:48:59 +0100
Subject: [PATCH] libutil: add an async copyNAR overload

Change-Id: I046ff704a2d2cb5115f7bf610feac3ed3247d992
---
 lix/libutil/archive.cc        | 82 +++++++++++++++++++++++++++++++++++
 lix/libutil/archive.hh        |  3 ++
 tests/unit/libutil/archive.cc | 13 ++++++
 3 files changed, 98 insertions(+)

diff --git a/lix/libutil/archive.cc b/lix/libutil/archive.cc
index cddf84943..ef38af6d7 100644
--- a/lix/libutil/archive.cc
+++ b/lix/libutil/archive.cc
@@ -13,8 +13,11 @@
 #include <fcntl.h>
 
 #include "lix/libutil/archive.hh"
+#include "lix/libutil/async-io.hh"
+#include "lix/libutil/box_ptr.hh"
 #include "lix/libutil/file-system.hh"
 #include "lix/libutil/finally.hh"
+#include "lix/libutil/generator.hh"
 #include "lix/libutil/serialise.hh"
 #include "lix/libutil/config.hh"
 #include "lix/libutil/logging.hh"
@@ -581,6 +584,80 @@ struct SyncParser
         }
     }
 };
+
+struct AsyncCopier : AsyncInputStream
+{
+    AsyncInputStream & source;
+    std::vector<char> buffer;
+    Parser parser{buffer};
+
+    struct Fragment
+    {
+        // how many bytes the parser requested but we haven't read from source yet
+        size_t pending = 0;
+        // whether the requested bytes are nar metadata (false) or contents (true)
+        bool pendingFileContents = false;
+    };
+
+    Generator<Fragment> stream{ignoreContents(parser.parseRoot())};
+    Fragment current;
+
+    explicit AsyncCopier(AsyncInputStream & source) : source(source) {}
+
+    kj::Promise<Result<size_t>> read(void * buffer, size_t size) override
+    try {
+        while (current.pending == 0) {
+            if (auto want = stream.next()) {
+                current = *want;
+            } else {
+                co_return 0;
+            }
+        }
+
+        size = std::min(current.pending, size);
+        if (size == 0) {
+            co_return 0;
+        }
+
+        auto got = TRY_AWAIT(source.read(buffer, size));
+        current.pending -= got;
+        if (got == 0) {
+            throw badArchive("truncated NAR encountered");
+        } else if (!current.pendingFileContents) {
+            auto end = this->buffer.size();
+            this->buffer.resize(end + size);
+            memcpy(this->buffer.data() + end, buffer, got);
+        }
+        co_return got;
+    } catch (...) {
+        co_return result::current_exception();
+    }
+
+    Generator<Fragment> ignoreContents(Generator<Parser::Response> stream)
+    {
+        while (auto i = stream.next()) {
+            if (auto want = std::get_if<Parser::WantBytes>(&*i)) {
+                co_yield Fragment{want->n, false};
+            } else if (auto f = std::get_if<Parser::FileHeader>(&*i)) {
+                co_yield Fragment{f->size, true};
+            } else if (auto sl = std::get_if<Parser::Symlink>(&*i)) {
+                // nothing to do
+            } else if (auto dir = std::get_if<Parser::Directory>(&*i)) {
+                while (auto e = dir->content.next()) {
+                    if (auto want = std::get_if<Parser::WantBytes>(&*e)) {
+                        co_yield Fragment{want->n, false};
+                    } else if (auto entry = std::get_if<Parser::Directory::Entry>(&*e)) {
+                        co_yield ignoreContents(std::move(entry->second));
+                    } else {
+                        assert(false && "expected parser response in dir");
+                    }
+                }
+            } else {
+                assert(false && "unhandled parser response");
+            }
+        }
+    }
+};
 }
 
 Generator<Entry> parse(Source & source)
@@ -818,4 +895,9 @@ WireFormatGenerator copyNAR(Source & source)
     assert(!items.next().has_value());
 }
 
+box_ptr<AsyncInputStream> copyNAR(AsyncInputStream & source)
+{
+    return make_box_ptr<nar::AsyncCopier>(source);
+}
+
 }
diff --git a/lix/libutil/archive.hh b/lix/libutil/archive.hh
index 4032b825f..cb72269c9 100644
--- a/lix/libutil/archive.hh
+++ b/lix/libutil/archive.hh
@@ -1,6 +1,8 @@
 #pragma once
 ///@file
 
+#include "lix/libutil/async-io.hh"
+#include "lix/libutil/box_ptr.hh"
 #include "lix/libutil/generator.hh"
 #include "lix/libutil/types.hh"
 #include "lix/libutil/serialise.hh"
@@ -209,6 +211,7 @@ void restorePath(const Path & path, Source & source);
  * Read a NAR from 'source' and return it as a generator.
  */
 WireFormatGenerator copyNAR(Source & source);
+box_ptr<AsyncInputStream> copyNAR(AsyncInputStream & source);
 
 
 inline constexpr std::string_view narVersionMagic1 = "nix-archive-1";
diff --git a/tests/unit/libutil/archive.cc b/tests/unit/libutil/archive.cc
index 89fc6be62..219cdc4a5 100644
--- a/tests/unit/libutil/archive.cc
+++ b/tests/unit/libutil/archive.cc
@@ -1,7 +1,9 @@
 #include "lix/libutil/archive.hh"
+#include "lix/libutil/async-io.hh"
 #include "lix/libutil/serialise.hh"
 #include <algorithm>
 #include <gtest/gtest.h>
+#include <kj/async.h>
 
 using namespace std::literals;
 
@@ -235,6 +237,17 @@ TEST_P(NarTest, copy)
     ASSERT_EQ(raw, copied);
 }
 
+TEST_P(NarTest, copyAsync)
+{
+    auto & [raw, _] = GetParam();
+    AsyncStringInputStream source(raw);
+
+    kj::EventLoop el;
+    kj::WaitScope ws{el};
+    auto copied = copyNAR(source)->drain().wait(ws).value();
+    ASSERT_EQ(raw, copied);
+}
+
 TEST_P(NarTest, index)
 {
     auto & [raw, entriesF] = GetParam();