From ea9df6fe51cf52a8e9106cd62a195bf185e6c5f6 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Thu, 23 Sep 2021 18:01:04 +0200
Subject: [PATCH] Shut down write side before draining the read side

This is important if the remote side *does* execute
nix-store/nix-daemon successfully, but stdout is polluted
(e.g. because the remote user's bashrc script prints something to
stdout). In that case we have to shutdown the write side to force the
remote nix process to exit.
---
 src/libstore/legacy-ssh-store.cc | 2 ++
 src/libstore/remote-store.cc     | 3 +++
 src/libstore/remote-store.hh     | 3 ++-
 src/libstore/ssh-store.cc        | 5 +++++
 src/libstore/uds-remote-store.cc | 6 ++++++
 src/libstore/uds-remote-store.hh | 6 ++++++
 6 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index 5d006f6e2..c660d7f5b 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -89,6 +89,8 @@ struct LegacySSHStore : public virtual LegacySSHStoreConfig, public virtual Stor
                 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'",
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index c71eb4631..b57a4cb05 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -169,6 +169,9 @@ void RemoteStore::initConnection(Connection & conn)
             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));
         }
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;
 };