From c7d97802e4f59b8621e67cf62275d6a7fde8fe62 Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Fri, 9 Aug 2024 21:17:52 +0200
Subject: [PATCH] libutil: rename and optimize closeMostFDs

this is only used to close non-stdio files in derivation sandboxes. we
may as well encode that in its name, drop the unnecessary integer set,
and use close_range to deal with the actual closing of files. not only
is this clearer, it also makes sandbox setup on linux fast by 1ms each

Change-Id: Id90e259a49c7bc896189e76bfbbf6ef2c0bcd3b2
---
 src/libstore/build/local-derivation-goal.cc |  2 +-
 src/libutil/file-descriptor.cc              | 20 +++++++++++++++-----
 src/libutil/file-descriptor.hh              |  4 ++--
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/libstore/build/local-derivation-goal.cc b/src/libstore/build/local-derivation-goal.cc
index 8b640e4ad..246cb68f6 100644
--- a/src/libstore/build/local-derivation-goal.cc
+++ b/src/libstore/build/local-derivation-goal.cc
@@ -1618,7 +1618,7 @@ void LocalDerivationGoal::runChild()
             throw SysError("changing into '%1%'", tmpDir);
 
         /* Close all other file descriptors. */
-        closeMostFDs({STDIN_FILENO, STDOUT_FILENO, STDERR_FILENO});
+        closeExtraFDs();
 
         setPersonality(drv->platform);
 
diff --git a/src/libutil/file-descriptor.cc b/src/libutil/file-descriptor.cc
index be9f8c889..4807869f9 100644
--- a/src/libutil/file-descriptor.cc
+++ b/src/libutil/file-descriptor.cc
@@ -218,13 +218,24 @@ void Pipe::close()
 }
 
 
-void closeMostFDs(const std::set<int> & exceptions)
+void closeExtraFDs()
 {
+    constexpr int MAX_KEPT_FD = 2;
+    static_assert(std::max({STDIN_FILENO, STDOUT_FILENO, STDERR_FILENO}) == MAX_KEPT_FD);
+
 #if __linux__
+    // first try to close_range everything we don't care about. if this
+    // returns an error with these parameters we're running on a kernel
+    // that does not implement close_range (i.e. pre 5.9) and fall back
+    // to the old method. we should remove that though, in some future.
+    if (close_range(3, ~0U, 0) == 0) {
+        return;
+    }
+
     try {
         for (auto & s : readDirectory("/proc/self/fd")) {
             auto fd = std::stoi(s.name);
-            if (!exceptions.count(fd)) {
+            if (fd > MAX_KEPT_FD) {
                 debug("closing leaked FD %d", fd);
                 close(fd);
             }
@@ -236,9 +247,8 @@ void closeMostFDs(const std::set<int> & exceptions)
 
     int maxFD = 0;
     maxFD = sysconf(_SC_OPEN_MAX);
-    for (int fd = 0; fd < maxFD; ++fd)
-        if (!exceptions.count(fd))
-            close(fd); /* ignore result */
+    for (int fd = MAX_KEPT_FD + 1; fd < maxFD; ++fd)
+        close(fd); /* ignore result */
 }
 
 
diff --git a/src/libutil/file-descriptor.hh b/src/libutil/file-descriptor.hh
index a83bc028f..7270b73b5 100644
--- a/src/libutil/file-descriptor.hh
+++ b/src/libutil/file-descriptor.hh
@@ -66,10 +66,10 @@ public:
 };
 
 /**
- * Close all file descriptors except those listed in the given set.
+ * Close all file descriptors except stdio fds (ie 0, 1, 2).
  * Good practice in child processes.
  */
-void closeMostFDs(const std::set<int> & exceptions);
+void closeExtraFDs();
 
 /**
  * Set the close-on-exec flag for the given file descriptor.