apply-local: Escalate privileges only during activation

Fixes #85.

manual/src/release-notes.md

@@ -4,6 +4,7 @@
 
 - `--reboot` is added to trigger a reboot and wait for the node to come back up.
 - The target user is no longer explicitly set when `deployment.targetUser` is null ([#91](https://github.com/zhaofengli/colmena/pull/91)).
+- In `apply-local`, we now only escalate privileges during activation ([#85](https://github.com/zhaofengli/colmena/issues/85)).
 
 ## [Release 0.3.0](https://github.com/zhaofengli/colmena/releases/tag/v0.3.0) (2022/04/27)
 

src/command/apply_local.rs

@@ -1,10 +1,8 @@
-use std::env;
 use regex::Regex;
 use std::collections::HashMap;
 
 use clap::{Arg, Command as ClapCommand, ArgMatches};
 use tokio::fs;
-use tokio::process::Command;
 
 use crate::error::ColmenaError;
 use crate::nix::deployment::{
@@ -13,7 +11,7 @@ use crate::nix::deployment::{
     TargetNode,
     Options,
 };
-use crate::nix::{NodeName, host};
+use crate::nix::{NodeName, host::Local as LocalHost};
 use crate::progress::SimpleProgressOutput;
 use crate::util;
 
@@ -29,12 +27,6 @@ pub fn subcommand() -> ClapCommand<'static> {
         .arg(Arg::new("sudo")
             .long("sudo")
             .help("Attempt to escalate privileges if not run as root"))
-        .arg(Arg::new("sudo-command")
-            .long("sudo-command")
-            .value_name("COMMAND")
-            .help("Command to use to escalate privileges")
-            .default_value("sudo")
-            .takes_value(true))
         .arg(Arg::new("verbose")
             .short('v')
             .long("verbose")
@@ -53,13 +45,22 @@ By default, Colmena will deploy keys set in `deployment.keys` before activating
             .long("node")
             .help("Override the node name to use")
             .takes_value(true))
-        .arg(Arg::new("we-are-launched-by-sudo")
-            .long("we-are-launched-by-sudo")
+
+        // Removed
+        .arg(Arg::new("sudo-command")
+            .long("sudo-command")
+            .value_name("COMMAND")
+            .help("Removed: Configure deployment.privilegeEscalationCommand in node configuration")
             .hide(true)
-            .takes_value(false))
+            .takes_value(true))
 }
 
 pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(), ColmenaError> {
+    if local_args.occurrences_of("sudo-command") > 0 {
+        log::error!("--sudo-command has been removed. Please configure it in deployment.privilegeEscalationCommand in the node configuration.");
+        quit::with_code(1);
+    }
+
     // Sanity check: Are we running NixOS?
     if let Ok(os_release) = fs::read_to_string("/etc/os-release").await {
         let re = Regex::new(r#"ID="?nixos"?"#).unwrap();
@@ -72,23 +73,14 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
         quit::with_code(5);
     }
 
-    // Escalate privileges?
+    let escalate_privileges = local_args.is_present("sudo");
+    let verbose = local_args.is_present("verbose") || escalate_privileges; // cannot use spinners with interactive sudo
+
     {
         let euid: u32 = unsafe { libc::geteuid() };
-        if euid != 0 {
-            if local_args.is_present("we-are-launched-by-sudo") {
-                log::error!("Failed to escalate privileges. We are still not root despite a successful sudo invocation.");
-                quit::with_code(3);
-            }
-
-            if local_args.is_present("sudo") {
-                let sudo = local_args.value_of("sudo-command").unwrap();
-
-                escalate(sudo).await;
-            } else {
-                log::warn!("Colmena was not started by root. This is probably not going to work.");
-                log::warn!("Hint: Add the --sudo flag.");
-            }
+        if euid != 0 && !escalate_privileges {
+            log::warn!("Colmena was not started by root. This is probably not going to work.");
+            log::warn!("Hint: Add the --sudo flag.");
         }
     }
 
@@ -113,9 +105,15 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
                 log::error!("Hint: Set deployment.allowLocalDeployment to true.");
                 quit::with_code(2);
             }
+            let mut host = LocalHost::new(nix_options);
+            if escalate_privileges {
+                let command = info.privilege_escalation_command().to_owned();
+                host.set_privilege_escalation_command(Some(command));
+            }
+
             TargetNode::new(
                 hostname.clone(),
-                Some(host::local(nix_options)),
+                Some(host.upcast()),
                 info.clone(),
             )
         } else {
@@ -127,7 +125,7 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
     let mut targets = HashMap::new();
     targets.insert(hostname.clone(), target);
 
-    let mut output = SimpleProgressOutput::new(local_args.is_present("verbose"));
+    let mut output = SimpleProgressOutput::new(verbose);
     let progress = output.get_sender();
 
     let mut deployment = Deployment::new(hive, targets, goal, progress);
@@ -149,21 +147,3 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
 
     Ok(())
 }
-
-async fn escalate(sudo: &str) -> ! {
-    // Restart ourselves with sudo
-    let argv: Vec<String> = env::args().collect();
-
-    let exit = Command::new(sudo)
-        .arg("--")
-        .args(argv)
-        .arg("--we-are-launched-by-sudo")
-        .spawn()
-        .expect("Failed to run sudo to escalate privileges")
-        .wait()
-        .await
-        .expect("Failed to wait on child");
-
-    // Exit with the same exit code
-    quit::with_code(exit.code().unwrap());
-}

src/nix/deployment/mod.rs

@@ -35,6 +35,7 @@ use super::{
     CopyDirection,
     CopyOptions,
     RebootOptions,
+    host::Local as LocalHost,
     key::{Key, UploadAt as UploadKeyAt},
     evaluator::{
         DrvSetEvaluator,
@@ -42,7 +43,6 @@ use super::{
         EvalError,
     },
 };
-use super::host;
 
 /// A deployment.
 pub type DeploymentHandle = Arc<Deployment>;
@@ -450,7 +450,7 @@ impl Deployment {
         let arc_self = self.clone();
         let profile: Profile = build_job.run(|job| async move {
             // FIXME: Remote builder?
-            let mut builder = host::local(arc_self.nix_options.clone());
+            let mut builder = LocalHost::new(arc_self.nix_options.clone()).upcast();
             builder.set_job(Some(job.clone()));
 
             let profile = profile_drv.realize(&mut builder).await?;

src/nix/host/local.rs

@@ -19,6 +19,7 @@ use super::{CopyDirection, CopyOptions, Host, key_uploader};
 pub struct Local {
     job: Option<JobHandle>,
     nix_options: NixOptions,
+    privilege_escalation_command: Option<Vec<String>>,
 }
 
 impl Local {
@@ -26,6 +27,7 @@ impl Local {
         Self {
             job: None,
             nix_options,
+            privilege_escalation_command: None,
         }
     }
 }
@@ -71,17 +73,15 @@ impl Host for Local {
 
         if goal.should_switch_profile() {
             let path = profile.as_path().to_str().unwrap();
-            Command::new("nix-env")
-                .args(&["--profile", SYSTEM_PROFILE])
-                .args(&["--set", path])
+            self.make_privileged_command(&["nix-env", "--profile", SYSTEM_PROFILE, "--set", path])
                 .passthrough()
                 .await?;
         }
 
-        let activation_command = profile.activation_command(goal).unwrap();
-        let mut command = Command::new(&activation_command[0]);
-        command
-            .args(&activation_command[1..]);
+        let command = {
+            let activation_command = profile.activation_command(goal).unwrap();
+            self.make_privileged_command(&activation_command)
+        };
 
         let mut execution = CommandExecution::new(command);
 
@@ -126,6 +126,14 @@ impl Host for Local {
 }
 
 impl Local {
+    pub fn set_privilege_escalation_command(&mut self, command: Option<Vec<String>>) {
+        self.privilege_escalation_command = command;
+    }
+
+    pub fn upcast(self) -> Box<dyn Host> {
+        Box::new(self)
+    }
+
     /// "Uploads" a single key.
     async fn upload_key(&mut self, name: &str, key: &Key, require_ownership: bool) -> ColmenaResult<()> {
         if let Some(job) = &self.job {
@@ -145,4 +153,20 @@ impl Local {
         let uploader = command.spawn()?;
         key_uploader::feed_uploader(uploader, key, self.job.clone()).await
     }
+
+    /// Constructs a command with privilege escalation.
+    fn make_privileged_command<S: AsRef<str>>(&self, command: &[S]) -> Command {
+        let mut full_command = Vec::new();
+        if let Some(esc) = &self.privilege_escalation_command {
+            full_command.extend(esc.iter().map(|s| s.as_str()));
+        }
+        full_command.extend(command.iter().map(|s| s.as_ref()));
+
+        let mut result = Command::new(full_command[0]);
+        if full_command.len() > 1 {
+            result.args(&full_command[1..]);
+        }
+
+        result
+    }
 }

src/nix/host/mod.rs

@@ -4,7 +4,7 @@ use async_trait::async_trait;
 
 use crate::error::{ColmenaError, ColmenaResult};
 use crate::job::JobHandle;
-use super::{StorePath, Profile, Goal, Key, NixOptions};
+use super::{StorePath, Profile, Goal, Key};
 
 mod ssh;
 pub use ssh::Ssh;
@@ -14,10 +14,6 @@ pub use local::Local;
 
 mod key_uploader;
 
-pub(crate) fn local(nix_options: NixOptions) -> Box<dyn Host + 'static> {
-    Box::new(Local::new(nix_options))
-}
-
 #[derive(Copy, Clone, Debug)]
 pub enum CopyDirection {
     ToRemote,

src/nix/mod.rs

@@ -156,6 +156,8 @@ impl NodeConfig {
     #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
     pub fn allows_local_deployment(&self) -> bool { self.allow_local_deployment }
 
+    pub fn privilege_escalation_command(&self) -> &Vec<String> { &self.privilege_escalation_command }
+
     pub fn build_on_target(&self) -> bool { self.build_on_target }
     pub fn set_build_on_target(&mut self, enable: bool) {
         self.build_on_target = enable;