about summary refs log tree commit diff
diff options
context:
space:
mode:
authorzimbatm <zimbatm@zimbatm.com>2019-08-12T19·03+0200
committerzimbatm <zimbatm@zimbatm.com>2019-08-16T13·05+0200
commit91b00b145f0c50e346d0250168cbcbcba7aa3b40 (patch)
treea29e7a9b454b29cd61dc12204ea4c370a02d1066
parentb7ea98bf3409a29bb6412c6a17a896ba1c1f524a (diff)
libutil: add SizedSource
Introduce the SizeSource which allows to bound how much data is being
read from a source. It also contains a drainAll() function to discard
the rest of the source, useful to keep the nix protocol in sync.
-rw-r--r--src/libutil/serialise.hh30
1 files changed, 30 insertions, 0 deletions
diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index 969e4dff383d..a344a5ac7520 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -179,6 +179,36 @@ struct TeeSource : Source
     }
 };
 
+/* A reader that consumes the original Source until 'size'. */
+struct SizedSource : Source
+{
+    Source & orig;
+    size_t remain;
+    SizedSource(Source & orig, size_t size)
+        : orig(orig), remain(size) { }
+    size_t read(unsigned char * data, size_t len)
+    {
+        if (this->remain <= 0) {
+            throw EndOfFile("sized: unexpected end-of-file");
+        }
+        len = std::min(len, this->remain);
+        size_t n = this->orig.read(data, len);
+        this->remain -= n;
+        return n;
+    }
+
+    /* Consume the original source until no remain data is left to consume. */
+    size_t drainAll()
+    {
+        std::vector<unsigned char> buf(8192);
+        size_t sum = 0;
+        while (this->remain > 0) {
+            size_t n = read(buf.data(), buf.size());
+            sum += n;
+        }
+        return sum;
+    }
+};
 
 /* Convert a function into a sink. */
 struct LambdaSink : Sink