Merge remote-tracking branch 'shlevy/ssh-substituter'

8 files changed, 271 insertions, 0 deletions

diff --git a/Makefile b/Makefile
index 0ad6bb236c82..78d85ebc58f2 100644
--- a/Makefile
+++ b/Makefile
@@ -10,6 +10,7 @@ makefiles = \
   src/nix-instantiate/local.mk \
   src/nix-env/local.mk \
   src/nix-daemon/local.mk \
+  src/download-via-ssh/local.mk \
   src/nix-log2xml/local.mk \
   src/bsdiff-4.3/local.mk \
   perl/local.mk \
diff --git a/src/download-via-ssh/download-via-ssh.cc b/src/download-via-ssh/download-via-ssh.cc
new file mode 100644
index 000000000000..d85f1572f33e
--- /dev/null
+++ b/src/download-via-ssh/download-via-ssh.cc
@@ -0,0 +1,138 @@
+#include "shared.hh"
+#include "util.hh"
+#include "serialise.hh"
+#include "archive.hh"
+#include "affinity.hh"
+#include "globals.hh"
+#include "serve-protocol.hh"
+
+#include <iostream>
+#include <unistd.h>
+
+using namespace nix;
+
+// !!! TODO:
+// * Respect more than the first host
+// * use a database
+// * show progress
+
+static std::pair<FdSink, FdSource> connect(string conn) {
+    Pipe to, from;
+    to.create();
+    from.create();
+    pid_t child = fork();
+    switch (child) {
+        case -1:
+            throw SysError("unable to fork");
+        case 0:
+            try {
+                restoreAffinity();
+                if (dup2(to.readSide, STDIN_FILENO) == -1)
+                    throw SysError("dupping stdin");
+                if (dup2(from.writeSide, STDOUT_FILENO) == -1)
+                    throw SysError("dupping stdout");
+                execlp("ssh"
+                      , "ssh"
+                      , "-x"
+                      , "-T"
+                      , conn.c_str()
+                      , "nix-store --serve"
+                      , NULL);
+                throw SysError("executing ssh");
+            } catch (std::exception & e) {
+                std::cerr << "error: " << e.what() << std::endl;
+            }
+            _exit(1);
+    }
+    // If child exits unexpectedly, we'll EPIPE or EOF early.
+    // If we exit unexpectedly, child will EPIPE or EOF early.
+    // So no need to keep track of it.
+
+    return std::pair<FdSink, FdSource>(to.writeSide.borrow(), from.readSide.borrow());
+}
+
+static void substitute(std::pair<FdSink, FdSource> & pipes, Path storePath, Path destPath) {
+    writeInt(cmdSubstitute, pipes.first);
+    writeString(storePath, pipes.first);
+    pipes.first.flush();
+    restorePath(destPath, pipes.second);
+    std::cout << std::endl;
+}
+
+static void query(std::pair<FdSink, FdSource> & pipes) {
+    writeInt(cmdQuery, pipes.first);
+    for (string line; getline(std::cin, line);) {
+        Strings tokenized = tokenizeString<Strings>(line);
+        string cmd = tokenized.front();
+        tokenized.pop_front();
+        if (cmd == "have") {
+            writeInt(qCmdHave, pipes.first);
+            writeStrings(tokenized, pipes.first);
+            pipes.first.flush();
+            PathSet paths = readStrings<PathSet>(pipes.second);
+            foreach (PathSet::iterator, i, paths)
+                std::cout << *i << std::endl;
+        } else if (cmd == "info") {
+            writeInt(qCmdInfo, pipes.first);
+            writeStrings(tokenized, pipes.first);
+            pipes.first.flush();
+            for (Path path = readString(pipes.second); !path.empty(); path = readString(pipes.second)) {
+                std::cout << path << std::endl;
+                std::cout << readString(pipes.second) << std::endl;
+                PathSet references = readStrings<PathSet>(pipes.second);
+                std::cout << references.size() << std::endl;
+                foreach (PathSet::iterator, i, references)
+                    std::cout << *i << std::endl;
+                std::cout << readLongLong(pipes.second) << std::endl;
+                std::cout << readLongLong(pipes.second) << std::endl;
+            }
+        } else
+            throw Error(format("unknown substituter query `%1%'") % cmd);
+        std::cout << std::endl;
+    }
+}
+
+void run(Strings args)
+{
+    if (args.empty())
+        throw UsageError("download-via-ssh requires an argument");
+
+    if (settings.sshSubstituterHosts.empty())
+        return;
+
+    std::cout << std::endl;
+
+    std::pair<FdSink, FdSource> pipes = connect(settings.sshSubstituterHosts.front());
+
+    /* Exchange the greeting */
+    writeInt(SERVE_MAGIC_1, pipes.first);
+    pipes.first.flush();
+    unsigned int magic = readInt(pipes.second);
+    if (magic != SERVE_MAGIC_2)
+        throw Error("protocol mismatch");
+    readInt(pipes.second); // Server version, unused for now
+    writeInt(SERVE_PROTOCOL_VERSION, pipes.first);
+    pipes.first.flush();
+
+    Strings::iterator i = args.begin();
+    if (*i == "--query")
+        query(pipes);
+    else if (*i == "--substitute")
+        if (args.size() != 3)
+            throw UsageError("download-via-ssh: --substitute takes exactly two arguments");
+        else {
+            Path storePath = *++i;
+            Path destPath = *++i;
+            substitute(pipes, storePath, destPath);
+        }
+    else
+        throw UsageError(format("download-via-ssh: unknown command `%1%'") % *i);
+}
+
+void printHelp()
+{
+    std::cerr << "Usage: download-via-ssh --query|--substitute store-path dest-path" << std::endl;
+}
+
+
+string programId = "download-via-ssh";
diff --git a/src/download-via-ssh/local.mk b/src/download-via-ssh/local.mk
new file mode 100644
index 000000000000..80f4c385acb3
--- /dev/null
+++ b/src/download-via-ssh/local.mk
@@ -0,0 +1,11 @@
+programs += download-via-ssh
+
+download-via-ssh_DIR := $(d)
+
+download-via-ssh_SOURCES := $(d)/download-via-ssh.cc
+
+download-via-ssh_INSTALL_DIR := $(libexecdir)/nix/substituters
+
+download-via-ssh_CXXFLAGS = -Isrc/nix-store
+
+download-via-ssh_LIBS = libmain libstore libutil libformat
diff --git a/src/libmain/shared.cc b/src/libmain/shared.cc
index fb70cb076732..30238c7fbe9a 100644
--- a/src/libmain/shared.cc
+++ b/src/libmain/shared.cc
@@ -223,6 +223,9 @@ static void initAndRun(int argc, char * * argv)
         else remaining.push_back(arg);
     }
 
