nix-store --serve: Use a versioned protocol

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

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();
+        if (cmd == "have") {
+            writeInt(qCmdHave, pipes.first);
             foreach (Strings::iterator, i, tokenized)
             writeStrings(tokenized, pipes.first);
             pipes.first.flush();
-        if (cmd == "have") {
             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);

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

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,13 +844,34 @@ static void opServe(Strings opFlags, Strings opArgs)
     FdSource in(STDIN_FILENO);
     FdSink out(STDOUT_FILENO);
 
-    string cmd = readString(in);
+    /* Exchange the greeting. */
-    if (cmd == "query") {
+    unsigned int magic = readInt(in);
-        for (cmd = readString(in); !cmd.empty(); cmd = readString(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);
-            if (cmd == "have") {
                             writeStrings(store->queryValidPaths(paths), out);
-            } else if (cmd == "info") {
+                        }
+                        break;
+                    case qCmdInfo:
+                        {
+                            PathSet paths = readStrings<PathSet>(in);
                             // !!! Maybe we want a queryPathInfos?
                             foreach (PathSet::iterator, i, paths) {
                                 if (!store->isValidPath(*i))
@@ -863,15 +885,21 @@ static void opServe(Strings opFlags, Strings opArgs)
                                 writeLongLong(info.narSize, out);
                             }
                             writeString("", out);
-            } else
+                        }
+                        break;
+                    default:
                         throw Error(format("Unknown serve query `%1%'") % cmd);
+                }
                 out.flush();
             }
-    } else if (cmd == "substitute")
+            break;
+        case cmdSubstitute:
             dumpPath(readString(in), out);
-    else
+            break;
+        default:
             throw Error(format("Unknown serve command `%1%'") % cmd);
     }
+}
 
 
 /* Scan the arguments; find the operation, set global flags, put all

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;
+
+}