nix-daemon: add --stdio flag for handling connections over stdin/stdout

1 files changed, 75 insertions, 1 deletions

diff --git a/src/nix-daemon/nix-daemon.cc b/src/nix-daemon/nix-daemon.cc
index f2b59c84a..fba522540 100644
--- a/src/nix-daemon/nix-daemon.cc
+++ b/src/nix-daemon/nix-daemon.cc
@@ -22,6 +22,7 @@
 #include <errno.h>
 #include <pwd.h>
 #include <grp.h>
+#include <fcntl.h>
 
 #if __APPLE__ || __FreeBSD__
 #include <sys/ucred.h>
@@ -29,6 +30,25 @@
 
 using namespace nix;
 
+#ifndef __linux__
+#define SPLICE_F_MOVE 0
+static ssize_t splice(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags)
+{
+    /* We ignore most parameters, we just have them for conformance with the linux syscall */
+    char buf[8192];
+    auto read_count = read(fd_in, buf, sizeof(buf));
+    if (read_count == -1)
+        return read_count;
+    auto write_count = decltype<read_count>(0);
+    while (write_count < read_count) {
+        auto res = write(fd_out, buf + write_count, read_count - write_count);
+        if (res == -1)
+            return res;
+        write_count += res;
+    }
+    return read_count;
+}
+#endif
 
 static FdSource from(STDIN_FILENO);
 static FdSink to(STDOUT_FILENO);
@@ -885,6 +905,8 @@ int main(int argc, char * * argv)
     return handleExceptions(argv[0], [&]() {
         initNix();
 
+        auto stdio = false;
+
         parseCmdLine(argc, argv, [&](Strings::iterator & arg, const Strings::iterator & end) {
             if (*arg == "--daemon")
                 ; /* ignored for backwards compatibility */
@@ -892,10 +914,62 @@ int main(int argc, char * * argv)
                 showManPage("nix-daemon");
             else if (*arg == "--version")
                 printVersion("nix-daemon");
+            else if (*arg == "--stdio")
+                stdio = true;
             else return false;
             return true;
         });
 
-        daemonLoop(argv);
+        if (stdio) {
+            if (getStoreType() == tDaemon) {
+                /* Forward on this connection to the real daemon */
+                auto socketPath = settings.nixDaemonSocketFile;
+                auto s = socket(PF_UNIX, SOCK_STREAM, 0);
+                if (s == -1)
+                    throw SysError("creating Unix domain socket");
+
+                auto socketDir = dirOf(socketPath);
+                if (chdir(socketDir.c_str()) == -1)
+                    throw SysError(format("changing to socket directory %1%") % socketDir);
+
+                auto socketName = baseNameOf(socketPath);
+                auto addr = sockaddr_un{};
+                addr.sun_family = AF_UNIX;
+                if (socketName.size() + 1 >= sizeof(addr.sun_path))
+                    throw Error(format("socket name %1% is too long") % socketName);
+                strcpy(addr.sun_path, socketName.c_str());
+
+                if (connect(s, (struct sockaddr *) &addr, sizeof(addr)) == -1)
+                    throw SysError(format("cannot connect to daemon at %1%") % socketPath);
+
+                auto nfds = (s > STDIN_FILENO ? s : STDIN_FILENO) + 1;
+                while (true) {
+                    fd_set fds;
+                    FD_ZERO(&fds);
+                    FD_SET(s, &fds);
+                    FD_SET(STDIN_FILENO, &fds);
+                    if (select(nfds, &fds, nullptr, nullptr, nullptr) == -1)
+                        throw SysError("waiting for data from client or server");
+                    if (FD_ISSET(s, &fds)) {
+                        auto res = splice(s, nullptr, STDOUT_FILENO, nullptr, SIZE_MAX, SPLICE_F_MOVE);
+                        if (res == -1)
+                            throw SysError("splicing data from daemon socket to stdout");
+                        else if (res == 0)
+                            throw EndOfFile("unexpected EOF from daemon socket");
+                    }
+                    if (FD_ISSET(STDIN_FILENO, &fds)) {
+                        auto res = splice(STDIN_FILENO, nullptr, s, nullptr, SIZE_MAX, SPLICE_F_MOVE);
+                        if (res == -1)
+                            throw SysError("splicing data from stdin to daemon socket");
+                        else if (res == 0)
+                            return;
+                    }
+                }
+            } else {
+                processConnection(true);
+            }
+        } else {
+            daemonLoop(argv);
+        }
     });
 }