populate: deliver sources other than flakes with the Rust program

The program delivers derivation, file, git, pass, passage, pipe, and
symlink sources, makes backups, and prunes, as the shell implementation
does, with the same messages.  Commands on the target are fixed shell
snippets in src/sh.rs, which get values as positional parameters, run
directly for local targets and over ssh, sharing one master connection,
for remote ones.  rsync's transfers and its check for .rsync-filter files
it can't read are in src/rsync.rs.  Given flake or flake-input sources, or
sources drawing inputs from others, the program fails before changing
anything, saying that these aren't implemented yet.

SIGHUP, SIGINT, and SIGTERM make populate stop once the command it waits
for ends, stop the ssh master, remove its temporary directory, and exit
with 128 plus the signal's number.

The spec's target gets local, which rust.nix sets during evaluation, like
the other functions of krops decide, and which the program decides when
it runs if it's null.  The spec comes on file descriptor 3, so the
commands populate runs keep standard input, and the tools populate runs
come first in PATH.

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(&quote(&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"));
+    }
+}