From e9487ced9e4e578253351aead795a77fbaddeac3 Mon Sep 17 00:00:00 2001 From: Zhaofeng Li Date: Wed, 17 Mar 2021 22:39:05 -0700 Subject: [PATCH] host: Use the key uploader script for both SSH and local --- src/nix/host/key_uploader.rs | 59 ++++++++++++++++++ ...-key.template => key_uploader.template.sh} | 0 src/nix/host/local.rs | 62 +++---------------- src/nix/host/mod.rs | 2 + src/nix/host/ssh.rs | 48 ++------------ 5 files changed, 77 insertions(+), 94 deletions(-) create mode 100644 src/nix/host/key_uploader.rs rename src/nix/host/{deploy-key.template => key_uploader.template.sh} (100%) diff --git a/src/nix/host/key_uploader.rs b/src/nix/host/key_uploader.rs new file mode 100644 index 0000000..58e9a90 --- /dev/null +++ b/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() }) + } +} diff --git a/src/nix/host/deploy-key.template b/src/nix/host/key_uploader.template.sh similarity index 100% rename from src/nix/host/deploy-key.template rename to src/nix/host/key_uploader.template.sh diff --git a/src/nix/host/local.rs b/src/nix/host/local.rs index c906f2c..2a092b6 100644 --- a/src/nix/host/local.rs +++ b/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 } } diff --git a/src/nix/host/mod.rs b/src/nix/host/mod.rs index 06321d4..c3e54ec 100644 --- a/src/nix/host/mod.rs +++ b/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 { Box::new(Local::new()) } diff --git a/src/nix/host/ssh.rs b/src/nix/host/ssh.rs index bd3dc2d..20316a5 100644 --- a/src/nix/host/ssh.rs +++ b/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 } }