2 files changed, 84 insertions, 14 deletions

diff --git a/third_party/nix/src/nix-daemon/CMakeLists.txt b/third_party/nix/src/nix-daemon/CMakeLists.txt
index 27a4a1254e77..63125a9b26b2 100644
--- a/third_party/nix/src/nix-daemon/CMakeLists.txt
+++ b/third_party/nix/src/nix-daemon/CMakeLists.txt
@@ -8,6 +8,8 @@ add_executable(nix-daemon)
 include_directories(${PROJECT_BINARY_DIR}) # for config.h
 set_property(TARGET nix-daemon PROPERTY CXX_STANDARD 17)
 
+pkg_check_modules(systemd REQUIRED)
+
 target_sources(nix-daemon
   PRIVATE
     nix-daemon-proto.hh
@@ -21,6 +23,7 @@ target_link_libraries(nix-daemon
   nixmain
   absl::flags
   absl::flags_parse
+  systemd
 )
 
 install(TARGETS nix-daemon DESTINATION bin)
diff --git a/third_party/nix/src/nix-daemon/nix-daemon.cc b/third_party/nix/src/nix-daemon/nix-daemon.cc
index f4128fa4b252..fd7a553e209e 100644
--- a/third_party/nix/src/nix-daemon/nix-daemon.cc
+++ b/third_party/nix/src/nix-daemon/nix-daemon.cc
@@ -8,9 +8,11 @@
 #include <glog/logging.h>
 #include <grpcpp/security/server_credentials.h>
 #include <grpcpp/server.h>
-#include <grpcpp/server_builder_impl.h>
+#include <grpcpp/server_builder.h>
+#include <grpcpp/server_posix.h>
 #include <sys/socket.h>
 #include <sys/un.h>
+#include <systemd/sd-daemon.h>
 
 #include "libmain/shared.hh"  // TODO(tazjin): can this be removed?
 #include "libstore/globals.hh"
@@ -27,6 +29,8 @@ namespace nix::daemon {
 using grpc::Server;
 using grpc_impl::ServerBuilder;
 
+namespace {
+
 // TODO(grfn): There has to be a better way to do this - this was ported
 // verbatim from the old daemon implementation without much critical evaluation.
 static int ForwardToSocket(nix::Path socket_path) {
@@ -86,35 +90,98 @@ static int ForwardToSocket(nix::Path socket_path) {
   }
 }
 
+void SetNonBlocking(int fd) {
+  int flags = fcntl(fd, F_GETFL);  // NOLINT
+  PCHECK(flags != 0) << "Error getting socket flags";
+  PCHECK(fcntl(  // NOLINT
+             fd, F_SETFL, flags | O_NONBLOCK) == 0)
+      << "Could not set socket flags";
+}
+
+}  // namespace
+
 int RunServer() {
   Store::Params params;
   params["path-info-cache-size"] = "0";
   auto store = openStore(settings.storeUri, params);
   auto worker = NewWorkerService(*store);
-
-  std::filesystem::path socket_path(settings.nixDaemonSocketFile);
-  std::filesystem::create_directories(socket_path.parent_path());
-  auto socket_addr = absl::StrFormat("unix://%s", socket_path);
-
   ServerBuilder builder;
-  builder.AddListeningPort(socket_addr, grpc::InsecureServerCredentials());
   builder.RegisterService(worker);
 
+  auto n_fds = sd_listen_fds(0);
+
+  if (n_fds > 1) {
+    LOG(FATAL) << "Too many file descriptors (" << n_fds
+               << ") received from systemd socket activation";
+  }
+
+  std::filesystem::path socket_path;
+
+  if (n_fds == 0) {
+    socket_path = settings.nixDaemonSocketFile;
+    std::filesystem::create_directories(socket_path.parent_path());
+    auto socket_addr = absl::StrFormat("unix://%s", socket_path);
+    builder.AddListeningPort(socket_addr, grpc::InsecureServerCredentials());
+  }
+
   std::unique_ptr<Server> server(builder.BuildAndStart());
-  if (server) {
-    LOG(INFO) << "Nix daemon listening at " << socket_addr;
-    server->Wait();
-    return 0;
-  } else {
+
+  if (!server) {
+    LOG(FATAL) << "Error building server";
     return 1;
   }
+
+  // We have been systemd socket-activated - instead of asking grpc to make the
+  // socket path for us, start our own accept loop and pass file descriptors to
+  // grpc.
+  //
+  // This approach was *somewhat* adapted from
+  // https://gist.github.com/yorickvP/8d523a4df2b10c5812fa7789e82b7c1b - at some
+  // point we'd like gRPC to do it for us, though - see
+  // https://github.com/grpc/grpc/issues/19133
+  if (n_fds == 1) {
+    int socket_fd = SD_LISTEN_FDS_START;
+    // Only used for logging
+    socket_path = readLink(absl::StrFormat("/proc/self/fd/%d", socket_fd));
+
+    PCHECK(sd_notify(0, "READY=1") == 0) << "Error notifying systemd";
+    for (;;) {
+      try {
+        struct sockaddr_un remote_addr {};
+        socklen_t remote_addr_len = sizeof(remote_addr);
+        int remote_fd =
+            accept(socket_fd,
+                   reinterpret_cast<struct sockaddr*>(&remote_addr),  // NOLINT
+                   &remote_addr_len);
+        checkInterrupt();
+        if (!remote_fd) {
+          if (errno == EINTR) {
+            continue;
+          }
+          PCHECK(false) << "error accepting connection";
+        }
+
+        LOG(INFO) << "Accepted remote connection on fd " << remote_fd;
+        SetNonBlocking(remote_fd);
+        grpc::AddInsecureChannelFromFd(server.get(), remote_fd);
+      } catch (Interrupted& e) {
+        return -1;
+      } catch (Error& e) {
+        LOG(ERROR) << "error processing connection: " << e.msg();
+      }
+    }
+  }
+
+  LOG(INFO) << "Nix daemon listening at " << socket_path;
+  server->Wait();
+  return 0;
 }
 
 }  // namespace nix::daemon
 
-int main(int argc, char** argv) {
+int main(int argc, char** argv) {  // NOLINT
   FLAGS_logtostderr = true;
-  google::InitGoogleLogging(argv[0]);
+  google::InitGoogleLogging(argv[0]);  // NOLINT
 
   absl::SetFlagsUsageConfig({.version_string = [] { return nix::nixVersion; }});
   absl::ParseCommandLine(argc, argv);