3 files changed, 48 insertions, 5 deletions
diff --git a/src/libutil/archive.hh b/src/libutil/archive.hh
index 302b1bb18..57780d16a 100644
--- a/src/libutil/archive.hh
+++ b/src/libutil/archive.hh
@@ -63,12 +63,29 @@ struct ParseSink
     virtual void createSymlink(const Path & path, const string & target) { };
 };
 
-struct TeeParseSink : ParseSink
+/* If the NAR archive contains a single file at top-level, then save
+   the contents of the file to `s'.  Otherwise barf. */
+struct RetrieveRegularNARSink : ParseSink
 {
-    StringSink saved;
-    TeeSource source;
+    bool regular = true;
+    Sink & sink;
 
-    TeeParseSink(Source & source) : source(source, saved) { }
+    RetrieveRegularNARSink(Sink & sink) : sink(sink) { }
+
+    void createDirectory(const Path & path)
+    {
+        regular = false;
+    }
+
+    void receiveContents(unsigned char * data, unsigned int len)
+    {
+        sink(data, len);
+    }
+
+    void createSymlink(const Path & path, const string & target)
+    {
+        regular = false;
+    }
 };
 
 void parseDump(ParseSink & sink, Source & source);
diff --git a/src/libutil/serialise.cc b/src/libutil/serialise.cc
index c8b71188f..00c945113 100644
--- a/src/libutil/serialise.cc
+++ b/src/libutil/serialise.cc
@@ -322,5 +322,18 @@ void StringSink::operator () (const unsigned char * data, size_t len)
     s->append((const char *) data, len);
 }
 
+size_t ChainSource::read(unsigned char * data, size_t len)
+{
+    if (useSecond) {
+        return source2.read(data, len);
+    } else {
+        try {
+            return source1.read(data, len);
+        } catch (EndOfFile &) {
+            useSecond = true;
+            return this->read(data, len);
+        }
+    }
+}
 
 }
diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index 8386a4991..5d9acf887 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -189,7 +189,7 @@ struct TeeSource : Source
     size_t read(unsigned char * data, size_t len)
     {
         size_t n = orig.read(data, len);
-        sink(data, len);
+        sink(data, n);
         return n;
     }
 };
@@ -256,6 +256,19 @@ struct LambdaSource : Source
     }
 };
 
+/* Chain two sources together so after the first is exhausted, the second is
+   used */
+struct ChainSource : Source
+{
+    Source & source1, & source2;
+    bool useSecond = false;
+    ChainSource(Source & s1, Source & s2)
+        : source1(s1), source2(s2)
+    { }
+
+    size_t read(unsigned char * data, size_t len) override;
+};
+
 
 /* Convert a function that feeds data into a Sink into a Source. The
    Source executes the function as a coroutine. */