Merge pull request #15399 from NixOS/fixed-length-drain-into

Introduce `Source::drainInto` with explicit length

src/libutil/archive.cc

@@ -135,18 +135,7 @@ static void parseContents(CreateRegularFileSink & sink, Source & source)
         return;
     }
 
-    uint64_t left = size;
-    std::array<char, 65536> buf;
-
-    while (left) {
-        checkInterrupt();
-        auto n = buf.size();
-        if ((uint64_t) n > left)
-            n = left;
-        source(buf.data(), n);
-        sink({buf.data(), n});
-        left -= n;
-    }
+    source.drainInto(sink, size);
 
     readPadding(size, source);
 }

src/libutil/git.cc

@@ -68,17 +68,7 @@ void parseBlob(
 
             crf.preallocateContents(size);
 
-            unsigned long long left = size;
-            std::string buf;
-            buf.reserve(65536);
-
-            while (left) {
-                checkInterrupt();
-                buf.resize(std::min((unsigned long long) buf.capacity(), left));
-                source(buf);
-                crf(buf);
-                left -= buf.size();
-            }
+            source.drainInto(crf, size);
         });
     };
 

src/libutil/include/nix/util/serialise.hh

@@ -96,8 +96,20 @@ struct Source
         return true;
     }
 
+    /**
+     * Read the rest of this `Source` into `sink`.
+     */
     void drainInto(Sink & sink);
 
+    /**
+     * Read exactly 'len' bytes and write them to 'sink'.
+     *
+     * Virtual in anticipation that some `Source` implementations, like
+     * `FdSource` may eventually be able to provide more performant
+     * implementations of this function.
+     */
+    virtual void drainInto(Sink & sink, uint64_t len);
+
     std::string drain();
 
     virtual void skip(size_t len);

src/libutil/serialise.cc

@@ -7,6 +7,7 @@
 
 #include <cstring>
 #include <cerrno>
+#include <limits>
 #include <memory>
 
 #include <boost/coroutine2/coroutine.hpp>
@@ -104,6 +105,20 @@ void Source::drainInto(Sink & sink)
     }
 }
 
+void Source::drainInto(Sink & sink, uint64_t len)
+{
+    std::array<char, 65536> buf;
+    while (len) {
+        checkInterrupt();
+        // Until std::saturate_cast is available (C++26)
+        auto lenTrunc = static_cast<size_t>(std::min<uint64_t>(len, std::numeric_limits<size_t>::max()));
+        auto n = read(buf.data(), std::min(lenTrunc, buf.size()));
+        sink({buf.data(), n});
+        assert(n <= len);
+        len -= n;
+    }
+}
+
 std::string Source::drain()
 {
     StringSink s;