3 files changed, 102 insertions, 4 deletions

diff --git a/src/libutil/serialise.cc b/src/libutil/serialise.cc
index 00c945113..a469a1e73 100644
--- a/src/libutil/serialise.cc
+++ b/src/libutil/serialise.cc
@@ -93,7 +93,7 @@ void Source::operator () (unsigned char * data, size_t len)
 }
 
 
-std::string Source::drain()
+void Source::drainInto(Sink & sink)
 {
     std::string s;
     std::vector<unsigned char> buf(8192);
@@ -101,12 +101,19 @@ std::string Source::drain()
         size_t n;
         try {
             n = read(buf.data(), buf.size());
-            s.append((char *) buf.data(), n);
+            sink(buf.data(), n);
         } catch (EndOfFile &) {
             break;
         }
     }
-    return s;
+}
+
+
+std::string Source::drain()
+{
+    StringSink s;
+    drainInto(s);
+    return *s.s;
 }
 
 
diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index 7682a0f19..b41e58f33 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -69,6 +69,8 @@ struct Source
 
     virtual bool good() { return true; }
 
+    void drainInto(Sink & sink);
+
     std::string drain();
 };
 
@@ -404,4 +406,93 @@ struct StreamToSourceAdapter : Source
 };
 
 
+/* A source that reads a distinct format of concatenated chunks back into its
+   logical form, in order to guarantee a known state to the original stream,
+   even in the event of errors.
+
+   Use with FramedSink, which also allows the logical stream to be terminated
+   in the event of an exception.
+*/
+struct FramedSource : Source
+{
+    Source & from;
+    bool eof = false;
+    std::vector<unsigned char> pending;
+    size_t pos = 0;
+
+    FramedSource(Source & from) : from(from)
+    { }
+
+    ~FramedSource()
+    {
+        if (!eof) {
+            while (true) {
+                auto n = readInt(from);
+                if (!n) break;
+                std::vector<unsigned char> data(n);
+                from(data.data(), n);
+            }
+        }
+    }
+
+    size_t read(unsigned char * data, size_t len) override
+    {
+        if (eof) throw EndOfFile("reached end of FramedSource");
+
+        if (pos >= pending.size()) {
+            size_t len = readInt(from);
+            if (!len) {
+                eof = true;
+                return 0;
+            }
+            pending = std::vector<unsigned char>(len);
+            pos = 0;
+            from(pending.data(), len);
+        }
+
+        auto n = std::min(len, pending.size() - pos);
+        memcpy(data, pending.data() + pos, n);
+        pos += n;
+        return n;
+    }
+};
+
+/* Write as chunks in the format expected by FramedSource.
+
+   The exception_ptr reference can be used to terminate the stream when you
+   detect that an error has occurred on the remote end.
+*/
+struct FramedSink : nix::BufferedSink
+{
+    BufferedSink & to;
+    std::exception_ptr & ex;
+
+    FramedSink(BufferedSink & to, std::exception_ptr & ex) : to(to), ex(ex)
+    { }
+
+    ~FramedSink()
+    {
+        try {
+            to << 0;
+            to.flush();
+        } catch (...) {
+            ignoreException();
+        }
+    }
+
+    void write(const unsigned char * data, size_t len) override
+    {
+        /* Don't send more data if the remote has
+            encountered an error. */
+        if (ex) {
+            auto ex2 = ex;
+            ex = nullptr;
+            std::rethrow_exception(ex2);
+        }
+        to << len;
+        to(data, len);
+    };
+};
+
+
 }
diff --git a/src/libutil/split.hh b/src/libutil/split.hh
index d19d7d8ed..87a23b13e 100644
--- a/src/libutil/split.hh
+++ b/src/libutil/split.hh
@@ -9,7 +9,7 @@ namespace nix {
 
 // If `separator` is found, we return the portion of the string before the
 // separator, and modify the string argument to contain only the part after the
-// separator. Otherwise, wer return `std::nullopt`, and we leave the argument
+// separator. Otherwise, we return `std::nullopt`, and we leave the argument
 // string alone.
 static inline std::optional<std::string_view> splitPrefixTo(std::string_view & string, char separator) {
     auto sepInstance = string.find(separator);