host: Use the key uploader script for both SSH and local

src/nix/host/key_uploader.rs

@@ -0,0 +1,59 @@
+//! Utilities for using the key uploader script.
+//!
+//! The key uploader is a simple shell script that reads the contents
+//! of the secret file from stdin into a temporary file then atomically
+//! replaces the destination file with the temporary file.
+
+use std::borrow::Cow;
+use std::path::Path;
+
+use futures::future::join3;
+use shell_escape::unix::escape;
+use tokio::io::{AsyncWriteExt, BufReader};
+use tokio::process::Child;
+
+use crate::nix::{Key, NixError, NixResult};
+use crate::progress::TaskProgress;
+use crate::util::capture_stream;
+
+const SCRIPT_TEMPLATE: &'static str = include_str!("./key_uploader.template.sh");
+
+pub fn generate_script<'a>(key: &'a Key, destination: &'a Path) -> Cow<'a, str> {
+    let key_script = SCRIPT_TEMPLATE.to_string()
+        .replace("%DESTINATION%", destination.to_str().unwrap())
+        .replace("%USER%", &escape(key.user().into()))
+        .replace("%GROUP%", &escape(key.group().into()))
+        .replace("%PERMISSIONS%", &escape(key.permissions().into()))
+        .trim_end_matches('\n').to_string();
+
+    escape(key_script.into())
+}
+
+pub async fn feed_uploader(mut uploader: Child, key: &Key, progress: TaskProgress, logs: &mut String) -> NixResult<()> {
+    let mut reader = key.reader().await?;
+    let mut stdin = uploader.stdin.take().unwrap();
+
+    tokio::io::copy(reader.as_mut(), &mut stdin).await?;
+    stdin.flush().await?;
+    drop(stdin);
+
+    let stdout = BufReader::new(uploader.stdout.take().unwrap());
+    let stderr = BufReader::new(uploader.stderr.take().unwrap());
+
+    let futures = join3(
+        capture_stream(stdout, progress.clone()),
+        capture_stream(stderr, progress.clone()),
+        uploader.wait(),
+    );
+    let (stdout_str, stderr_str, exit) = futures.await;
+    logs.push_str(&stdout_str);
+    logs.push_str(&stderr_str);
+
+    let exit = exit?;
+
+    if exit.success() {
+        Ok(())
+    } else {
+        Err(NixError::NixFailure { exit_code: exit.code().unwrap() })
+    }
+}

src/nix/host/local.rs

@@ -1,13 +1,11 @@
 use std::convert::TryInto;
 use std::collections::HashMap;
-use std::fs;
+use std::process::Stdio;
 
 use async_trait::async_trait;
-use tokio::fs::OpenOptions;
 use tokio::process::Command;
-use tempfile::NamedTempFile;
 
-use super::{CopyDirection, CopyOptions, Host};
+use super::{CopyDirection, CopyOptions, Host, key_uploader};
 use crate::nix::{StorePath, Profile, Goal, NixResult, NixCommand, Key, SYSTEM_PROFILE};
 use crate::util::CommandExecution;
 use crate::progress::TaskProgress;
@@ -108,56 +106,16 @@ impl Local {
         self.progress_bar.log(&format!("Deploying key {}", name));
 
         let dest_path = key.dest_dir().join(name);
+        let key_script = format!("'{}'", key_uploader::generate_script(key, &dest_path));
 
-        let temp = NamedTempFile::new()?;
-        let (_, temp_path) = temp.keep().map_err(|pe| pe.error)?;
-        let mut reader = key.reader().await?;
-        let mut writer = OpenOptions::new().write(true).open(&temp_path).await?;
-        tokio::io::copy(reader.as_mut(), &mut writer).await?;
+        let mut command = Command::new("sh");
 
-        // Well, we need the userspace chmod program to parse the
-        // permission, for NixOps compatibility
-        {
-            let mut command = Command::new("chmod");
-            command
-                .arg(&key.permissions())
-                .arg(&temp_path);
+        command.args(&["-c", &key_script]);
+        command.stdin(Stdio::piped());
+        command.stderr(Stdio::piped());
+        command.stdout(Stdio::piped());
 
-            let mut execution = CommandExecution::new(command);
-            let exit = execution.run().await;
-
-            let (stdout, stderr) = execution.get_logs();
-            self.logs += stdout.unwrap();
-            self.logs += stderr.unwrap();
-
-            exit?;
-        }
-        {
-            let mut command = Command::new("chown");
-            command
-                .arg(&format!("{}:{}", key.user(), key.group()))
-                .arg(&temp_path);
-
-            let mut execution = CommandExecution::new(command);
-            let exit = execution.run().await;
-
-            let (stdout, stderr) = execution.get_logs();
-            self.logs += stdout.unwrap();
-            self.logs += stderr.unwrap();
-
-            exit?;
-        }
-
-        let parent_dir = dest_path.parent().unwrap();
-        fs::create_dir_all(parent_dir)?;
-
-        if fs::rename(&temp_path, &dest_path).is_err() {
-            // Linux can not rename across different filesystems, try copy-then-remove
-            let copy_result = fs::copy(&temp_path, &dest_path);
-            fs::remove_file(&temp_path)?;
-            copy_result?;
-        }
-
-        Ok(())
+        let uploader = command.spawn()?;
+        key_uploader::feed_uploader(uploader, key, self.progress_bar.clone(), &mut self.logs).await
     }
 }

