libutil: add an async nar parser

sadly this is a visitor-only interface; async generators are not yet a
thing and preliminary benchmarks say that overhead would be too large.

Change-Id: I0460d18eba94441cf3101d46bfffdb69cdda81d1

lix/libutil/archive.cc

@@ -15,12 +15,13 @@
 #include "lix/libutil/archive.hh"
 #include "lix/libutil/async-io.hh"
 #include "lix/libutil/box_ptr.hh"
+#include "lix/libutil/config.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"
+#include "lix/libutil/result.hh"
+#include "lix/libutil/serialise.hh"
 #include "lix/libutil/signals.hh"
 
 namespace nix {
@@ -658,6 +659,86 @@ struct AsyncCopier : AsyncInputStream
         }
     }
 };
+
+// sadly async parsers can't be written to produce a tree of generators
+// the way sync parsers can. once we have async generators that may not
+// be as hard, but async generators in kj might have too much overhead.
+struct AsyncParser
+{
+    AsyncInputStream & source;
+    std::vector<char> buffer;
+
+    kj::Promise<Result<void>> parse(NARParseVisitor & target)
+    try {
+        Parser parser{buffer};
+        auto stream = parser.parseRoot();
+        TRY_AWAIT(parse(stream, target, ""));
+        co_return result::success();
+    } catch (...) {
+        co_return result::current_exception();
+    }
+
+    kj::Promise<Result<void>> read(char * buffer, size_t n)
+    try {
+        while (n > 0) {
+            auto got = TRY_AWAIT(source.read(buffer, n));
+            if (got == 0) {
+                throw badArchive("unexpected end of nar encountered");
+            }
+            buffer += got;
+            n -= got;
+        }
+        co_return result::success();
+    } catch (...) {
+        co_return result::current_exception();
+    }
+
+    kj::Promise<Result<void>> feed(size_t n)
+    {
+        auto end = buffer.size();
+        buffer.resize(end + n);
+        return read(buffer.data() + end, n);
+    }
+
+    kj::Promise<Result<void>>
+    parse(Generator<Parser::Response> & stream, NARParseVisitor & target, const std::string & name)
+    try {
+        while (auto i = stream.next()) {
+            if (auto want = std::get_if<Parser::WantBytes>(&*i)) {
+                TRY_AWAIT(feed(want->n));
+            } else if (auto f = std::get_if<Parser::FileHeader>(&*i)) {
+                auto file = target.createRegularFile(name, f->size, f->executable);
+                auto left = f->size;
+                std::array<char, 65536> buf;
+
+                while (left) {
+                    auto n = size_t(std::min<uint64_t>(buf.size(), left));
+                    TRY_AWAIT(read(buf.data(), n));
+                    file->receiveContents({buf.data(), n});
+                    left -= n;
+                }
+            } else if (auto sl = std::get_if<Parser::Symlink>(&*i)) {
+                target.createSymlink(name, sl->target);
+            } else if (auto d = std::get_if<Parser::Directory>(&*i)) {
+                auto dir = target.createDirectory(name);
+                while (auto e = d->content.next()) {
+                    if (auto want = std::get_if<Parser::WantBytes>(&*e)) {
+                        TRY_AWAIT(feed(want->n));
+                    } else if (auto entry = std::get_if<Parser::Directory::Entry>(&*e)) {
+                        TRY_AWAIT(parse(entry->second, *dir, entry->first));
+                    } else {
+                        assert(false && "expected parser response in dir");
+                    }
+                }
+            } else {
+                assert(false && "unhandled parser response");
+            }
+        }
+        co_return result::success();
+    } catch (...) {
+        co_return result::current_exception();
+    }
+};
 }
 
 Generator<Entry> parse(Source & source)
@@ -762,6 +843,15 @@ void parseDump(NARParseVisitor & sink, Source & source)
     }
 }
 
