From 7d21863bb3f3d4c42b273a8648a65bec83fe0b60 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Wed, 28 Mar 2018 13:32:44 +0200
Subject: [PATCH] Make <nix/fetchurl.nix> run in constant memory

E.g.

  nix-build --store ~/my-nix/ -E 'import <nix/fetchurl.nix> { url = https://cache.nixos.org/nar/0nwi996rgq4b914qyx0mv2wq4k80hjac7xilikavagw7kxmn2iiv.nar.xz; sha256 = "0nwi996rgq4b914qyx0mv2wq4k80hjac7xilikavagw7kxmn2iiv"; }'

now runs in 17 MiB (was 70 MiB), while

  nix-build --store ~/my-nix/ -E 'import <nix/fetchurl.nix> { url = https://cache.nixos.org/nar/0nwi996rgq4b914qyx0mv2wq4k80hjac7xilikavagw7kxmn2iiv.nar.xz; sha256 = "0d2fxljdih3nc5dqx41hjzic3141ajil94m8kdbpryq569dpsbvb"; unpack = true; }'

runs in 17 MiB (was 346 MiB).
---
 src/libstore/builtins/fetchurl.cc | 74 +++++++++++++++++--------------
 src/libutil/util.cc               | 17 +++++++
 src/libutil/util.hh               |  2 +
 3 files changed, 60 insertions(+), 33 deletions(-)

diff --git a/src/libstore/builtins/fetchurl.cc b/src/libstore/builtins/fetchurl.cc
index 4ca4a838e..1f4abd374 100644
--- a/src/libstore/builtins/fetchurl.cc
+++ b/src/libstore/builtins/fetchurl.cc
@@ -22,52 +22,60 @@ void builtinFetchurl(const BasicDerivation & drv, const std::string & netrcData)
         return i->second;
     };
 
-    auto fetch = [&](const string & url) {
-        /* No need to do TLS verification, because we check the hash of
-           the result anyway. */
-        DownloadRequest request(url);
-        request.verifyTLS = false;
-        request.decompress = false;
+    Path storePath = getAttr("out");
+    auto mainUrl = getAttr("url");
 
-        /* Note: have to use a fresh downloader here because we're in
-           a forked process. */
-        auto data = makeDownloader()->download(request);
-        assert(data.data);
+    /* Note: have to use a fresh downloader here because we're in
+       a forked process. */
+    auto downloader = makeDownloader();
 
-        return data.data;
+    auto fetch = [&](const std::string & url) {
+
+        auto source = sinkToSource([&](Sink & sink) {
+
+            /* No need to do TLS verification, because we check the hash of
+               the result anyway. */
+            DownloadRequest request(url);
+            request.verifyTLS = false;
+            request.decompress = false;
+
+            downloader->download(std::move(request), sink);
+        });
+
+        if (get(drv.env, "unpack", "") == "1") {
+
+            if (hasSuffix(mainUrl, ".xz")) {
+                auto source2 = sinkToSource([&](Sink & sink) {
+                    decompress("xz", *source, sink);
+                });
+                restorePath(storePath, *source2);
+            } else
+                restorePath(storePath, *source);
+
+        } else
+              writeFile(storePath, *source);
+
+        auto executable = drv.env.find("executable");
+        if (executable != drv.env.end() && executable->second == "1") {
+            if (chmod(storePath.c_str(), 0755) == -1)
+                throw SysError(format("making '%1%' executable") % storePath);
+        }
     };
 
-    std::shared_ptr<std::string> data;
-
+    /* Try the hashed mirrors first. */
     if (getAttr("outputHashMode") == "flat")
         for (auto hashedMirror : settings.hashedMirrors.get())
             try {
                 if (!hasSuffix(hashedMirror, "/")) hashedMirror += '/';
                 auto ht = parseHashType(getAttr("outputHashAlgo"));
-                data = fetch(hashedMirror + printHashType(ht) + "/" + Hash(getAttr("outputHash"), ht).to_string(Base16, false));
-                break;
+                fetch(hashedMirror + printHashType(ht) + "/" + Hash(getAttr("outputHash"), ht).to_string(Base16, false));
+                return;
             } catch (Error & e) {
                 debug(e.what());
             }
 
-    if (!data) data = fetch(getAttr("url"));
-
-    Path storePath = getAttr("out");
-
-    auto unpack = drv.env.find("unpack");
-    if (unpack != drv.env.end() && unpack->second == "1") {
-        if (string(*data, 0, 6) == string("\xfd" "7zXZ\0", 6))
-            data = decompress("xz", *data);
-        StringSource source(*data);
-        restorePath(storePath, source);
-    } else
-        writeFile(storePath, *data);
-
-    auto executable = drv.env.find("executable");
-    if (executable != drv.env.end() && executable->second == "1") {
-        if (chmod(storePath.c_str(), 0755) == -1)
-            throw SysError(format("making '%1%' executable") % storePath);
-    }
+    /* Otherwise try the specified URL. */
+    fetch(mainUrl);
 }
 
 }
diff --git a/src/libutil/util.cc b/src/libutil/util.cc
index 22753f4b7..6bc64ae75 100644
--- a/src/libutil/util.cc
+++ b/src/libutil/util.cc
@@ -328,6 +328,23 @@ void writeFile(const Path & path, const string & s, mode_t mode)
 }
 
 
+void writeFile(const Path & path, Source & source, mode_t mode)
+{
+    AutoCloseFD fd = open(path.c_str(), O_WRONLY | O_TRUNC | O_CREAT | O_CLOEXEC, mode);
+    if (!fd)
+        throw SysError(format("opening file '%1%'") % path);
+
+    std::vector<unsigned char> buf(64 * 1024);
+
+    while (true) {
+        try {
+            auto n = source.read(buf.data(), buf.size());
+            writeFull(fd.get(), (unsigned char *) buf.data(), n);
+        } catch (EndOfFile &) { break; }
+    }
+}
+
+
 string readLine(int fd)
 {
     string s;
diff --git a/src/libutil/util.hh b/src/libutil/util.hh
index af97b76d9..fc25d2775 100644
--- a/src/libutil/util.hh
+++ b/src/libutil/util.hh
@@ -103,6 +103,8 @@ void readFile(const Path & path, Sink & sink);
 /* Write a string to a file. */
 void writeFile(const Path & path, const string & s, mode_t mode = 0666);
 
+void writeFile(const Path & path, Source & source, mode_t mode = 0666);
+
 /* Read a line from a file descriptor. */
 string readLine(int fd);