src/nix/host/mod.rs

@@ -11,6 +11,8 @@ pub use ssh::Ssh;
 mod local;
 pub use local::Local;
 
+mod key_uploader;
+
 pub(crate) fn local() -> Box<dyn Host + 'static> {
     Box::new(Local::new())
 }

src/nix/host/ssh.rs

@@ -4,18 +4,13 @@ use std::path::PathBuf;
 use std::process::Stdio;
 
 use async_trait::async_trait;
-use futures::future::join3;
-use shell_escape::unix::escape;
 use tokio::process::Command;
-use tokio::io::{AsyncWriteExt, BufReader};
 
-use super::{CopyDirection, CopyOptions, Host};
-use crate::nix::{StorePath, Profile, Goal, NixResult, NixCommand, NixError, Key, SYSTEM_PROFILE};
-use crate::util::{CommandExecution, capture_stream};
+use super::{CopyDirection, CopyOptions, Host, key_uploader};
+use crate::nix::{StorePath, Profile, Goal, NixResult, NixCommand, Key, SYSTEM_PROFILE};
+use crate::util::CommandExecution;
 use crate::progress::TaskProgress;
 
-const DEPLOY_KEY_TEMPLATE: &'static str = include_str!("./deploy-key.template");
-
 /// A remote machine connected over SSH.
 #[derive(Debug)]
 pub struct Ssh {
@@ -203,13 +198,7 @@ impl Ssh {
         self.progress_bar.log(&format!("Deploying key {}", name));
 
         let dest_path = key.dest_dir().join(name);
-
-        let key_script = DEPLOY_KEY_TEMPLATE.to_string()
-            .replace("%DESTINATION%", dest_path.to_str().unwrap())
-            .replace("%USER%", &escape(key.user().into()))
-            .replace("%GROUP%", &escape(key.group().into()))
-            .replace("%PERMISSIONS%", &escape(key.permissions().into()));
-        let key_script = escape(key_script.into());
+        let key_script = key_uploader::generate_script(key, &dest_path);
 
         let mut command = self.ssh(&["sh", "-c", &key_script]);
 
@@ -217,32 +206,7 @@ impl Ssh {
         command.stderr(Stdio::piped());
         command.stdout(Stdio::piped());
 
-        let mut reader = key.reader().await?;
-
-        let mut child = command.spawn()?;
-        let mut stdin = child.stdin.take().unwrap();
-        tokio::io::copy(reader.as_mut(), &mut stdin).await?;
-        stdin.flush().await?;
-        drop(stdin);
-
-        let stdout = BufReader::new(child.stdout.take().unwrap());
-        let stderr = BufReader::new(child.stderr.take().unwrap());
-
-        let futures = join3(
-            capture_stream(stdout, self.progress_bar.clone()),
-            capture_stream(stderr, self.progress_bar.clone()),
-            child.wait(),
-        );
-        let (stdout_str, stderr_str, exit) = futures.await;
-        self.logs += &stdout_str;
-        self.logs += &stderr_str;
-
-        let exit = exit?;
-
-        if exit.success() {
-            Ok(())
-        } else {
-            Err(NixError::NixFailure { exit_code: exit.code().unwrap() })
-        }
+        let uploader = command.spawn()?;
+        key_uploader::feed_uploader(uploader, key, self.progress_bar.clone(), &mut self.logs).await
     }
 }