From ecba88de9367f733610121fc0153310f92e05b65 Mon Sep 17 00:00:00 2001
From: Shea Levy <shea@shealevy.com>
Date: Fri, 2 Sep 2016 14:31:38 -0400
Subject: [PATCH] Add ssh store implementation

---
 src/libstore/remote-store.hh |  8 ++--
 src/libstore/ssh-store.cc    | 92 ++++++++++++++++++++++++++++++++++++
 src/libstore/ssh-store.hh    | 40 ++++++++++++++++
 3 files changed, 136 insertions(+), 4 deletions(-)
 create mode 100644 src/libstore/ssh-store.cc
 create mode 100644 src/libstore/ssh-store.hh

diff --git a/src/libstore/remote-store.hh b/src/libstore/remote-store.hh
index a69a4f2a3..5c9c617d9 100644
--- a/src/libstore/remote-store.hh
+++ b/src/libstore/remote-store.hh
@@ -90,20 +90,20 @@ protected:
         FdSource from;
         unsigned int daemonVersion;
 
-        ~Connection();
+        virtual ~Connection();
 
         void processStderr(Sink * sink = 0, Source * source = 0);
     };
 
     virtual ref<Connection> openConnection() = 0;
 
-    void setOptions(Connection & conn);
-
     void initConnection(Connection & conn);
 
+    ref<Pool<Connection>> connections;
+
 private:
 
-    ref<Pool<Connection>> connections;
+    void setOptions(Connection & conn);
 };
 
 class UDSRemoteStore : public LocalFSStore, public RemoteStore
diff --git a/src/libstore/ssh-store.cc b/src/libstore/ssh-store.cc
new file mode 100644
index 000000000..e1eb76659
--- /dev/null
+++ b/src/libstore/ssh-store.cc
@@ -0,0 +1,92 @@
+#include "ssh-store.hh"
+#include "remote-fs-accessor.hh"
+#include "archive.hh"
+#include "worker-protocol.hh"
+#include "pool.hh"
+
+namespace nix {
+
+SSHStore::SSHStore(string uri, const Params & params, size_t maxConnections)
+    : Store(params)
+    , RemoteStore(params, maxConnections)
+    , tmpDir(createTempDir("", "nix", true, true, 0700))
+    , socketPath((Path) tmpDir + "/ssh.sock")
+    , sshMaster(startProcess([&]() {
+                auto key = get(params, "ssh-key", "");
+                if (key.empty())
+                    execlp("ssh", "ssh", "-N", "-M", "-S", socketPath.c_str(), uri.c_str(), NULL);
+                else
+                    execlp("ssh", "ssh", "-N", "-M", "-S", socketPath.c_str(), "-i", key.c_str(), uri.c_str(), NULL);
+                throw SysError("starting ssh master");
+      }))
+    , uri(std::move(uri))
+{
+}
+
+string SSHStore::getUri()
+{
+    return "ssh://" + uri;
+}
+
+class ForwardSource : public Source
+{
+    Source & readSource;
+    Sink & writeSink;
+public:
+    ForwardSource(Source & readSource, Sink & writeSink) : readSource(readSource), writeSink(writeSink) {}
+    size_t read(unsigned char * data, size_t len) override
+    {
+        auto res = readSource.read(data, len);
+        writeSink(data, len);
+        return res;
+    }
+};
+
+void SSHStore::narFromPath(const Path & path, Sink & sink)
+{
+    auto conn(connections->get());
+    conn->to << wopNarFromPath << path;
+    conn->processStderr();
+    ParseSink ps;
+    auto fwd = ForwardSource(conn->from, sink);
+    parseDump(ps, fwd);
+}
+
+ref<FSAccessor> SSHStore::getFSAccessor()
+{
+    return make_ref<RemoteFSAccessor>(ref<Store>(shared_from_this()));
+}
+
+ref<RemoteStore::Connection> SSHStore::openConnection()
+{
+    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(), uri.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());
+    initConnection(*conn);
+    return conn;
+}
+
+static RegisterStoreImplementation regStore([](
+    const std::string & uri, const Store::Params & params)
+    -> std::shared_ptr<Store>
+{
+    if (std::string(uri, 0, 6) != "ssh://") return 0;
+    return std::make_shared<SSHStore>(uri.substr(6), params);
+});
+
+}
diff --git a/src/libstore/ssh-store.hh b/src/libstore/ssh-store.hh
new file mode 100644
index 000000000..44ece5598
--- /dev/null
+++ b/src/libstore/ssh-store.hh
@@ -0,0 +1,40 @@
+#pragma once
+
+#include "store-api.hh"
+#include "remote-store.hh"
+
+namespace nix {
+
+class SSHStore : public RemoteStore
+{
+public:
+
+    SSHStore(string uri, const Params & params, size_t maxConnections = std::numeric_limits<size_t>::max());
+
+    std::string getUri() override;
+
+    void narFromPath(const Path & path, Sink & sink) override;
+
+    ref<FSAccessor> getFSAccessor() override;
+
+private:
+
+    struct Connection : RemoteStore::Connection
+    {
+        Pid sshPid;
+        AutoCloseFD out;
+        AutoCloseFD in;
+    };
+
+    ref<RemoteStore::Connection> openConnection() override;
+
+    AutoDelete tmpDir;
+
+    Path socketPath;
+
+    Pid sshMaster;
+
+    string uri;
+};
+
+}