+    if (char *pack = getenv("_NIX_OPTIONS"))
+        settings.unpack(pack);
+
     settings.update();
 
     run(remaining);
diff --git a/src/libstore/globals.cc b/src/libstore/globals.cc
index 68add1982f43..1d4bcd94f7b8 100644
--- a/src/libstore/globals.cc
+++ b/src/libstore/globals.cc
@@ -79,6 +79,7 @@ void Settings::processEnvironment()
 #endif
         substituters.push_back(nixLibexecDir + "/nix/substituters/download-using-manifests.pl");
         substituters.push_back(nixLibexecDir + "/nix/substituters/download-from-binary-cache.pl");
+        substituters.push_back(nixLibexecDir + "/nix/substituters/download-via-ssh");
     } else
         substituters = tokenizeString<Strings>(subs, ":");
 }
@@ -151,6 +152,7 @@ void Settings::update()
     get(gcKeepDerivations, "gc-keep-derivations");
     get(autoOptimiseStore, "auto-optimise-store");
     get(envKeepDerivations, "env-keep-derivations");
+    get(sshSubstituterHosts, "ssh-substituter-hosts");
 }
 
 
@@ -182,6 +184,13 @@ void Settings::get(StringSet & res, const string & name)
     res.insert(ss.begin(), ss.end());
 }
 
+void Settings::get(Strings & res, const string & name)
+{
+    SettingsMap::iterator i = settings.find(name);
+    if (i == settings.end()) return;
+    res = tokenizeString<Strings>(i->second);
+}
+
 
 template<class N> void Settings::get(N & res, const string & name)
 {
@@ -206,6 +215,17 @@ string Settings::pack()
 }
 
 
