From 960af8f793ba8a00edb978dde7127a8b07f5f18f Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Mon, 24 May 2021 00:15:38 -0700
Subject: [PATCH] Add deployment.privilegeEscalationCommand

This adds a NixOps-equivalent option for non-root deployment
on remote hosts.

Fixes #27.
---
 src/nix/eval.nix    | 10 ++++++++++
 src/nix/host/ssh.rs | 14 ++++++++++++++
 src/nix/mod.rs      |  4 ++++
 3 files changed, 28 insertions(+)

diff --git a/src/nix/eval.nix b/src/nix/eval.nix
index dae8372..704ea1c 100644
--- a/src/nix/eval.nix
+++ b/src/nix/eval.nix
@@ -155,6 +155,16 @@ let
           type = types.bool;
           default = true;
         };
+        privilegeEscalationCommand = lib.mkOption {
+          description = ''
+            Command to use to elevate privileges when activating the new profiles on SSH hosts.
+
+            This is used on SSH hosts when `deployment.targetUser` is not `root`.
+            The user must be allowed to use the command non-interactively.
+          '';
+          type = types.listOf types.str;
+          default = [ "sudo" "-H" "--" ];
+        };
       };
     };
   };
diff --git a/src/nix/host/ssh.rs b/src/nix/host/ssh.rs
index f61a559..4aa1d97 100644
--- a/src/nix/host/ssh.rs
+++ b/src/nix/host/ssh.rs
@@ -26,6 +26,9 @@ pub struct Ssh {
     /// Local path to a ssh_config file.
     ssh_config: Option<PathBuf>,
 
+    /// Command to elevate privileges with.
+    privilege_escalation_command: Vec<String>,
+
     friendly_name: String,
     progress_bar: TaskProgress,
     logs: String,
@@ -107,6 +110,7 @@ impl Ssh {
             port: None,
             ssh_config: None,
             friendly_name,
+            privilege_escalation_command: Vec::new(),
             progress_bar: TaskProgress::default(),
             logs: String::new(),
         }
@@ -120,6 +124,10 @@ impl Ssh {
         self.ssh_config = Some(ssh_config);
     }
 
+    pub fn set_privilege_escalation_command(&mut self, command: Vec<String>) {
+        self.privilege_escalation_command = command;
+    }
+
     pub fn upcast(self) -> Box<dyn Host> {
         Box::new(self)
     }
@@ -128,6 +136,11 @@ impl Ssh {
     pub fn ssh(&self, command: &[&str]) -> Command {
         let options = self.ssh_options();
         let options_str = options.join(" ");
+        let privilege_escalation_command = if self.user != "root" {
+            self.privilege_escalation_command.as_slice()
+        } else {
+            &[]
+        };
 
         let mut cmd = Command::new("ssh");
 
@@ -135,6 +148,7 @@ impl Ssh {
             .arg(self.ssh_target())
             .args(&options)
             .arg("--")
+            .args(privilege_escalation_command)
             .args(command)
             .env("NIX_SSHOPTS", options_str);
 
diff --git a/src/nix/mod.rs b/src/nix/mod.rs
index 740699f..fba1dfa 100644
--- a/src/nix/mod.rs
+++ b/src/nix/mod.rs
@@ -121,6 +121,9 @@ pub struct NodeConfig {
     #[serde(rename = "replaceUnknownProfiles")]
     replace_unknown_profiles: bool,
 
+    #[serde(rename = "privilegeEscalationCommand")]
+    privilege_escalation_command: Vec<String>,
+
     #[validate(custom = "validate_keys")]
     keys: HashMap<String, Key>,
 }
@@ -132,6 +135,7 @@ impl NodeConfig {
     pub fn to_ssh_host(&self) -> Option<Ssh> {
         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());
 
             if let Some(target_port) = self.target_port {
                 host.set_port(target_port);