From 4a5ca576da511fcc64039c2494f41f710d662478 Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Sun, 20 Feb 2022 19:24:07 +0000
Subject: [PATCH] Factor out `ServeProto::BasicClientConnection::handshake`
 Hydra to share

---
 src/libstore/legacy-ssh-store.cc    | 11 ++---------
 src/libstore/serve-protocol-impl.cc | 19 +++++++++++++++++++
 src/libstore/serve-protocol-impl.hh | 22 ++++++++++++++++++++++
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index 058b1affd..4f020c452 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -64,15 +64,8 @@ ref<LegacySSHStore::Connection> LegacySSHStore::openConnection()
     StringSink saved;
     TeeSource tee(conn->from, saved);
     try {
-        conn->to << SERVE_MAGIC_1 << SERVE_PROTOCOL_VERSION;
-        conn->to.flush();
-
-        unsigned int magic = readInt(conn->from);
-        if (magic != SERVE_MAGIC_2)
-            throw Error("'nix-store --serve' protocol mismatch from '%s'", host);
-        conn->remoteVersion = readInt(conn->from);
-        if (GET_PROTOCOL_MAJOR(conn->remoteVersion) != 0x200)
-            throw Error("unsupported 'nix-store --serve' protocol version on '%s'", host);
+        conn->remoteVersion = ServeProto::BasicClientConnection::handshake(
+            conn->to, tee, SERVE_PROTOCOL_VERSION, host);
     } catch (SerialisationError & e) {
         // in.close(): Don't let the remote block on us not writing.
         conn->sshConn->in.close();
diff --git a/src/libstore/serve-protocol-impl.cc b/src/libstore/serve-protocol-impl.cc
index e65de7650..6bf6c8cf6 100644
--- a/src/libstore/serve-protocol-impl.cc
+++ b/src/libstore/serve-protocol-impl.cc
@@ -4,6 +4,25 @@
 
 namespace nix {
 
+ServeProto::Version ServeProto::BasicClientConnection::handshake(
+    BufferedSink & to,
+    Source & from,
+    ServeProto::Version localVersion,
+    std::string_view host)
+{
+    to << SERVE_MAGIC_1 << localVersion;
+    to.flush();
+
+    unsigned int magic = readInt(from);
+    if (magic != SERVE_MAGIC_2)
+        throw Error("'nix-store --serve' protocol mismatch from '%s'", host);
+    auto remoteVersion = readInt(from);
+    if (GET_PROTOCOL_MAJOR(remoteVersion) != 0x200)
+        throw Error("unsupported 'nix-store --serve' protocol version on '%s'", host);
+    return remoteVersion;
+}
+
+
 StorePathSet ServeProto::BasicClientConnection::queryValidPaths(
     const Store & store,
     bool lock, const StorePathSet & paths,
diff --git a/src/libstore/serve-protocol-impl.hh b/src/libstore/serve-protocol-impl.hh
index 312f5d47a..8cd241fd3 100644
--- a/src/libstore/serve-protocol-impl.hh
+++ b/src/libstore/serve-protocol-impl.hh
@@ -63,6 +63,28 @@ struct ServeProto::BasicClientConnection
     FdSource from;
     ServeProto::Version remoteVersion;
 
+    /**
+     * Establishes connection, negotiating version.
+     *
+     * @return the version provided by the other side of the
+     * connection.
+     *
+     * @param to Taken by reference to allow for various error handling
+     * mechanisms.
+     *
+     * @param from Taken by reference to allow for various error
+     * handling mechanisms.
+     *
+     * @param localVersion Our version which is sent over
+     *
+     * @param host Just used to add context to thrown exceptions.
+     */
+    static ServeProto::Version handshake(
+        BufferedSink & to,
+        Source & from,
+        ServeProto::Version localVersion,
+        std::string_view host);
+
     /**
      * Coercion to `ServeProto::ReadConn`. This makes it easy to use the
      * factored out serve protocol serializers with a