Split the root finding in a separate library

flake.nix

@@ -412,7 +412,7 @@
             CXXFLAGS = prev.lib.optionalString prev.stdenv.hostPlatform.isStatic "-static";
 
             buildPhase = ''
-              $CXX $CXXFLAGS -std=c++17 nix-find-roots.cc -o nix-find-roots
+              $CXX $CXXFLAGS -std=c++17 *.cc **/*.cc -I lib -o nix-find-roots
             '';
 
             installPhase = ''

src/nix-find-roots/lib/find-roots.cc

@@ -7,114 +7,19 @@
  * This program intentionnally doesnt depend on any Nix library to reduce the attack surface.
  */
 
-#include <filesystem>
-#include <iostream>
 #include <regex>
-#include <set>
 #include <unistd.h>
 #include <vector>
 #include <algorithm>
-#include <getopt.h>
 #include <fstream>
-#include <sys/socket.h>
-#include <sys/un.h>
 
+#include "find-roots.hh"
+
+
+namespace nix::roots_tracer {
 namespace fs = std::filesystem;
 using std::set, std::string;
 
-class Error : public std::exception {
-private:
-    const string message;
-
-public:
-    Error(std::string message)
-        : message(message)
-    {}
-
-    const char* what() const noexcept override
-    {
-        return message.c_str();
-    }
-};
-
-struct GlobalOpts {
-    fs::path storeDir = "/nix/store";
-    fs::path stateDir = "/nix/var/nix";
-    fs::path socketPath = "/nix/var/nix/gc-socket/socket";
-    enum VerbosityLvl {
-        Quiet = 0,
-        Normal = 1,
-        Verbose = 2
-    };
-    VerbosityLvl verbosity = Quiet;
-};
-
-void log(GlobalOpts::VerbosityLvl currentVerbosity, GlobalOpts::VerbosityLvl minVerbosity, std::string_view msg)
-{
-    if (currentVerbosity < minVerbosity)
-        return;
-    std::cerr << msg << std::endl;
-}
-
-void debug(GlobalOpts::VerbosityLvl currentVerbosity, std::string_view msg)
-{
-    log(currentVerbosity, GlobalOpts::Verbose, msg);
-}
-
-GlobalOpts parseCmdLine(int argc, char** argv)
-{
-    GlobalOpts res;
-    auto usage = [&]() {
-        std::cerr << "Usage: " << string(argv[0]) << " [--verbose|-v] [-s storeDir] [-d stateDir] [-l socketPath]" << std::endl;
-        exit(1);
-    };
-    static struct option long_options[] = {
-        { "verbose", no_argument, (int*)&res.verbosity, GlobalOpts::Verbose },
-        { "socket_path", required_argument, 0, 'l' },
-        { "store_dir", required_argument, 0, 's' },
-        { "state_dir", required_argument, 0, 'd' },
-    };
-
-    int option_index = 0;
-    int opt_char;
-    while((opt_char = getopt_long(argc, argv, "vd:s:l:",
-                    long_options, &option_index)) != -1) {
-        switch (opt_char) {
-            case 0:
-                break;
-          break;
-            case '?':
-                usage();
-                break;
-            case 'v':
-                res.verbosity = GlobalOpts::Verbose;
-                break;
-            case 's':
-                res.storeDir = fs::path(optarg);
-                break;
-            case 'd':
-                res.stateDir = fs::path(optarg);
-                break;
-            case 'l':
-                res.socketPath = fs::path(optarg);
-                break;
-            default:
-                std::cerr << "Got invalid char: " << (char)opt_char << std::endl;
-                abort();
-        }
-    };
-    return res;
-}
-
-/*
- * A value of type `Roots` is a mapping from a store path to the set of roots that keep it alive
- */
-typedef std::map<fs::path, std::set<fs::path>> Roots;
-struct TraceResult {
-    Roots storeRoots;
-    set<fs::path> deadLinks;
-};
-
 string quoteRegexChars(const string & raw)
 {
     static auto specialRegex = std::regex(R"([.^$\\*+?()\[\]{}|])");
@@ -130,8 +35,8 @@ bool isInStore(fs::path storeDir, fs::path dir)
     return (std::search(dir.begin(), dir.end(), storeDir.begin(), storeDir.end()) == dir.begin());
 }
 
-void followPathToStore(
+void traceStaticRoot(
-    const GlobalOpts & opts,
+    const TracerConfig & opts,
     int recursionsLeft,
     TraceResult & res,
     const fs::path & root,
@@ -139,7 +44,7 @@ void followPathToStore(
     const std::optional<const fs::path> & parent = std::nullopt
     )
 {
-    debug(opts.verbosity, "Considering file " + root.string());
+    opts.debug("Considering file " + root.string());
 
     if (recursionsLeft < 0)
         return;
@@ -149,15 +54,15 @@ void followPathToStore(
             {
                 auto directory_iterator = fs::recursive_directory_iterator(root);
                 for (auto & child : directory_iterator)
-                    followPathToStore(opts, recursionsLeft, res, child.path(), child.symlink_status());
+                    traceStaticRoot(opts, recursionsLeft, res, child.path(), child.symlink_status());
                 break;
             }
         case fs::file_type::symlink:
             {
                 auto target = root.parent_path() / fs::read_symlink(root);
                 auto not_found = [&](std::string msg) {
-                    debug(opts.verbosity, "Error accessing the file " + target.string() + ": " + msg);
+                    opts.debug("Error accessing the file " + target.string() + ": " + msg);
-                    debug(opts.verbosity, "(When resolving the symlink " + root.string() + ")");
+                    opts.debug("(When resolving the symlink " + root.string() + ")");
                     res.deadLinks.insert(root);
                 };
                 try {
@@ -169,7 +74,7 @@ void followPathToStore(
                         res.storeRoots[target].insert(root);
                         return;
                     } else {
-                        followPathToStore(opts, recursionsLeft - 1, res, target, target_status);
+                        traceStaticRoot(opts, recursionsLeft - 1, res, target, target_status);
                     }
 
                 } catch (fs::filesystem_error & e) {
@@ -187,17 +92,17 @@ void followPathToStore(
     }
 }
 
-void followPathToStore(
+void traceStaticRoot(
-    const GlobalOpts & opts,
+    const TracerConfig & opts,
     int recursionsLeft,
     TraceResult & res,
     const fs::path & root)
 {
     try {
         auto status = fs::symlink_status(root);
-        followPathToStore(opts, recursionsLeft, res, root, status);
+        traceStaticRoot(opts, recursionsLeft, res, root, status);
     } catch (fs::filesystem_error & e) {
-        debug(opts.verbosity, "Error accessing the file " + root.string() + ": " + e.what());
+        opts.debug("Error accessing the file " + root.string() + ": " + e.what());
     }
 }
 
@@ -212,12 +117,12 @@ void followPathToStore(
  * Also returns the set of all dead links encountered during the process (so
  * that they can be removed if it makes sense).
  */
-TraceResult followPathsToStore(GlobalOpts opts, set<fs::path> roots)
+TraceResult traceStaticRoots(TracerConfig opts, set<fs::path> roots)
 {
     int maxRecursionLevel = 2;
     TraceResult res;
     for (auto & root : roots)
-        followPathToStore(opts, maxRecursionLevel, res, root);
+        traceStaticRoot(opts, maxRecursionLevel, res, root);
     return res;
 }
 
@@ -226,7 +131,7 @@ TraceResult followPathsToStore(GlobalOpts opts, set<fs::path> roots)
  * like a store path (i.e. that matches `storePathRegex(opts.storeDir)`) and add them
  * to `res`
  */
-void scanFileContent(const GlobalOpts & opts, const fs::path & fileToScan, Roots & res)
+void scanFileContent(const TracerConfig & opts, const fs::path & fileToScan, Roots & res)
 {
     if (!fs::exists(fileToScan))
         return;
@@ -250,7 +155,7 @@ void scanFileContent(const GlobalOpts & opts, const fs::path & fileToScan, Roots
  * Scan the content of a `/proc/[pid]/maps` file for regions that are mmaped to
  * a store path
  */
-void scanMapsFile(const GlobalOpts & opts, const fs::path & mapsFile, Roots & res)
+void scanMapsFile(const TracerConfig & opts, const fs::path & mapsFile, Roots & res)
 {
     if (!fs::exists(mapsFile))
         return;
@@ -274,7 +179,7 @@ void scanMapsFile(const GlobalOpts & opts, const fs::path & mapsFile, Roots & re
 
 }
 
-Roots getRuntimeRoots(GlobalOpts opts)
+Roots getRuntimeRoots(TracerConfig opts)
 {
     auto procDir = fs::path("/proc");
     if (!fs::exists(procDir))
@@ -288,7 +193,7 @@ Roots getRuntimeRoots(GlobalOpts opts)
             || !procEntry.is_directory())
             continue;
 
-        debug(opts.verbosity, "Considering path " + procEntry.path().string());
+        opts.debug("Considering path " + procEntry.path().string());
 
         // A set of paths used by the executable and possibly symlinks to a
         // path in the store
@@ -308,7 +213,7 @@ Roots getRuntimeRoots(GlobalOpts opts)
             if (isInStore(opts.storeDir, realPath))
                 res[realPath].insert(path);
         } catch (fs::filesystem_error &e) {
-            debug(opts.verbosity, e.what());
+            opts.debug(e.what());
         }
 
         // Scan the environment of the executable
@@ -325,66 +230,4 @@ Roots getRuntimeRoots(GlobalOpts opts)
     return res;
 }
 
-int main(int argc, char * * argv)
-{
-    const GlobalOpts opts = parseCmdLine(argc, argv);
-    const set<fs::path> standardRoots = {
-        opts.stateDir / fs::path("profiles"),
-        opts.stateDir / fs::path("gcroots"),
-    };
-
-    int mySock = socket(PF_UNIX, SOCK_STREAM, 0);
-    if (mySock == 0) {
-        throw Error("Cannot create Unix domain socket");
-    }
-    struct sockaddr_un addr;
-    addr.sun_family = AF_UNIX;
-
-    unlink(opts.socketPath.c_str());
-    strcpy(addr.sun_path, opts.socketPath.c_str());
-    if (bind(mySock, (struct sockaddr*) &addr, sizeof(addr)) == -1) {
-        throw Error("Cannot bind to socket");
-    }
-
-    if (listen(mySock, 5) == -1)
-        throw Error("cannot listen on socket " + opts.socketPath.string());
-
-    addr.sun_family = AF_UNIX;
-    while (1) {
-        struct sockaddr_un remoteAddr;
-        socklen_t remoteAddrLen = sizeof(remoteAddr);
-        int remoteSocket = accept(
-                mySock,
-                (struct sockaddr*) & remoteAddr,
-                &remoteAddrLen
-                );
-
-        if (remoteSocket == -1) {
-            if (errno == EINTR) continue;
-            throw Error("Error accepting the connection");
-        }
-
-        log(opts.verbosity, GlobalOpts::Quiet, "accepted connection");
-
-        auto traceResult = followPathsToStore(opts, standardRoots);
-        auto runtimeRoots = getRuntimeRoots(opts);
-        traceResult.storeRoots.insert(runtimeRoots.begin(), runtimeRoots.end());
-        for (auto & [rootInStore, externalRoots] : traceResult.storeRoots) {
-            for (auto & externalRoot : externalRoots) {
-                send(remoteSocket, rootInStore.string().c_str(), rootInStore.string().size(), 0);
-                send(remoteSocket, "\t", strlen("\t"), 0);
-                send(remoteSocket, externalRoot.string().c_str(), externalRoot.string().size(), 0);
-                send(remoteSocket, "\n", strlen("\n"), 0);
-            }
-
-        }
-        send(remoteSocket, "\n", strlen("\n"), 0);
-        for (auto & deadLink : traceResult.deadLinks) {
-            send(remoteSocket, deadLink.string().c_str(), deadLink.string().size(), 0);
-            send(remoteSocket, "\n", strlen("\n"), 0);
-        }
-
-        close(remoteSocket);
-    }
 }
-

src/nix-find-roots/lib/find-roots.hh

@@ -0,0 +1,48 @@
+#include <filesystem>
+#include <set>
+#include <map>
+#include <functional>
+
+namespace nix::roots_tracer {
+namespace fs = std::filesystem;
+using std::set, std::map, std::string;
+
+class Error : public std::exception {
+private:
+    const string message;
+
+public:
+    Error(std::string message)
+        : message(message)
+    {}
+
+    const char* what() const noexcept override
+    {
+        return message.c_str();
+    }
+};
+
+inline void logNone(std::string_view)
+{ }
+
+struct TracerConfig {
+    fs::path storeDir = "/nix/store";
+    fs::path stateDir = "/nix/var/nix";
+    fs::path socketPath = "/nix/var/nix/gc-socket/socket";
+
+    std::function<void(std::string_view msg)> log = logNone;
+    std::function<void(std::string_view msg)> debug = logNone;
+};
+
+/*
+ * A value of type `Roots` is a mapping from a store path to the set of roots that keep it alive
+ */
+typedef map<fs::path, std::set<fs::path>> Roots;
+struct TraceResult {
+    Roots storeRoots;
+    set<fs::path> deadLinks;
+};
+
+TraceResult traceStaticRoots(TracerConfig opts, set<fs::path> initialRoots);
+Roots getRuntimeRoots(TracerConfig opts);
+}

src/nix-find-roots/local.mk

@@ -1,9 +1,22 @@
 ifndef HOST_DARWIN
+libraries += find-roots
+
+find-roots_NAME = libfindroots
+
+find-roots_DIR := $(d)/lib
+
+find-roots_SOURCES := $(wildcard $(d)/lib/*.cc)
+
 programs += nix-find-roots
 
 nix-find-roots_DIR := $(d)
 
-nix-find-roots_SOURCES := $(wildcard $(d)/*.cc)
+nix-find-roots_SOURCES := $(d)/main.cc
+
+nix-find-roots_LIBS := find-roots
+
+nix-find-roots_CXXFLAGS += \
+	-I src/nix-find-roots/lib
 
 nix-find-roots_INSTALL_DIR := $(libexecdir)/nix
 endif

src/nix-find-roots/main.cc

@@ -0,0 +1,124 @@
+#include "find-roots.hh"
+#include <getopt.h>
+#include <sys/socket.h>
+#include <sys/un.h>
+#include <iostream>
+#include <unistd.h>
+#include <cstring>
+
+using namespace nix::roots_tracer;
+
+void logStderr(std::string_view msg)
+{
+  std::cerr << msg << std::endl;
+}
+
+TracerConfig parseCmdLine(int argc, char** argv)
+{
+    TracerConfig res;
+    res.log = logStderr;
+    auto usage = [&]() {
+        std::cerr << "Usage: " << string(argv[0]) << " [--verbose|-v] [-s storeDir] [-d stateDir] [-l socketPath]" << std::endl;
+        exit(1);
+    };
+    static struct option long_options[] = {
+        { "verbose", no_argument, 0, 'v' },
+        { "socket_path", required_argument, 0, 'l' },
+        { "store_dir", required_argument, 0, 's' },
+        { "state_dir", required_argument, 0, 'd' },
+    };
+
+    int option_index = 0;
+    int opt_char;
+    while((opt_char = getopt_long(argc, argv, "vd:s:l:",
+                    long_options, &option_index)) != -1) {
+        switch (opt_char) {
+            case 0:
+                break;
+          break;
+            case '?':
+                usage();
+                break;
+            case 'v':
+                res.debug = logStderr;
+                break;
+            case 's':
+                res.storeDir = fs::path(optarg);
+                break;
+            case 'd':
+                res.stateDir = fs::path(optarg);
+                break;
+            case 'l':
+                res.socketPath = fs::path(optarg);
+                break;
+            default:
+                std::cerr << "Got invalid char: " << (char)opt_char << std::endl;
+                abort();
+        }
+    };
+    return res;
+}
+
+
+int main(int argc, char * * argv)
+{
+    const TracerConfig opts = parseCmdLine(argc, argv);
+    const set<fs::path> standardRoots = {
+        opts.stateDir / fs::path("profiles"),
+        opts.stateDir / fs::path("gcroots"),
+    };
+
+    int mySock = socket(PF_UNIX, SOCK_STREAM, 0);
+    if (mySock == 0) {
+        throw Error("Cannot create Unix domain socket");
+    }
+    struct sockaddr_un addr;
+    addr.sun_family = AF_UNIX;
+
+    unlink(opts.socketPath.c_str());
+    strcpy(addr.sun_path, opts.socketPath.c_str());
+    if (bind(mySock, (struct sockaddr*) &addr, sizeof(addr)) == -1) {
+        throw Error("Cannot bind to socket");
+    }
+
+    if (listen(mySock, 5) == -1)
+        throw Error("cannot listen on socket " + opts.socketPath.string());
+
+    addr.sun_family = AF_UNIX;
+    while (1) {
+        struct sockaddr_un remoteAddr;
+        socklen_t remoteAddrLen = sizeof(remoteAddr);
+        int remoteSocket = accept(
+                mySock,
+                (struct sockaddr*) & remoteAddr,
+                &remoteAddrLen
+                );
+
+        if (remoteSocket == -1) {
+            if (errno == EINTR) continue;
+            throw Error("Error accepting the connection");
+        }
+
+        opts.log("accepted connection");
+
+        auto traceResult = traceStaticRoots(opts, standardRoots);
+        auto runtimeRoots = getRuntimeRoots(opts);
+        traceResult.storeRoots.insert(runtimeRoots.begin(), runtimeRoots.end());
+        for (auto & [rootInStore, externalRoots] : traceResult.storeRoots) {
+            for (auto & externalRoot : externalRoots) {
+                send(remoteSocket, rootInStore.string().c_str(), rootInStore.string().size(), 0);
+                send(remoteSocket, "\t", strlen("\t"), 0);
+                send(remoteSocket, externalRoot.string().c_str(), externalRoot.string().size(), 0);
+                send(remoteSocket, "\n", strlen("\n"), 0);
+            }
+
+        }
+        send(remoteSocket, "\n", strlen("\n"), 0);
+        for (auto & deadLink : traceResult.deadLinks) {
+            send(remoteSocket, deadLink.string().c_str(), deadLink.string().size(), 0);
+            send(remoteSocket, "\n", strlen("\n"), 0);
+        }
+
+        close(remoteSocket);
+    }
+}