+kj::Promise<Result<void>> parseDump(NARParseVisitor & sink, AsyncInputStream & source)
+try {
+    nar::AsyncParser parser{source};
+    TRY_AWAIT(parser.parse(sink));
+    co_return result::success();
+} catch (...) {
+    co_return result::current_exception();
+}
+
 /*
  * Note [NAR restoration security]:
  * It's *critical* that NAR restoration will never overwrite anything even if

lix/libutil/archive.hh

@@ -204,6 +204,7 @@ Entry create(Source & source);
 }
 
 void parseDump(NARParseVisitor & sink, Source & source);
+kj::Promise<Result<void>> parseDump(NARParseVisitor & sink, AsyncInputStream & source);
 
 void restorePath(const Path & path, Source & source);
 

tests/unit/libutil/archive.cc

@@ -1,5 +1,6 @@
 #include "lix/libutil/archive.hh"
 #include "lix/libutil/async-io.hh"
+#include "lix/libutil/box_ptr.hh"
 #include "lix/libutil/serialise.hh"
 #include <algorithm>
 #include <gtest/gtest.h>
@@ -228,6 +229,103 @@ TEST_P(NarTest, parse)
     }
 }
 
+TEST_P(NarTest, parseAsync)
+{
+    struct File
+    {
+        bool executable;
+        uint64_t size;
+        std::string contents;
+        operator nar::Entry() const
+        {
+            auto stream = [](auto contents) -> Generator<Bytes> { co_yield contents; };
+            return nar::File{executable, size, stream(std::span{contents})};
+        }
+    };
+    struct Directory;
+    using Entry = std::variant<File, Symlink, Directory>;
+
+    struct Directory : std::map<std::string, Entry>
+    {
+        operator nar::Entry() const
+        {
+            return nar::Directory{
+                [](const Directory & d) -> Generator<std::pair<const std::string &, nar::Entry>> {
+                    for (auto & [name, entry] : d) {
+                        co_yield std::pair{std::cref(name), *toNar(entry).next()};
+                    }
+                }(*this),
+            };
+        }
+
+        static Entries toNar(const Entry & e)
+        {
+            co_yield std::visit([](const auto & e) -> Entries { co_yield e; }, e);
+        }
+    };
+
+    struct ReconstructVisitor : NARParseVisitor
+    {
+        std::map<std::string, Entry> & parent;
+
+        struct FileReader : NARParseVisitor::FileHandle
+        {
+            File & file;
+
+            FileReader(File & file) : file(file) {}
+
+            void receiveContents(std::string_view data) override
+            {
+                file.contents += data;
+            }
+
+            void close() override {}
+        };
+
+        explicit ReconstructVisitor(std::map<std::string, Entry> & parent) : parent(parent) {}
+
+        box_ptr<NARParseVisitor> createDirectory(const std::string & name) override
+        {
+            auto & dir = std::get<Directory>(parent.emplace(name, Directory{}).first->second);
+            return make_box_ptr<ReconstructVisitor>(dir);
+        }
+
+        box_ptr<FileHandle>
+        createRegularFile(const std::string & name, uint64_t size, bool executable) override
+        {
+            auto & file = std::get<File>(parent.emplace(name, File{executable, size}).first->second);
+            return make_box_ptr<FileReader>(file);
+        }
+
+        void createSymlink(const std::string & name, const std::string & target) override
+        {
+            parent.emplace(name, Symlink{target});
+        }
+    };
+
+    auto & [raw, entriesF] = GetParam();
+    AsyncStringInputStream source(raw);
+
+    std::map<std::string, Entry> contents;
+    ReconstructVisitor rv{contents};
+
+    kj::EventLoop el;
+    kj::WaitScope ws{el};
+
+    auto entries = entriesF();
+    parseDump(rv, source).wait(ws).value();
+    auto parsed = 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);
+    }
+}
+
 TEST_P(NarTest, copy)
 {
     auto & [raw, _] = GetParam();