diff options
Diffstat (limited to 'pkgs/populate/src/run.rs')
| -rw-r--r-- | pkgs/populate/src/run.rs | 348 |
1 files changed, 348 insertions, 0 deletions
diff --git a/pkgs/populate/src/run.rs b/pkgs/populate/src/run.rs new file mode 100644 index 0000000..9ea05b3 --- /dev/null +++ b/pkgs/populate/src/run.rs @@ -0,0 +1,348 @@ +//! Running commands: on the deploying machine, and on the target, directly +//! or over ssh, sharing one master connection. + +use crate::spec::Target; +use std::fmt; +use std::fs; +use std::io::Write; +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; +use std::process::{Command, ExitStatus, Stdio}; +use std::sync::atomic::{AtomicI32, Ordering}; + +/// Why populate stops before delivering everything. +#[derive(Debug)] +pub enum Failure { + /// A message to print. + Message(String), + /// A command failed, after printing why. + Command, + /// A signal arrived: SIGHUP, SIGINT, or SIGTERM. + Signal(i32), +} + +impl fmt::Display for Failure { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Failure::Message(message) => f.write_str(message), + Failure::Command => f.write_str("a command failed"), + Failure::Signal(signal) => write!(f, "got signal {signal}"), + } + } +} + +pub type Result<T> = std::result::Result<T, Failure>; + +/// The signal that arrived last, or 0. +static SIGNAL: AtomicI32 = AtomicI32::new(0); + +extern "C" fn record_signal(signal: libc::c_int) { + SIGNAL.store(signal, Ordering::SeqCst); +} + +/// Makes SIGHUP, SIGINT, and SIGTERM get recorded instead of ending +/// populate, so it can clean up. A signal to the process group, like the +/// one of an interrupt key, ends the command populate waits for, too, after +/// which populate stops. A signal to populate alone takes effect once that +/// command ends. +pub fn trap_signals() { + for signal in [libc::SIGHUP, libc::SIGINT, libc::SIGTERM] { + // SAFETY: the handler only stores into an atomic. + unsafe { + let mut action: libc::sigaction = std::mem::zeroed(); + action.sa_sigaction = record_signal as extern "C" fn(libc::c_int) as usize; + action.sa_flags = libc::SA_RESTART; + libc::sigemptyset(&mut action.sa_mask); + libc::sigaction(signal, &action, std::ptr::null_mut()); + } + } +} + +/// Ignores further signals, so they don't end cleaning up halfway. +pub fn ignore_signals() { + for signal in [libc::SIGHUP, libc::SIGINT, libc::SIGTERM] { + // SAFETY: setting a disposition has no other effect. + unsafe { + libc::signal(signal, libc::SIG_IGN); + } + } +} + +/// Fails if a signal arrived. +pub fn check_signal() -> Result<()> { + match SIGNAL.load(Ordering::SeqCst) { + 0 => Ok(()), + signal => Err(Failure::Signal(signal)), + } +} + +/// Quotes a word for POSIX shells. +pub fn quote(word: &str) -> String { + let safe = |c: char| c.is_ascii_alphanumeric() || "-+./:=_@%,".contains(c); + if !word.is_empty() && word.chars().all(safe) { + word.to_string() + } else { + format!("'{}'", word.replace('\'', r"'\''")) + } +} + +/// Where populate delivers to, and how it gets there. +pub struct Context { + pub target: Target, + /// Whether the target is the machine populate runs on, as the user + /// running it, so commands run directly instead of over ssh. + pub local: bool, + /// Populate's private temporary directory, which is $TMPDIR for the + /// commands it runs. + pub tmpdir: PathBuf, + /// The ssh command, with the options for the target and for sharing a + /// master connection, as a script in tmpdir, for rsync's -e as well. + pub ssh: PathBuf, +} + +impl Context { + /// Creates the temporary directory, and the ssh script in it. + pub fn new(target: Target, local: Option<bool>) -> Result<Context> { + let local = local.unwrap_or_else(|| is_local(&target)); + let output = Command::new("mktemp") + .args(["-dt", "populate.XXXXXXXX"]) + .stderr(Stdio::inherit()) + .output() + .map_err(|e| Failure::Message(format!("error: can't run mktemp: {e}")))?; + if !output.status.success() { + return Err(Failure::Command); + } + let tmpdir = PathBuf::from(String::from_utf8_lossy(&output.stdout).trim_end()); + let ssh = tmpdir.join("ssh"); + let context = Context { + target, + local, + tmpdir, + ssh, + }; + let mut script = String::from("#!/bin/sh\nexec ssh"); + for word in context.ssh_options() { + script.push(' '); + script.push_str("e(&word)); + } + script.push_str(" \"$@\"\n"); + let write = || -> std::io::Result<()> { + let mut file = fs::File::create(&context.ssh)?; + file.write_all(script.as_bytes())?; + fs::set_permissions(&context.ssh, fs::Permissions::from_mode(0o700)) + }; + write().map_err(|e| { + Failure::Message(format!("error: can't write {}: {e}", context.ssh.display())) + })?; + Ok(context) + } + + fn ssh_options(&self) -> Vec<String> { + let mut options = Vec::new(); + if let Some(user) = &self.target.user { + options.extend(["-l".to_string(), user.clone()]); + } + if let Some(port) = &self.target.port { + options.extend(["-p".to_string(), port.clone()]); + } + let control_path = format!("ControlPath={}/ssh-%r@%h:%p.sock", self.tmpdir.display()); + for option in ["ControlMaster=auto", &control_path, "ControlPersist=1h"] { + options.extend(["-o".to_string(), option.to_string()]); + } + options.push("-T".to_string()); + options.extend(self.target.extra_options.iter().cloned()); + options + } + + /// The path on the target, as populate names it in what it prints: + /// prefixed by the host if the target isn't local, and by the user as + /// well if that isn't root. + pub fn name(&self, path: &str) -> String { + if self.local { + return path.to_string(); + } + let user = match self.target.user.as_deref() { + None | Some("") | Some("root") => String::new(), + Some(user) => format!("{user}@"), + }; + format!("{user}{}:{path}", self.target.host) + } + + /// The path on the target, as rsync names it. + pub fn rsync_destination(&self, path: &str) -> String { + if self.local { + return path.to_string(); + } + let user = match self.target.user.as_deref() { + None | Some("") => String::new(), + Some(user) => format!("{user}@"), + }; + format!("{user}{}:{path}", self.target.host) + } + + /// A command running the shell code script on the target, with args as + /// its positional parameters. Over ssh, the remote shell gets it as + /// one quoted command, run via sudo if the target says so. + pub fn shell(&self, script: &str, args: &[&str]) -> Command { + let script = format!("set -efu\n{script}"); + if self.local { + let mut command = Command::new("sh"); + command.arg("-c").arg(script).arg("sh").args(args); + command + } else { + let mut words = vec![ + if self.target.sudo { + "sudo sh -c" + } else { + "sh -c" + } + .to_string(), + quote(&script), + "sh".to_string(), + ]; + words.extend(args.iter().map(|arg| quote(arg))); + let mut command = Command::new(&self.ssh); + command.arg(&self.target.host).arg(words.join(" ")); + command + } + } + + /// Runs shell code on the target, with its output going where + /// populate's errors go. + pub fn run(&self, script: &str, args: &[&str]) -> Result<()> { + let mut command = self.shell(script, args); + command.stdout(stderr()); + status(&mut command) + } + + /// Runs shell code on the target, returning its output. + pub fn output(&self, script: &str, args: &[&str]) -> Result<String> { + output(&mut self.shell(script, args)) + } + + /// Stops the ssh master, if any command started one, and removes the + /// temporary directory. + pub fn clean_up(&self) { + if !self.local { + let check = Command::new(&self.ssh) + .args(["-q", "-O", "check", &self.target.host]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + if matches!(check, Ok(status) if status.success()) { + let _ = Command::new(&self.ssh) + .args(["-q", "-O", "stop", &self.target.host]) + .stdin(Stdio::null()) + .status(); + } + } + let _ = Command::new("rm").arg("-R").arg(&self.tmpdir).status(); + } +} + +/// A command running shell code on the deploying machine, in dir, with +/// args as its positional parameters. +pub fn local_shell(script: &str, args: &[&str], dir: &Path) -> Command { + let mut command = Command::new("sh"); + command + .arg("-c") + .arg(format!("set -efu\n{script}")) + .arg("sh") + .args(args) + .current_dir(dir); + command +} + +/// Where populate's errors go, for commands whose output goes there too. +pub fn stderr() -> Stdio { + use std::os::fd::AsFd; + match std::io::stderr().as_fd().try_clone_to_owned() { + Ok(fd) => Stdio::from(fd), + Err(_) => Stdio::inherit(), + } +} + +fn spawn_error(command: &Command, e: std::io::Error) -> Failure { + Failure::Message(format!( + "error: can't run {}: {e}", + command.get_program().to_string_lossy() + )) +} + +/// Runs a command, failing if it fails, or if a signal arrived meanwhile. +pub fn status(command: &mut Command) -> Result<()> { + let result: std::io::Result<ExitStatus> = command.status(); + check_signal()?; + match result { + Ok(status) if status.success() => Ok(()), + Ok(_) => Err(Failure::Command), + Err(e) => Err(spawn_error(command, e)), + } +} + +/// Runs a command, returning its output, failing like status. +pub fn output(command: &mut Command) -> Result<String> { + command.stderr(Stdio::inherit()); + let result = command.output(); + check_signal()?; + match result { + Ok(output) if output.status.success() => { + Ok(String::from_utf8_lossy(&output.stdout).into_owned()) + } + Ok(_) => Err(Failure::Command), + Err(e) => Err(spawn_error(command, e)), + } +} + +/// Whether the target is the machine populate runs on, as the user running +/// it: its user is that user, and its host is localhost or the machine's +/// name, like krops' lib's isLocalTarget decides during evaluation. +fn is_local(target: &Target) -> bool { + let user = std::env::var("LOGNAME").ok(); + let host = std::env::var("HOSTNAME").ok().or_else(host_name); + target.user == user && (target.host == "localhost" || Some(&target.host) == host.as_ref()) +} + +/// The machine's name, from /etc/hostname, as krops' lib reads it. +fn host_name() -> Option<String> { + let text = fs::read_to_string("/etc/hostname").ok()?; + let label = |line: &&str| { + !line.is_empty() + && line.len() <= 63 + && line.chars().all(|c| c.is_ascii_alphanumeric() || c == '-') + && !line.starts_with('-') + && !line.ends_with('-') + }; + let names: Vec<&str> = text.lines().filter(label).collect(); + match names.as_slice() { + [name] => Some(name.to_string()), + _ => None, + } +} + +/// Tells whether path is a path in the store itself, not one below it. +pub fn is_store_path(path: &str) -> bool { + Path::new(path).parent() == Some(Path::new("/nix/store")) && path.len() > "/nix/store/".len() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn quoting() { + assert_eq!(quote("abc/d.e"), "abc/d.e"); + assert_eq!(quote(""), "''"); + assert_eq!(quote("a b"), "'a b'"); + assert_eq!(quote("it's"), r"'it'\''s'"); + } + + #[test] + fn store_paths() { + assert!(is_store_path("/nix/store/abc-x")); + assert!(!is_store_path("/nix/store/abc-x/y")); + assert!(!is_store_path("/nix/store")); + assert!(!is_store_path("/etc")); + } +} |
