6 files changed, 46 insertions, 6 deletions

diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index 45eed5707..c660d7f5b 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -82,9 +82,20 @@ struct LegacySSHStore : public virtual LegacySSHStoreConfig, public virtual Stor
             conn->to << SERVE_MAGIC_1 << SERVE_PROTOCOL_VERSION;
             conn->to.flush();
 
-            unsigned int magic = readInt(conn->from);
-            if (magic != SERVE_MAGIC_2)
-                throw Error("protocol mismatch with 'nix-store --serve' on '%s'", host);
+            StringSink saved;
+            try {
+                TeeSource tee(conn->from, saved);
+                unsigned int magic = readInt(tee);
+                if (magic != SERVE_MAGIC_2)
+                    throw Error("'nix-store --serve' protocol mismatch from '%s'", host);
+            } catch (SerialisationError & e) {
+                /* In case the other side is waiting for our input,
+                   close it. */
+                conn->sshConn->in.close();
+                auto msg = conn->from.drain();
+                throw Error("'nix-store --serve' protocol mismatch from '%s', got '%s'",
+                    host, chomp(*saved.s + msg));
+            }
             conn->remoteVersion = readInt(conn->from);
             if (GET_PROTOCOL_MAJOR(conn->remoteVersion) != 0x200)
                 throw Error("unsupported 'nix-store --serve' protocol version on '%s'", host);
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index 73f590e7b..b57a4cb05 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -162,8 +162,19 @@ void RemoteStore::initConnection(Connection & conn)
     try {
         conn.to << WORKER_MAGIC_1;
         conn.to.flush();
-        unsigned int magic = readInt(conn.from);
-        if (magic != WORKER_MAGIC_2) throw Error("protocol mismatch");
+        StringSink saved;
+        try {
+            TeeSource tee(conn.from, saved);
+            unsigned int magic = readInt(tee);
+            if (magic != WORKER_MAGIC_2)
+                throw Error("protocol mismatch");
+        } catch (SerialisationError & e) {
+            /* In case the other side is waiting for our input, close
+               it. */
+            conn.closeWrite();
+            auto msg = conn.from.drain();
+            throw Error("protocol mismatch, got '%s'", chomp(*saved.s + msg));
+        }
 
         conn.from >> conn.daemonVersion;
         if (GET_PROTOCOL_MAJOR(conn.daemonVersion) != GET_PROTOCOL_MAJOR(PROTOCOL_VERSION))
diff --git a/src/libstore/remote-store.hh b/src/libstore/remote-store.hh
index 8901c79fc..ac1eaa19e 100644
--- a/src/libstore/remote-store.hh
+++ b/src/libstore/remote-store.hh
@@ -125,7 +125,6 @@ public:
 
     struct Connection
     {
-        AutoCloseFD fd;
         FdSink to;
         FdSource from;
         unsigned int daemonVersion;
@@ -133,6 +132,8 @@ public:
 
         virtual ~Connection();
 
+        virtual void closeWrite() = 0;
+
         std::exception_ptr processStderr(Sink * sink = 0, Source * source = 0, bool flush = true);
     };
 
diff --git a/src/libstore/ssh-store.cc b/src/libstore/ssh-store.cc
index f2caf2aeb..bb03daef4 100644
--- a/src/libstore/ssh-store.cc
+++ b/src/libstore/ssh-store.cc
@@ -57,6 +57,11 @@ private:
     struct Connection : RemoteStore::Connection
     {
         std::unique_ptr<SSHMaster::Connection> sshConn;
+
+        void closeWrite() override
+        {
+            sshConn->in.close();
+        }
     };
 
     ref<RemoteStore::Connection> openConnection() override;
diff --git a/src/libstore/uds-remote-store.cc b/src/libstore/uds-remote-store.cc
index cac4fa036..cfadccf68 100644
--- a/src/libstore/uds-remote-store.cc
+++ b/src/libstore/uds-remote-store.cc
@@ -45,6 +45,12 @@ std::string UDSRemoteStore::getUri()
 }
 
 
+void UDSRemoteStore::Connection::closeWrite()
+{
+    shutdown(fd.get(), SHUT_WR);
+}
+
+
 ref<RemoteStore::Connection> UDSRemoteStore::openConnection()
 {
     auto conn = make_ref<Connection>();
diff --git a/src/libstore/uds-remote-store.hh b/src/libstore/uds-remote-store.hh
index ddc7716cd..f8dfcca70 100644
--- a/src/libstore/uds-remote-store.hh
+++ b/src/libstore/uds-remote-store.hh
@@ -40,6 +40,12 @@ public:
 
 private:
 
+    struct Connection : RemoteStore::Connection
+    {
+        AutoCloseFD fd;
+        void closeWrite() override;
+    };
+
     ref<RemoteStore::Connection> openConnection() override;
     std::optional<std::string> path;
 };