Improve SSH handling

* Unify SSH code in SSHStore and LegacySSHStore.

* Fix a race starting the SSH master. We now wait synchronously for
  the SSH master to finish starting. This prevents the SSH clients
  from starting their own connections.

* Don't use a master if max-connections == 1.

* Add a "max-connections" store parameter.

* Add a "compress" store parameter.

7 files changed, 185 insertions, 113 deletions

diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index 031fcac95..4a2ac42f3 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -4,6 +4,7 @@
 #include "serve-protocol.hh"
 #include "store-api.hh"
 #include "worker-protocol.hh"
+#include "ssh.hh"
 
 namespace nix {
 
@@ -11,73 +12,42 @@ static std::string uriScheme = "legacy-ssh://";
 
 struct LegacySSHStore : public Store
 {
-    string host;
-
     struct Connection
     {
-        Pid sshPid;
-        AutoCloseFD out;
-        AutoCloseFD in;
+        std::unique_ptr<SSHMaster::Connection> sshConn;
         FdSink to;
         FdSource from;
     };
 
-    AutoDelete tmpDir;
-
-    Path socketPath;
-
-    Pid sshMaster;
+    std::string host;
 
     ref<Pool<Connection>> connections;
 
-    Path key;
+    SSHMaster master;
 
-    LegacySSHStore(const string & host, const Params & params,
-        size_t maxConnections = std::numeric_limits<size_t>::max())
+    LegacySSHStore(const string & host, const Params & params)
         : Store(params)
         , host(host)
-        , tmpDir(createTempDir("", "nix", true, true, 0700))
-        , socketPath((Path) tmpDir + "/ssh.sock")
         , connections(make_ref<Pool<Connection>>(
-            maxConnections,
+            std::max(1, std::stoi(get(params, "max-connections", "1"))),
             [this]() { return openConnection(); },
             [](const ref<Connection> & r) { return true; }
             ))
-        , key(get(params, "ssh-key", ""))
+        , master(
+            host,
+            get(params, "ssh-key", ""),
+            // Use SSH master only if using more than 1 connection.
+            connections->capacity() > 1,
+            get(params, "compress", "") == "true")
     {
     }
 
     ref<Connection> openConnection()
     {
-        if ((pid_t) sshMaster == -1) {
-            sshMaster = startProcess([&]() {
-                restoreSignals();
-                Strings args{ "ssh", "-M", "-S", socketPath, "-N", "-x", "-a", host };
-                if (!key.empty())
-                    args.insert(args.end(), {"-i", key});
-                execvp("ssh", stringsToCharPtrs(args).data());
-                throw SysError("starting SSH master connection to host ‘%s’", host);
-            });
-        }
-
         auto conn = make_ref<Connection>();
-        Pipe in, out;
-        in.create();
-        out.create();
-        conn->sshPid = startProcess([&]() {
-            if (dup2(in.readSide.get(), STDIN_FILENO) == -1)
-                throw SysError("duping over STDIN");
-            if (dup2(out.writeSide.get(), STDOUT_FILENO) == -1)
-                throw SysError("duping over STDOUT");
-            execlp("ssh", "ssh", "-S", socketPath.c_str(), host.c_str(), "nix-store", "--serve", "--write", nullptr);
-            throw SysError("executing ‘nix-store --serve’ on remote host ‘%s’", host);
-        });
-        in.readSide = -1;
-        out.writeSide = -1;
-        conn->out = std::move(out.readSide);
-        conn->in = std::move(in.writeSide);
-        conn->to = FdSink(conn->in.get());
-        conn->from = FdSource(conn->out.get());
+        conn->sshConn = master.startCommand("nix-store --serve");
+        conn->to = FdSink(conn->sshConn->in.get());
+        conn->from = FdSource(conn->sshConn->out.get());
 
         int remoteVersion;
 
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index 47413d573..5e62bd3d5 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -40,10 +40,10 @@ template PathSet readStorePaths(Store & store, Source & from);
 template Paths readStorePaths(Store & store, Source & from);
 
 /* TODO: Separate these store impls into different files, give them better names */
-RemoteStore::RemoteStore(const Params & params, size_t maxConnections)
+RemoteStore::RemoteStore(const Params & params)
     : Store(params)
     , connections(make_ref<Pool<Connection>>(
-            maxConnections,
+            std::max(1, std::stoi(get(params, "max-connections", "1"))),
             [this]() { return openConnection(); },
             [](const ref<Connection> & r) { return r->to.good() && r->from.good(); }
             ))
@@ -51,10 +51,10 @@ RemoteStore::RemoteStore(const Params & params, size_t maxConnections)
 }
 
 
-UDSRemoteStore::UDSRemoteStore(const Params & params, size_t maxConnections)
+UDSRemoteStore::UDSRemoteStore(const Params & params)
     : Store(params)
     , LocalFSStore(params)
-    , RemoteStore(params, maxConnections)
+    , RemoteStore(params)
 {
 }
 
@@ -129,7 +129,7 @@ void RemoteStore::initConnection(Connection & conn)
         conn.processStderr();
     }
     catch (Error & e) {
-        throw Error(format("cannot start daemon worker: %1%") % e.msg());
+        throw Error("cannot open connection to remote store ‘%s’: %s", getUri(), e.what());
     }
 
     setOptions(conn);
diff --git a/src/libstore/remote-store.hh b/src/libstore/remote-store.hh
index 40f17da30..ed7b27c88 100644
--- a/src/libstore/remote-store.hh
+++ b/src/libstore/remote-store.hh
@@ -22,7 +22,7 @@ class RemoteStore : public virtual Store
 {
 public:
 
-    RemoteStore(const Params & params, size_t maxConnections = std::numeric_limits<size_t>::max());
+    RemoteStore(const Params & params);
 
     /* Implementations of abstract store API methods. */
 
@@ -113,7 +113,7 @@ class UDSRemoteStore : public LocalFSStore, public RemoteStore
 {
 public:
 
-    UDSRemoteStore(const Params & params, size_t maxConnections = std::numeric_limits<size_t>::max());
+    UDSRemoteStore(const Params & params);
 
     std::string getUri() override;
 
diff --git a/src/libstore/ssh-store.cc b/src/libstore/ssh-store.cc
index 6f1862afa..20f020bda 100644
--- a/src/libstore/ssh-store.cc
+++ b/src/libstore/ssh-store.cc
@@ -4,6 +4,7 @@
 #include "archive.hh"
 #include "worker-protocol.hh"
 #include "pool.hh"
+#include "ssh.hh"
 
 namespace nix {
 
@@ -13,9 +14,23 @@ class SSHStore : public RemoteStore
 {
 public:
 
-    SSHStore(string host, const Params & params, size_t maxConnections = std::numeric_limits<size_t>::max());
+    SSHStore(const std::string & host, const Params & params)
+        : Store(params)
+        , RemoteStore(params)
+        , host(host)
+        , master(
+            host,
+            get(params, "ssh-key", ""),
+            // Use SSH master only if using more than 1 connection.
+            connections->capacity() > 1,
+            get(params, "compress", "") == "true")
+    {
+    }
 
-    std::string getUri() override;
+    std::string getUri() override
+    {
+        return uriScheme + host;
+    }
 
     void narFromPath(const Path & path, Sink & sink) override;
 
@@ -25,43 +40,16 @@ private:
 
     struct Connection : RemoteStore::Connection
     {
-        Pid sshPid;
-        AutoCloseFD out;
-        AutoCloseFD in;
+        std::unique_ptr<SSHMaster::Connection> sshConn;
     };
 
     ref<RemoteStore::Connection> openConnection() override;
 
-    AutoDelete tmpDir;
-
-    Path socketPath;
-
-    Pid sshMaster;
-
-    string host;
-
-    Path key;
+    std::string host;
 
-    bool compress;
+    SSHMaster master;
 };
 
-SSHStore::SSHStore(string host, const Params & params, size_t maxConnections)
-    : Store(params)
-    , RemoteStore(params, maxConnections)
-    , tmpDir(createTempDir("", "nix", true, true, 0700))
-    , socketPath((Path) tmpDir + "/ssh.sock")
-    , host(std::move(host))
-    , key(get(params, "ssh-key", ""))
-    , compress(get(params, "compress", "") == "true")
-{
-    /* open a connection and perform the handshake to verify all is well */
-    connections->get();
-}
-
-string SSHStore::getUri()
-{
-    return uriScheme + host;
-}
 
 class ForwardSource : public Source
 {
@@ -94,35 +82,10 @@ ref<FSAccessor> SSHStore::getFSAccessor()
 
 ref<RemoteStore::Connection> SSHStore::openConnection()
 {
-    if ((pid_t) sshMaster == -1) {
-        sshMaster = startProcess([&]() {
-            restoreSignals();
-            if (key.empty())
-                execlp("ssh", "ssh", "-N", "-M", "-S", socketPath.c_str(), host.c_str(), NULL);
-            else
-                execlp("ssh", "ssh", "-N", "-M", "-S", socketPath.c_str(), "-i", key.c_str(), host.c_str(), NULL);
-            throw SysError("starting ssh master");
-        });
-    }
-
     auto conn = make_ref<Connection>();
-    Pipe in, out;
-    in.create();
-    out.create();
-    conn->sshPid = startProcess([&]() {
-        if (dup2(in.readSide.get(), STDIN_FILENO) == -1)
-            throw SysError("duping over STDIN");
-        if (dup2(out.writeSide.get(), STDOUT_FILENO) == -1)
-            throw SysError("duping over STDOUT");
-        execlp("ssh", "ssh", "-S", socketPath.c_str(), host.c_str(), "nix-daemon", "--stdio", NULL);
-        throw SysError("executing nix-daemon --stdio over ssh");
-    });
-    in.readSide = -1;
-    out.writeSide = -1;
-    conn->out = std::move(out.readSide);
-    conn->in = std::move(in.writeSide);
-    conn->to = FdSink(conn->in.get());
-    conn->from = FdSource(conn->out.get());
+    conn->sshConn = master.startCommand("nix-daemon --stdio");
+    conn->to = FdSink(conn->sshConn->in.get());
+    conn->from = FdSource(conn->sshConn->out.get());
     initConnection(*conn);
     return conn;
 }
diff --git a/src/libstore/ssh.cc b/src/libstore/ssh.cc
new file mode 100644
index 000000000..7c3de4a48
--- /dev/null
+++ b/src/libstore/ssh.cc
@@ -0,0 +1,93 @@
+#include "ssh.hh"
+
+namespace nix {
+
+std::unique_ptr<SSHMaster::Connection> SSHMaster::startCommand(const std::string & command)
+{
+    startMaster();
+
+    Pipe in, out;
+    in.create();
+    out.create();
+
+    auto conn = std::make_unique<Connection>();
+    conn->sshPid = startProcess([&]() {
+        restoreSignals();
+
+        close(in.writeSide.get());
+        close(out.readSide.get());
+
+        if (dup2(in.readSide.get(), STDIN_FILENO) == -1)
+            throw SysError("duping over stdin");
+        if (dup2(out.writeSide.get(), STDOUT_FILENO) == -1)
+            throw SysError("duping over stdout");
+
+        Strings args = { "ssh", host.c_str(), "-x", "-a" };
+        if (!keyFile.empty())
+            args.insert(args.end(), {"-i", keyFile});
+        if (compress)
+            args.push_back("-C");
+        if (useMaster)
+            args.insert(args.end(), {"-S", socketPath});
+        args.push_back(command);
+        execvp(args.begin()->c_str(), stringsToCharPtrs(args).data());
+
+        throw SysError("executing ‘%s’ on ‘%s’", command, host);
+    });
+
+
+    in.readSide = -1;
+    out.writeSide = -1;
+
+    conn->out = std::move(out.readSide);
+    conn->in = std::move(in.writeSide);
+
+    return conn;
+}
+
+void SSHMaster::startMaster()
+{
+    if (!useMaster || sshMaster != -1) return;
+
+    tmpDir = std::make_unique<AutoDelete>(createTempDir("", "nix", true, true, 0700));
+
+    socketPath = (Path) *tmpDir + "/ssh.sock";
+
+    Pipe out;
+    out.create();
+
+    sshMaster = startProcess([&]() {
+        restoreSignals();
+
+        close(out.readSide.get());
+
+        if (dup2(out.writeSide.get(), STDOUT_FILENO) == -1)
+            throw SysError("duping over stdout");
+
+        Strings args =
+            { "ssh", host.c_str(), "-M", "-N", "-S", socketPath
+            , "-o", "LocalCommand=echo started"
+            , "-o", "PermitLocalCommand=yes"
+            };
+        if (!keyFile.empty())
+            args.insert(args.end(), {"-i", keyFile});
+        if (compress)
+            args.push_back("-C");
+
+        execvp(args.begin()->c_str(), stringsToCharPtrs(args).data());
+
+        throw SysError("starting SSH master");
+    });
+
+    out.writeSide = -1;
+
+    std::string reply;
+    try {
+        reply = readLine(out.readSide.get());
+    } catch (EndOfFile & e) { }
+
+    if (reply != "started")
+        throw Error("failed to start SSH master connection to ‘%s’", host);
+}
+
+}
diff --git a/src/libstore/ssh.hh b/src/libstore/ssh.hh
new file mode 100644
index 000000000..2d2b98370
--- /dev/null
+++ b/src/libstore/ssh.hh
@@ -0,0 +1,41 @@
+#pragma once
+
+#include "util.hh"
+
+namespace nix {
+
+class SSHMaster
+{
+private:
+
+    std::string host;
+    std::string keyFile;
+    bool useMaster;
+    bool compress;
+    Pid sshMaster;
+    std::unique_ptr<AutoDelete> tmpDir;
+    Path socketPath;
+
+public:
+
+    SSHMaster(const std::string & host, const std::string & keyFile, bool useMaster, bool compress)
+        : host(host)
+        , keyFile(keyFile)
+        , useMaster(useMaster)
+        , compress(compress)
+    {
+    }
+
+    struct Connection
+    {
+        Pid sshPid;
+        AutoCloseFD out, in;
+    };
+
+    std::unique_ptr<Connection> startCommand(const std::string & command);
+
+    void startMaster();
+
+};
+
+}
diff --git a/src/libutil/pool.hh b/src/libutil/pool.hh
index f291cd578..3c3dd4b07 100644
--- a/src/libutil/pool.hh
+++ b/src/libutil/pool.hh
@@ -141,11 +141,16 @@ public:
         }
     }
 
-    unsigned int count()
+    size_t count()
     {
         auto state_(state.lock());
         return state_->idle.size() + state_->inUse;
     }
+
+    size_t capacity()
+    {
+        return state.lock()->max;
+    }
 };
 
 }