From 8b369f90fd2422a03310a0d3c623668e86e00e01 Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Mon, 20 May 2024 17:41:12 -0400
Subject: [PATCH] Query path infos (plural) and handshake version minimum for
 hydra

1. Hydra currently queries for multiple path infos at once, so let us
   make a connection item for that.

2. The minimum of the two versions should always be used, see #9584.
   (The issue remains open because the daemon protocol needs to be
   likewise updated.)
---
 src/libstore/legacy-ssh-store.cc      | 32 ++++++++++++++-------------
 src/libstore/serve-protocol-impl.cc   | 29 ++++++++++++++++++++++--
 src/libstore/serve-protocol-impl.hh   |  4 ++++
 tests/unit/libstore/serve-protocol.cc |  3 ++-
 4 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index e422adeec..9bc986bfa 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -105,24 +105,26 @@ void LegacySSHStore::queryPathInfoUncached(const StorePath & path,
 
         debug("querying remote host '%s' for info on '%s'", host, printStorePath(path));
 
-        conn->to << ServeProto::Command::QueryPathInfos << PathSet{printStorePath(path)};
-        conn->to.flush();
+        auto infos = conn->queryPathInfos(*this, {path});
 
-        auto p = readString(conn->from);
-        if (p.empty()) return callback(nullptr);
-        auto path2 = parseStorePath(p);
-        assert(path == path2);
-        auto info = std::make_shared<ValidPathInfo>(
-            path,
-            ServeProto::Serialise<UnkeyedValidPathInfo>::read(*this, *conn));
+        switch (infos.size()) {
+        case 0:
+            return callback(nullptr);
+        case 1: {
+            auto & [path2, info] = *infos.begin();
 
-        if (info->narHash == Hash::dummy)
-            throw Error("NAR hash is now mandatory");
+            if (info.narHash == Hash::dummy)
+                throw Error("NAR hash is now mandatory");
 
-        auto s = readString(conn->from);
-        assert(s == "");
-
-        callback(std::move(info));
+            assert(path == path2);
+            return callback(std::make_shared<ValidPathInfo>(
+                std::move(path),
+                std::move(info)
+            ));
+        }
+        default:
+            throw Error("More path infos returned than queried");
+        }
     } catch (...) { callback.rethrow(); }
 }
 
diff --git a/src/libstore/serve-protocol-impl.cc b/src/libstore/serve-protocol-impl.cc
index b39212884..5e1d3f1e8 100644
--- a/src/libstore/serve-protocol-impl.cc
+++ b/src/libstore/serve-protocol-impl.cc
@@ -19,7 +19,7 @@ ServeProto::Version ServeProto::BasicClientConnection::handshake(
     auto remoteVersion = readInt(from);
     if (GET_PROTOCOL_MAJOR(remoteVersion) != 0x200)
         throw Error("unsupported 'nix-store --serve' protocol version on '%s'", host);
-    return remoteVersion;
+    return std::min(remoteVersion, localVersion);
 }
 
 ServeProto::Version ServeProto::BasicServerConnection::handshake(
@@ -31,7 +31,8 @@ ServeProto::Version ServeProto::BasicServerConnection::handshake(
     if (magic != SERVE_MAGIC_1) throw Error("protocol mismatch");
     to << SERVE_MAGIC_2 << localVersion;
     to.flush();
-    return readInt(from);
+    auto remoteVersion = readInt(from);
+    return std::min(remoteVersion, localVersion);
 }
 
 
@@ -51,6 +52,30 @@ StorePathSet ServeProto::BasicClientConnection::queryValidPaths(
 }
 
 
+std::map<StorePath, UnkeyedValidPathInfo> ServeProto::BasicClientConnection::queryPathInfos(
+    const Store & store,
+    const StorePathSet & paths)
+{
+    std::map<StorePath, UnkeyedValidPathInfo> infos;
+
+    to << ServeProto::Command::QueryPathInfos;
+    ServeProto::write(store, *this, paths);
+    to.flush();
+
+    while (true) {
+        auto storePathS = readString(from);
+        if (storePathS == "") break;
+
+        auto storePath = store.parseStorePath(storePathS);
+        assert(paths.count(storePath) == 1);
+        auto info = ServeProto::Serialise<UnkeyedValidPathInfo>::read(store, *this);
+        infos.insert_or_assign(std::move(storePath), std::move(info));
+    }
+
+    return infos;
+}
+
+
 void ServeProto::BasicClientConnection::putBuildDerivationRequest(
     const Store & store,
     const StorePath & drvPath, const BasicDerivation & drv,
diff --git a/src/libstore/serve-protocol-impl.hh b/src/libstore/serve-protocol-impl.hh
index fd8d94697..1d9c79ef0 100644
--- a/src/libstore/serve-protocol-impl.hh
+++ b/src/libstore/serve-protocol-impl.hh
@@ -122,6 +122,10 @@ struct ServeProto::BasicClientConnection
         bool lock, const StorePathSet & paths,
         SubstituteFlag maybeSubstitute);
 
+    std::map<StorePath, UnkeyedValidPathInfo> queryPathInfos(
+        const Store & store,
+        const StorePathSet & paths);
+
     /**
      * Just the request half, because Hydra may do other things between
      * issuing the request and reading the `BuildResult` response.
diff --git a/tests/unit/libstore/serve-protocol.cc b/tests/unit/libstore/serve-protocol.cc
index b2fd0fb82..61dc15528 100644
--- a/tests/unit/libstore/serve-protocol.cc
+++ b/tests/unit/libstore/serve-protocol.cc
@@ -505,7 +505,8 @@ TEST_F(ServeProtoTest, handshake_client_corrupted_throws)
             } else {
                 auto ver = ServeProto::BasicClientConnection::handshake(
                     nullSink, in, defaultVersion, "blah");
-                EXPECT_NE(ver, defaultVersion);
+                // `std::min` of this and the other version saves us
+                EXPECT_EQ(ver, defaultVersion);
             }
         }
     });