+void Settings::unpack(const string &pack) {
+    Strings lines = tokenizeString<Strings>(pack, "\n");
+    foreach (Strings::iterator, i, lines) {
+        string::size_type eq = i->find('=');
+        if (eq == string::npos)
+            throw Error("illegal option name/value");
+        set(i->substr(0, eq), i->substr(eq + 1));
+    }
+}
+
+
 Settings::SettingsMap Settings::getOverrides()
 {
     return overrides;
diff --git a/src/libstore/globals.hh b/src/libstore/globals.hh
index 9300edbe9615..13772c65cbac 100644
--- a/src/libstore/globals.hh
+++ b/src/libstore/globals.hh
@@ -25,6 +25,8 @@ struct Settings {
 
     string pack();
 
+    void unpack(const string &pack);
+
     SettingsMap getOverrides();
 
     /* The directory where we store sources and derived files. */
@@ -144,6 +146,9 @@ struct Settings {
        chroot. */
     StringSet dirsInChroot;
 
+    /* Set of ssh connection strings for the ssh substituter */
+    Strings sshSubstituterHosts;
+
     /* Whether to impersonate a Linux 2.6 machine on newer kernels. */
     bool impersonateLinux26;
 
@@ -195,6 +200,7 @@ private:
     void get(string & res, const string & name);
     void get(bool & res, const string & name);
     void get(StringSet & res, const string & name);
+    void get(Strings & res, const string & name);
     template<class N> void get(N & res, const string & name);
 };
 
diff --git a/src/nix-store/nix-store.cc b/src/nix-store/nix-store.cc
index 69a98fe4726d..350a4ce0fa37 100644
--- a/src/nix-store/nix-store.cc
+++ b/src/nix-store/nix-store.cc
@@ -6,6 +6,7 @@
 #include "xmlgraph.hh"
 #include "local-store.hh"
 #include "util.hh"
+#include "serve-protocol.hh"
 
 #include <iostream>
 #include <algorithm>
@@ -834,6 +835,71 @@ static void opClearFailedPaths(Strings opFlags, Strings opArgs)
 }
 
 
+// Serve the nix store in a way usable by a restricted ssh user
+static void opServe(Strings opFlags, Strings opArgs)
+{
+    if (!opArgs.empty() || !opFlags.empty())
+        throw UsageError("no arguments or flags expected");
+
+    FdSource in(STDIN_FILENO);
+    FdSink out(STDOUT_FILENO);
+
+    /* Exchange the greeting. */
+    unsigned int magic = readInt(in);
+    if (magic != SERVE_MAGIC_1) throw Error("protocol mismatch");
+    writeInt(SERVE_MAGIC_2, out);
+    writeInt(SERVE_PROTOCOL_VERSION, out);
+    out.flush();
+    readInt(in); // Client version, unused for now
+
+    ServeCommand cmd = (ServeCommand) readInt(in);
+    switch (cmd) {
+        case cmdQuery:
+            while (true) {
+                QueryCommand qCmd;
+                try {
+                    qCmd = (QueryCommand) readInt(in);
+                } catch (EndOfFile & e) {
+                    break;
+                }
+                switch (qCmd) {
+                    case qCmdHave: {
+                        PathSet paths = readStrings<PathSet>(in);
+                        writeStrings(store->queryValidPaths(paths), out);
+                        break;
+                    }
+                    case qCmdInfo: {
+                        PathSet paths = readStrings<PathSet>(in);
+                        // !!! Maybe we want a queryPathInfos?
+                        foreach (PathSet::iterator, i, paths) {
+                            if (!store->isValidPath(*i))
+                                continue;
+                            ValidPathInfo info = store->queryPathInfo(*i);
+                            writeString(info.path, out);
+                            writeString(info.deriver, out);
+                            writeStrings(info.references, out);
+                            // !!! Maybe we want compression?
+                            writeLongLong(info.narSize, out); // downloadSize
+                            writeLongLong(info.narSize, out);
+                        }
+                        writeString("", out);
+                        break;
+                    }
+                    default:
+                        throw Error(format("unknown serve query `%1%'") % cmd);
+                }
+                out.flush();
+            }
+            break;
+        case cmdSubstitute:
+            dumpPath(readString(in), out);
+            break;
+        default:
+            throw Error(format("unknown serve command `%1%'") % cmd);
+    }
+}
+
+
 /* Scan the arguments; find the operation, set global flags, put all
    other flags in a list, and put all other arguments in another
    list. */
@@ -904,6 +970,8 @@ void run(Strings args)
             indirectRoot = true;
         else if (arg == "--no-output")
             noOutput = true;
+        else if (arg == "--serve")
+            op = opServe;
         else if (arg[0] == '-') {
             opFlags.push_back(arg);
             if (arg == "--max-freed" || arg == "--max-links" || arg == "--max-atime") { /* !!! hack */
diff --git a/src/nix-store/serve-protocol.hh b/src/nix-store/serve-protocol.hh
new file mode 100644
index 000000000000..69277bc1b99b
--- /dev/null
+++ b/src/nix-store/serve-protocol.hh
@@ -0,0 +1,24 @@
+#pragma once
+
+namespace nix {
+
+
+#define SERVE_MAGIC_1 0x390c9deb
+#define SERVE_MAGIC_2 0x5452eecb
+
+#define SERVE_PROTOCOL_VERSION 0x101
+#define GET_PROTOCOL_MAJOR(x) ((x) & 0xff00)
+#define GET_PROTOCOL_MINOR(x) ((x) & 0x00ff)
+
+
+typedef enum {
+    cmdQuery = 0,
+    cmdSubstitute = 1,
+} ServeCommand;
+
+typedef enum {
+    qCmdHave = 0,
+    qCmdInfo = 1,
+} QueryCommand;
+
+}