From 0e8fc118b3d2d3bb6f9b0a918bf8ceb3927774cd Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <eelco.dolstra@logicblox.com>
Date: Wed, 10 Dec 2014 16:35:42 +0100
Subject: [PATCH] Use vfork

---
 src/libutil/util.cc          | 31 +++++++++++++++++++++++--------
 src/libutil/util.hh          | 12 ++++++++++--
 src/nix-daemon/nix-daemon.cc |  7 ++++++-
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/src/libutil/util.cc b/src/libutil/util.cc
index 305e470eb..bdd114c5e 100644
--- a/src/libutil/util.cc
+++ b/src/libutil/util.cc
@@ -863,31 +863,46 @@ void killUser(uid_t uid)
 //////////////////////////////////////////////////////////////////////
 
 
-pid_t startProcess(std::function<void()> fun,
-    bool dieWithParent, const string & errorPrefix, bool runExitHandlers)
+/* Wrapper around vfork to prevent the child process from clobbering
+   the caller's stack frame in the parent. */
+static pid_t doFork(bool allowVfork, std::function<void()> fun) __attribute__((noinline));
+static pid_t doFork(bool allowVfork, std::function<void()> fun)
 {
+#ifdef __linux__
+    pid_t pid = allowVfork ? vfork() : fork();
+#else
     pid_t pid = fork();
-    if (pid == -1) throw SysError("unable to fork");
+#endif
+    if (pid != 0) return pid;
+    fun();
+    abort();
+}
 
-    if (pid == 0) {
+
+pid_t startProcess(std::function<void()> fun, const ProcessOptions & options)
+{
+    auto wrapper = [&]() {
         _writeToStderr = 0;
         try {
 #if __linux__
-            if (dieWithParent && prctl(PR_SET_PDEATHSIG, SIGKILL) == -1)
+            if (options.dieWithParent && prctl(PR_SET_PDEATHSIG, SIGKILL) == -1)
                 throw SysError("setting death signal");
 #endif
             restoreAffinity();
             fun();
         } catch (std::exception & e) {
             try {
-                std::cerr << errorPrefix << e.what() << "\n";
+                std::cerr << options.errorPrefix << e.what() << "\n";
             } catch (...) { }
         } catch (...) { }
-        if (runExitHandlers)
+        if (options.runExitHandlers)
             exit(1);
         else
             _exit(1);
-    }
+    };
+
+    pid_t pid = doFork(options.allowVfork, wrapper);
+    if (pid == -1) throw SysError("unable to fork");
 
     return pid;
 }
diff --git a/src/libutil/util.hh b/src/libutil/util.hh
index 628b8a0e1..09eb4d6e2 100644
--- a/src/libutil/util.hh
+++ b/src/libutil/util.hh
@@ -269,8 +269,16 @@ void killUser(uid_t uid);
 
 /* Fork a process that runs the given function, and return the child
    pid to the caller. */
-pid_t startProcess(std::function<void()> fun, bool dieWithParent = true,
-    const string & errorPrefix = "error: ", bool runExitHandlers = false);
+struct ProcessOptions
+{
+    string errorPrefix;
+    bool dieWithParent;
+    bool runExitHandlers;
+    bool allowVfork;
+    ProcessOptions() : errorPrefix("error: "), dieWithParent(true), runExitHandlers(false), allowVfork(true) { };
+};
+
+pid_t startProcess(std::function<void()> fun, const ProcessOptions & options = ProcessOptions());
 
 
 /* Run a program and return its stdout in a string (i.e., like the
diff --git a/src/nix-daemon/nix-daemon.cc b/src/nix-daemon/nix-daemon.cc
index 3864ab935..9a8f469ad 100644
--- a/src/nix-daemon/nix-daemon.cc
+++ b/src/nix-daemon/nix-daemon.cc
@@ -799,6 +799,11 @@ static void daemonLoop(char * * argv)
                 % (peer.uidKnown ? user : "<unknown>"));
 
             /* Fork a child to handle the connection. */
+            ProcessOptions options;
+            options.errorPrefix = "unexpected Nix daemon error: ";
+            options.dieWithParent = false;
+            options.runExitHandlers = true;
+            options.allowVfork = false;
             startProcess([&]() {
                 fdSocket.close();
 
@@ -821,7 +826,7 @@ static void daemonLoop(char * * argv)
                 processConnection(trusted);
 
                 exit(0);
-            }, false, "unexpected Nix daemon error: ", true);
+            }, options);
 
         } catch (Interrupted & e) {
             throw;