From 4c7f8eb838d0ab7c721ba65d7776d14d0b426758 Mon Sep 17 00:00:00 2001
From: Justinas Stankevicius <justinas@justinas.org>
Date: Thu, 11 Feb 2021 21:01:31 +0200
Subject: [PATCH] keyCommand: on error, do not upload key, report

---
 src/nix/deployment.rs |  4 ++--
 src/nix/host/ssh.rs   |  6 +++---
 src/nix/key.rs        | 38 ++++++++++++++++++++++++++++++++------
 src/nix/mod.rs        |  9 +++++++++
 src/progress.rs       |  4 ++++
 5 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/src/nix/deployment.rs b/src/nix/deployment.rs
index 60e438c..0bcdae1 100644
--- a/src/nix/deployment.rs
+++ b/src/nix/deployment.rs
@@ -281,7 +281,7 @@ impl Deployment {
                         task.log("Uploading keys...");
 
                         if let Err(e) = target.host.upload_keys(&target.config.keys).await {
-                            task.failure(&format!("Failed to upload keys: {}", e));
+                            task.failure_err(&e);
 
                             let mut results = arc_self.results.lock().await;
                             let stage = Stage::Apply(node.to_string());
@@ -481,7 +481,7 @@ impl Deployment {
             bar.log("Uploading keys...");
 
             if let Err(e) = target.host.upload_keys(&target.config.keys).await {
-                bar.failure(&format!("Failed to upload keys: {}", e));
+                bar.failure_err(&e);
 
                 let mut results = self.results.lock().await;
                 let stage = Stage::Apply(name.to_string());
diff --git a/src/nix/host/ssh.rs b/src/nix/host/ssh.rs
index 02ba6d2..7c7d97a 100644
--- a/src/nix/host/ssh.rs
+++ b/src/nix/host/ssh.rs
@@ -211,10 +211,10 @@ impl Ssh {
         command.stderr(Stdio::piped());
         command.stdout(Stdio::piped());
 
-        let mut child = command.spawn()?;
-
-        let mut stdin = child.stdin.take().unwrap();
         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);
diff --git a/src/nix/key.rs b/src/nix/key.rs
index 76f22f8..a076de8 100644
--- a/src/nix/key.rs
+++ b/src/nix/key.rs
@@ -2,11 +2,12 @@ use std::{
     convert::TryFrom,
     io::{self, Cursor},
     path::{Path, PathBuf},
-    process::Stdio,
+    process::{ExitStatus, Stdio},
 };
 
 use regex::Regex;
 use serde::{Deserialize, Serialize};
+use snafu::Snafu;
 use tokio::{
     fs::File,
     io::AsyncRead,
@@ -14,6 +15,21 @@ use tokio::{
 };
 use validator::{Validate, ValidationError};
 
+#[non_exhaustive]
+#[derive(Debug, Snafu)]
+pub enum KeyError {
+    #[snafu(display("I/O Error: {}", error))]
+    IoError { error: io::Error },
+    #[snafu(display("Key command failed: {}, stderr: {}", status, stderr))]
+    KeyCommandStatus { status: ExitStatus, stderr: String },
+}
+
+impl From<std::io::Error> for KeyError {
+    fn from(error: std::io::Error) -> Self {
+        Self::IoError { error }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(try_from = "KeySources")]
 enum KeySource {
@@ -78,7 +94,7 @@ pub struct Key {
 }
 
 impl Key {
-    pub async fn reader(&'_ self,) -> Result<Box<dyn AsyncRead + Send + Unpin + '_>, io::Error> {
+    pub async fn reader(&'_ self) -> Result<Box<dyn AsyncRead + Send + Unpin + '_>, KeyError> {
         match &self.source {
             KeySource::Text(content) => {
                 Ok(Box::new(Cursor::new(content)))
@@ -87,15 +103,25 @@ impl Key {
                 let pathname = &command[0];
                 let argv = &command[1..];
 
-                let stdout = Command::new(pathname)
+                let output = Command::new(pathname)
                     .args(argv)
                     .stdin(Stdio::null())
                     .stdout(Stdio::piped())
-                    .stderr(Stdio::null())
+                    .stderr(Stdio::piped())
                     .spawn()?
-                    .stdout.take().unwrap();
+                    .wait_with_output().await?;
 
-                Ok(Box::new(stdout))
+                if output.status.success() {
+                    Ok(Box::new(Cursor::new(output.stdout)))
+                } else {
+                    Err(KeyError::KeyCommandStatus {
+                        status: output.status,
+                        stderr: std::str::from_utf8(&output.stderr)
+                            .unwrap_or_default()
+                            .trim_end()
+                            .into(),
+                    })
+                }
             }
             KeySource::File(path) => {
                 Ok(Box::new(File::open(path).await?))
diff --git a/src/nix/mod.rs b/src/nix/mod.rs
index 1169ade..cadd36d 100644
--- a/src/nix/mod.rs
+++ b/src/nix/mod.rs
@@ -59,6 +59,9 @@ pub enum NixError {
     #[snafu(display("Validation error"))]
     ValidationError { errors: ValidationErrors },
 
+    #[snafu(display("Failed to upload keys: {}", error))]
+    KeyError { error: key::KeyError },
+
     #[snafu(display("Invalid NixOS system profile"))]
     InvalidProfile,
 
@@ -72,6 +75,12 @@ impl From<std::io::Error> for NixError {
     }
 }
 
+impl From<key::KeyError> for NixError {
+    fn from(error: key::KeyError) -> Self {
+        Self::KeyError { error }
+    }
+}
+
 impl From<ValidationErrors> for NixError {
     fn from(errors: ValidationErrors) -> Self {
         Self::ValidationError { errors }
diff --git a/src/progress.rs b/src/progress.rs
index bcac4dd..28e45c6 100644
--- a/src/progress.rs
+++ b/src/progress.rs
@@ -224,6 +224,10 @@ impl TaskProgress {
         }
     }
 
+    pub fn failure_err<E: std::error::Error>(self, error: &E) {
+        self.failure(&error.to_string())
+    }
+
     fn plain_print(&self, style: Style, line: &str) {
         eprintln!("{:>width$} | {}", style.apply_to(&self.label), line, width = self.label_width);
     }