distributed builds: load remote builder host key from the machines file

This is already used by Hydra, and is very useful when materializing
a remote builder list from service discovery. This allows the service
discovery tool to only sync one file instead of two.

5 files changed, 26 insertions, 3 deletions

diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index 253c0033e..99b0bb5a8 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -15,6 +15,7 @@ struct LegacySSHStoreConfig : virtual StoreConfig
     using StoreConfig::StoreConfig;
     const Setting<int> maxConnections{(StoreConfig*) this, 1, "max-connections", "maximum number of concurrent SSH connections"};
     const Setting<Path> sshKey{(StoreConfig*) this, "", "ssh-key", "path to an SSH private key"};
+    const Setting<std::string> sshPublicHostKey{(StoreConfig*) this, "", "base64-ssh-public-host-key", "The public half of the host's SSH key"};
     const Setting<bool> compress{(StoreConfig*) this, false, "compress", "whether to compress the connection"};
     const Setting<Path> remoteProgram{(StoreConfig*) this, "nix-store", "remote-program", "path to the nix-store executable on the remote system"};
     const Setting<std::string> remoteStore{(StoreConfig*) this, "", "remote-store", "URI of the store on the remote system"};
@@ -59,6 +60,7 @@ struct LegacySSHStore : public virtual LegacySSHStoreConfig, public virtual Stor
         , master(
             host,
             sshKey,
+            sshPublicHostKey,
             // Use SSH master only if using more than 1 connection.
             connections->capacity() > 1,
             compress,
diff --git a/src/libstore/machines.cc b/src/libstore/machines.cc
index 7db2556f4..b42e5e434 100644
--- a/src/libstore/machines.cc
+++ b/src/libstore/machines.cc
@@ -54,9 +54,15 @@ ref<Store> Machine::openStore() const {
     if (hasPrefix(storeUri, "ssh://")) {
         storeParams["max-connections"] = "1";
         storeParams["log-fd"] = "4";
+    }
+
+    if (hasPrefix(storeUri, "ssh://") || hasPrefix(storeUri, "ssh-ng://")) {
         if (sshKey != "")
             storeParams["ssh-key"] = sshKey;
+        if (sshPublicHostKey != "")
+            storeParams["base64-ssh-public-host-key"] = sshPublicHostKey;
     }
+
     {
         auto & fs = storeParams["system-features"];
         auto append = [&](auto feats) {
diff --git a/src/libstore/ssh-store.cc b/src/libstore/ssh-store.cc
index 17c258201..f2caf2aeb 100644
--- a/src/libstore/ssh-store.cc
+++ b/src/libstore/ssh-store.cc
@@ -13,6 +13,7 @@ struct SSHStoreConfig : virtual RemoteStoreConfig
     using RemoteStoreConfig::RemoteStoreConfig;
 
     const Setting<Path> sshKey{(StoreConfig*) this, "", "ssh-key", "path to an SSH private key"};
+    const Setting<std::string> sshPublicHostKey{(StoreConfig*) this, "", "base64-ssh-public-host-key", "The public half of the host's SSH key"};
     const Setting<bool> compress{(StoreConfig*) this, false, "compress", "whether to compress the connection"};
     const Setting<Path> remoteProgram{(StoreConfig*) this, "nix-daemon", "remote-program", "path to the nix-daemon executable on the remote system"};
     const Setting<std::string> remoteStore{(StoreConfig*) this, "", "remote-store", "URI of the store on the remote system"};
@@ -34,6 +35,7 @@ public:
         , master(
             host,
             sshKey,
+            sshPublicHostKey,
             // Use SSH master only if using more than 1 connection.
             connections->capacity() > 1,
             compress)
diff --git a/src/libstore/ssh.cc b/src/libstore/ssh.cc
index 84548a6e4..235eed37a 100644
--- a/src/libstore/ssh.cc
+++ b/src/libstore/ssh.cc
@@ -2,24 +2,37 @@
 
 namespace nix {
 
-SSHMaster::SSHMaster(const std::string & host, const std::string & keyFile, bool useMaster, bool compress, int logFD)
+SSHMaster::SSHMaster(const std::string & host, const std::string & keyFile, const std::string & sshPublicHostKey, bool useMaster, bool compress, int logFD)
     : host(host)
     , fakeSSH(host == "localhost")
     , keyFile(keyFile)
+    , sshPublicHostKey(sshPublicHostKey)
     , useMaster(useMaster && !fakeSSH)
     , compress(compress)
     , logFD(logFD)
 {
     if (host == "" || hasPrefix(host, "-"))
         throw Error("invalid SSH host name '%s'", host);
+
+    auto state(state_.lock());
+    state->tmpDir = std::make_unique<AutoDelete>(createTempDir("", "nix", true, true, 0700));
 }
 
 void SSHMaster::addCommonSSHOpts(Strings & args)
 {
+    auto state(state_.lock());
+
     for (auto & i : tokenizeString<Strings>(getEnv("NIX_SSHOPTS").value_or("")))
         args.push_back(i);
     if (!keyFile.empty())
         args.insert(args.end(), {"-i", keyFile});
+    if (!sshPublicHostKey.empty()) {
+        Path fileName = (Path) *state->tmpDir + "/host-key";
+        auto p = host.rfind("@");
+        string thost = p != string::npos ? string(host, p + 1) : host;
+        writeFile(fileName, thost + " " + base64Decode(sshPublicHostKey) + "\n");
+        args.insert(args.end(), {"-oUserKnownHostsFile=" + fileName});
+    }
     if (compress)
         args.push_back("-C");
 }
@@ -87,7 +100,6 @@ Path SSHMaster::startMaster()
 
     if (state->sshMaster != -1) return state->socketPath;
 
-    state->tmpDir = std::make_unique<AutoDelete>(createTempDir("", "nix", true, true, 0700));
 
     state->socketPath = (Path) *state->tmpDir + "/ssh.sock";
 
diff --git a/src/libstore/ssh.hh b/src/libstore/ssh.hh
index 4f0f0bd29..dabbcedda 100644
--- a/src/libstore/ssh.hh
+++ b/src/libstore/ssh.hh
@@ -12,6 +12,7 @@ private:
     const std::string host;
     bool fakeSSH;
     const std::string keyFile;
+    const std::string sshPublicHostKey;
     const bool useMaster;
     const bool compress;
     const int logFD;
@@ -29,7 +30,7 @@ private:
 
 public:
 
-    SSHMaster(const std::string & host, const std::string & keyFile, bool useMaster, bool compress, int logFD = -1);
+    SSHMaster(const std::string & host, const std::string & keyFile, const std::string & sshPublicHostKey, bool useMaster, bool compress, int logFD = -1);
 
     struct Connection
     {