nix-store --serve: Use a versioned protocol

Signed-off-by: Shea Levy <shea@shealevy.com>

4 files changed, 102 insertions, 34 deletions

diff --git a/src/download-via-ssh/download-via-ssh.cc b/src/download-via-ssh/download-via-ssh.cc
index 003d7de2bf..be70a374fe 100644
--- a/src/download-via-ssh/download-via-ssh.cc
+++ b/src/download-via-ssh/download-via-ssh.cc
@@ -4,6 +4,7 @@
 #include "archive.hh"
 #include "affinity.hh"
 #include "globals.hh"
+#include "serve-protocol.hh"
 
 #include <iostream>
 #include <unistd.h>
@@ -54,7 +55,7 @@ static pair<FdSink, FdSource> connect(string conn) {
 }
 
 static void substitute(pair<FdSink, FdSource> & pipes, Path storePath, Path destPath) {
-    writeString("substitute", pipes.first);
+    writeInt(cmdSubstitute, pipes.first);
     writeString(storePath, pipes.first);
     pipes.first.flush();
     restorePath(destPath, pipes.second);
@@ -63,20 +64,24 @@ static void substitute(pair<FdSink, FdSource> & pipes, Path storePath, Path dest
 
 static void query(pair<FdSink, FdSource> & pipes) {
     using std::cin;
-    writeString("query", pipes.first);
+    writeInt(cmdQuery, pipes.first);
     for (string line; getline(cin, line);) {
         Strings tokenized = tokenizeString<Strings>(line);
         string cmd = tokenized.front();
-        writeString(cmd, pipes.first);
         tokenized.pop_front();
-        foreach (Strings::iterator, i, tokenized)
-        writeStrings(tokenized, pipes.first);
-        pipes.first.flush();
         if (cmd == "have") {
+            writeInt(qCmdHave, pipes.first);
+            foreach (Strings::iterator, i, tokenized)
+            writeStrings(tokenized, pipes.first);
+            pipes.first.flush();
             PathSet paths = readStrings<PathSet>(pipes.second);
             foreach (PathSet::iterator, i, paths)
                 cout << *i << endl;
         } else if (cmd == "info") {
+            writeInt(qCmdInfo, pipes.first);
+            foreach (Strings::iterator, i, tokenized)
+            writeStrings(tokenized, pipes.first);
+            pipes.first.flush();
             for (Path path = readString(pipes.second); !path.empty(); path = readString(pipes.second)) {
                 cout << path << endl;
                 cout << readString(pipes.second) << endl;
@@ -91,7 +96,6 @@ static void query(pair<FdSink, FdSource> & pipes) {
             throw Error(format("Unknown substituter query `%1%'") % cmd);
         cout << endl;
     }
-    writeString("", pipes.first);
 }
 
 void run(Strings args)
@@ -106,6 +110,16 @@ void run(Strings args)
 
     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);
diff --git a/src/download-via-ssh/local.mk b/src/download-via-ssh/local.mk
index 92bf115946..80f4c385ac 100644
--- a/src/download-via-ssh/local.mk
+++ b/src/download-via-ssh/local.mk
@@ -6,4 +6,6 @@ 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/nix-store/nix-store.cc b/src/nix-store/nix-store.cc
index 68ad902674..638d244988 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>
@@ -843,34 +844,61 @@ static void opServe(Strings opFlags, Strings opArgs)
     FdSource in(STDIN_FILENO);
     FdSink out(STDOUT_FILENO);
 
-    string cmd = readString(in);
-    if (cmd == "query") {
-        for (cmd = readString(in); !cmd.empty(); cmd = readString(in)) {
-            PathSet paths = readStrings<PathSet>(in);
-            if (cmd == "have") {
-                writeStrings(store->queryValidPaths(paths), out);
-            } else if (cmd == "info") {
-                // !!! 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);
+    /* 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;
                 }
-                writeString("", out);
-            } else
-                throw Error(format("Unknown serve query `%1%'") % cmd);
-            out.flush();
-        }
-    } else if (cmd == "substitute")
-        dumpPath(readString(in), out);
-    else
-        throw Error(format("Unknown serve command `%1%'") % cmd);
+                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);
+    }
 }
 
 
diff --git a/src/nix-store/serve-protocol.hh b/src/nix-store/serve-protocol.hh
new file mode 100644
index 0000000000..69277bc1b9
--- /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;
+
+}