From 906adadacd2d1c98346a2f42c0b42a32d2806d94 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Wed, 7 Apr 2021 13:40:13 +0200
Subject: [PATCH] Restore stack size in child processes

Fixes #4673.
---
 src/libutil/util.cc | 27 +++++++++++++++++++++++++++
 src/libutil/util.hh |  4 ++++
 src/nix/main.cc     | 13 +------------
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/src/libutil/util.cc b/src/libutil/util.cc
index 06821abb1..60b318559 100644
--- a/src/libutil/util.cc
+++ b/src/libutil/util.cc
@@ -32,6 +32,7 @@
 
 #ifdef __linux__
 #include <sys/prctl.h>
+#include <sys/resource.h>
 #endif
 
 
@@ -1618,11 +1619,37 @@ static void restoreSignals()
         throw SysError("restoring signals");
 }
 
+#if __linux__
+rlim_t savedStackSize = 0;
+#endif
+
+void setStackSize(size_t stackSize)
+{
+    #if __linux__
+    struct rlimit limit;
+    if (getrlimit(RLIMIT_STACK, &limit) == 0 && limit.rlim_cur < stackSize) {
+        savedStackSize = limit.rlim_cur;
+        limit.rlim_cur = stackSize;
+        setrlimit(RLIMIT_STACK, &limit);
+    }
+    #endif
+}
+
 void restoreProcessContext()
 {
     restoreSignals();
 
     restoreAffinity();
+
+    #if __linux__
+    if (savedStackSize) {
+        struct rlimit limit;
+        if (getrlimit(RLIMIT_STACK, &limit) == 0) {
+            limit.rlim_cur = savedStackSize;
+            setrlimit(RLIMIT_STACK, &limit);
+        }
+    }
+    #endif
 }
 
 /* RAII helper to automatically deregister a callback. */
diff --git a/src/libutil/util.hh b/src/libutil/util.hh
index 4762d1e55..f84d0fb31 100644
--- a/src/libutil/util.hh
+++ b/src/libutil/util.hh
@@ -300,6 +300,10 @@ std::pair<int, std::string> runProgram(const RunOptions & options);
 void runProgram2(const RunOptions & options);
 
 
+/* Change the stack size. */
+void setStackSize(size_t stackSize);
+
+
 /* Restore the original inherited Unix process context (such as signal
    masks, stack size, CPU affinity). */
 void restoreProcessContext();
diff --git a/src/nix/main.cc b/src/nix/main.cc
index f8701ee56..008482be3 100644
--- a/src/nix/main.cc
+++ b/src/nix/main.cc
@@ -17,10 +17,6 @@
 #include <netdb.h>
 #include <netinet/in.h>
 
-#if __linux__
-#include <sys/resource.h>
-#endif
-
 #include <nlohmann/json.hpp>
 
 extern std::string chrootHelperName;
@@ -335,14 +331,7 @@ int main(int argc, char * * argv)
 {
     // Increase the default stack size for the evaluator and for
     // libstdc++'s std::regex.
-    #if __linux__
-    rlim_t stackSize = 64 * 1024 * 1024;
-    struct rlimit limit;
-    if (getrlimit(RLIMIT_STACK, &limit) == 0 && limit.rlim_cur < stackSize) {
-        limit.rlim_cur = stackSize;
-        setrlimit(RLIMIT_STACK, &limit);
-    }
-    #endif
+    nix::setStackSize(64 * 1024 * 1024);
 
     return nix::handleExceptions(argv[0], [&]() {
         nix::mainWrapped(argc, argv);