From afa7439c58ad9903cae7ad39b3274418dc8f00a8 Mon Sep 17 00:00:00 2001
From: minhuw <wangmh15@gmail.com>
Date: Tue, 27 Feb 2024 10:54:43 +0800
Subject: [PATCH] Support adding extra ssh options on node configurations

---
 src/nix/hive/options.nix | 7 +++++++
 src/nix/host/ssh.rs      | 9 +++++++++
 src/nix/mod.rs           | 4 ++++
 src/nix/node_filter.rs   | 1 +
 4 files changed, 21 insertions(+)

diff --git a/src/nix/hive/options.nix b/src/nix/hive/options.nix
index 8c0feba..95dcf1a 100644
--- a/src/nix/hive/options.nix
+++ b/src/nix/hive/options.nix
@@ -208,6 +208,13 @@ with builtins; rec {
           type = types.listOf types.str;
           default = [ "sudo" "-H" "--" ];
         };
+        sshOptions = lib.mkOption {
+          description = mdDoc ''
+            Extra SSH options to pass to the SSH command.
+          '';
+          type = types.listOf types.str;
+          default = [];
+        };
       };
     };
   };
diff --git a/src/nix/host/ssh.rs b/src/nix/host/ssh.rs
index 6ce6226..64ebb68 100644
--- a/src/nix/host/ssh.rs
+++ b/src/nix/host/ssh.rs
@@ -33,6 +33,9 @@ pub struct Ssh {
     /// Command to elevate privileges with.
     privilege_escalation_command: Vec<String>,
 
+    /// extra SSH options
+    extra_ssh_options: Vec<String>,
+
     /// Whether to use the experimental `nix copy` command.
     use_nix3_copy: bool,
 
@@ -189,6 +192,7 @@ impl Ssh {
             port: None,
             ssh_config: None,
             privilege_escalation_command: Vec::new(),
+            extra_ssh_options: Vec::new(),
             use_nix3_copy: false,
             job: None,
         }
@@ -206,6 +210,10 @@ impl Ssh {
         self.privilege_escalation_command = command;
     }
 
+    pub fn set_extra_ssh_options(&mut self, options: Vec<String>) {
+        self.extra_ssh_options = options;
+    }
+
     pub fn set_use_nix3_copy(&mut self, enable: bool) {
         self.use_nix3_copy = enable;
     }
@@ -346,6 +354,7 @@ impl Ssh {
         ]
         .iter()
         .map(|s| s.to_string())
+        .chain(self.extra_ssh_options.clone())
         .collect();
 
         if let Some(port) = self.port {
diff --git a/src/nix/mod.rs b/src/nix/mod.rs
index d06cbd5..8479ce6 100644
--- a/src/nix/mod.rs
+++ b/src/nix/mod.rs
@@ -78,6 +78,9 @@ pub struct NodeConfig {
     #[serde(rename = "privilegeEscalationCommand")]
     privilege_escalation_command: Vec<String>,
 
+    #[serde(rename = "sshOptions")]
+    extra_ssh_options: Vec<String>,
+
     #[validate(custom = "validate_keys")]
     keys: HashMap<String, Key>,
 }
@@ -181,6 +184,7 @@ impl NodeConfig {
         self.target_host.as_ref().map(|target_host| {
             let mut host = Ssh::new(self.target_user.clone(), target_host.clone());
             host.set_privilege_escalation_command(self.privilege_escalation_command.clone());
+            host.set_extra_ssh_options(self.extra_ssh_options.clone());
 
             if let Some(target_port) = self.target_port {
                 host.set_port(target_port);
diff --git a/src/nix/node_filter.rs b/src/nix/node_filter.rs
index 0e99801..6e459fa 100644
--- a/src/nix/node_filter.rs
+++ b/src/nix/node_filter.rs
@@ -248,6 +248,7 @@ mod tests {
             build_on_target: false,
             replace_unknown_profiles: false,
             privilege_escalation_command: vec![],
+            extra_ssh_options: vec![],
             keys: HashMap::new(),
         };