From a83694c7a1212567ec2f032f84c0d72722bcf5ea Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Wed, 19 Aug 2020 19:34:47 +0000
Subject: [PATCH] Use `RemoteStore` to open connection for proxying daemon

Removes duplicate websocket opening code, and also means we should be
able to to ssh-ssh-... daemon relays, not just uds-uds-... ones.
---
 src/libstore/remote-store.cc | 13 ++++++++++---
 src/libstore/remote-store.hh |  4 ++--
 src/libstore/ssh-store.cc    |  1 -
 src/nix-daemon/nix-daemon.cc | 36 +++++++++---------------------------
 4 files changed, 21 insertions(+), 33 deletions(-)

diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index 553069b89..ff7149643 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -86,7 +86,16 @@ RemoteStore::RemoteStore(const Params & params)
     : Store(params)
     , connections(make_ref<Pool<Connection>>(
             std::max(1, (int) maxConnections),
-            [this]() { return openConnectionWrapper(); },
+            [this]() {
+                auto conn = openConnectionWrapper();
+                try {
+                    initConnection(*conn);
+                } catch (...) {
+                    failed = true;
+                    throw;
+                }
+                return conn;
+            },
             [this](const ref<Connection> & r) {
                 return
                     r->to.good()
@@ -169,8 +178,6 @@ ref<RemoteStore::Connection> UDSRemoteStore::openConnection()
 
     conn->startTime = std::chrono::steady_clock::now();
 
-    initConnection(*conn);
-
     return conn;
 }
 
diff --git a/src/libstore/remote-store.hh b/src/libstore/remote-store.hh
index 72d2a6689..5ad2dc75b 100644
--- a/src/libstore/remote-store.hh
+++ b/src/libstore/remote-store.hh
@@ -102,8 +102,6 @@ public:
 
     void flushBadConnections();
 
-protected:
-
     struct Connection
     {
         AutoCloseFD fd;
@@ -119,6 +117,8 @@ protected:
 
     ref<Connection> openConnectionWrapper();
 
+protected:
+
     virtual ref<Connection> openConnection() = 0;
 
     void initConnection(Connection & conn);
diff --git a/src/libstore/ssh-store.cc b/src/libstore/ssh-store.cc
index caae6b596..7e9b44afe 100644
--- a/src/libstore/ssh-store.cc
+++ b/src/libstore/ssh-store.cc
@@ -89,7 +89,6 @@ ref<RemoteStore::Connection> SSHStore::openConnection()
         + (remoteStore.get() == "" ? "" : " --store " + shellEscape(remoteStore.get())));
     conn->to = FdSink(conn->sshConn->in.get());
     conn->from = FdSource(conn->sshConn->out.get());
-    initConnection(*conn);
     return conn;
 }
 
diff --git a/src/nix-daemon/nix-daemon.cc b/src/nix-daemon/nix-daemon.cc
index 9613cb7d3..bdf56d2e2 100644
--- a/src/nix-daemon/nix-daemon.cc
+++ b/src/nix-daemon/nix-daemon.cc
@@ -286,46 +286,28 @@ static int _main(int argc, char * * argv)
         initPlugins();
 
         if (stdio) {
-            if (openUncachedStore().dynamic_pointer_cast<UDSRemoteStore>()) {
-                // FIXME Use the connection the UDSRemoteStore opened
+            if (auto store = openUncachedStore().dynamic_pointer_cast<RemoteStore>()) {
+                auto conn = store->openConnectionWrapper();
+                int from = conn->from.fd;
+                int to = conn->to.fd;
 
-                //  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("changing to socket directory '%1%'", socketDir);
-
-                auto socketName = std::string(baseNameOf(socketPath));
-                auto addr = sockaddr_un{};
-                addr.sun_family = AF_UNIX;
-                if (socketName.size() + 1 >= sizeof(addr.sun_path))
-                    throw Error("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("cannot connect to daemon at %1%", socketPath);
-
-                auto nfds = (s > STDIN_FILENO ? s : STDIN_FILENO) + 1;
+                auto nfds = std::max(from, to) + 1;
                 while (true) {
                     fd_set fds;
                     FD_ZERO(&fds);
-                    FD_SET(s, &fds);
+                    FD_SET(from, &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, SSIZE_MAX, SPLICE_F_MOVE);
+                    if (FD_ISSET(from, &fds)) {
+                        auto res = splice(from, nullptr, STDOUT_FILENO, nullptr, SSIZE_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, SSIZE_MAX, SPLICE_F_MOVE);
+                        auto res = splice(STDIN_FILENO, nullptr, to, nullptr, SSIZE_MAX, SPLICE_F_MOVE);
                         if (res == -1)
                             throw SysError("splicing data from stdin to daemon socket");
                         else if (res == 0)