From ce9f639a53d497eaa761c755932bd2ca5b93c75e Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Wed, 10 Feb 2021 17:34:52 -0800
Subject: [PATCH] key: Make the key source better typed

---
 src/nix/host/local.rs |  6 ++---
 src/nix/host/ssh.rs   |  8 +++----
 src/nix/key.rs        | 54 +++++++++++++++++++++++++++++++------------
 3 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/src/nix/host/local.rs b/src/nix/host/local.rs
index 0456fc5..c906f2c 100644
--- a/src/nix/host/local.rs
+++ b/src/nix/host/local.rs
@@ -107,7 +107,7 @@ impl Local {
     async fn upload_key(&mut self, name: &str, key: &Key) -> NixResult<()> {
         self.progress_bar.log(&format!("Deploying key {}", name));
 
-        let dest_path = key.dest_dir.join(name);
+        let dest_path = key.dest_dir().join(name);
 
         let temp = NamedTempFile::new()?;
         let (_, temp_path) = temp.keep().map_err(|pe| pe.error)?;
@@ -120,7 +120,7 @@ impl Local {
         {
             let mut command = Command::new("chmod");
             command
-                .arg(&key.permissions)
+                .arg(&key.permissions())
                 .arg(&temp_path);
 
             let mut execution = CommandExecution::new(command);
@@ -135,7 +135,7 @@ impl Local {
         {
             let mut command = Command::new("chown");
             command
-                .arg(&format!("{}:{}", key.user, key.group))
+                .arg(&format!("{}:{}", key.user(), key.group()))
                 .arg(&temp_path);
 
             let mut execution = CommandExecution::new(command);
diff --git a/src/nix/host/ssh.rs b/src/nix/host/ssh.rs
index 74322ce..02ba6d2 100644
--- a/src/nix/host/ssh.rs
+++ b/src/nix/host/ssh.rs
@@ -197,13 +197,13 @@ impl Ssh {
     async fn upload_key(&mut self, name: &str, key: &Key) -> NixResult<()> {
         self.progress_bar.log(&format!("Deploying key {}", name));
 
-        let dest_path = key.dest_dir.join(name);
+        let dest_path = key.dest_dir().join(name);
 
         let remote_command = DEPLOY_KEY_TEMPLATE.to_string()
             .replace("%DESTINATION%", dest_path.to_str().unwrap())
-            .replace("%USER%", &key.user)
-            .replace("%GROUP%", &key.group)
-            .replace("%PERMISSIONS%", &key.permissions);
+            .replace("%USER%", &key.user())
+            .replace("%GROUP%", &key.group())
+            .replace("%PERMISSIONS%", &key.permissions());
 
         let mut command = self.ssh(&["sh", "-c", &remote_command]);
 
diff --git a/src/nix/key.rs b/src/nix/key.rs
index eb8021e..2cf637a 100644
--- a/src/nix/key.rs
+++ b/src/nix/key.rs
@@ -1,6 +1,6 @@
 use std::{
     io::{self, Cursor},
-    path::PathBuf,
+    path::{Path, PathBuf},
 };
 
 use regex::Regex;
@@ -8,31 +8,55 @@ use serde::{Deserialize, Serialize};
 use tokio::{fs::File, io::AsyncRead};
 use validator::{Validate, ValidationError};
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+enum KeySource {
+    #[serde(rename = "text")]
+    Text(String),
+
+    #[serde(rename = "keyCommand")]
+    Command(Vec<String>),
+
+    #[serde(rename = "keyFile")]
+    File(PathBuf),
+}
+
 #[derive(Debug, Clone, Validate, Serialize, Deserialize)]
 pub struct Key {
-    pub(crate) text: Option<String>,
-    #[serde(rename = "keyFile")]
-    pub(crate) key_file: Option<String>,
+    #[serde(flatten)]
+    source: KeySource,
+
     #[validate(custom = "validate_dest_dir")]
     #[serde(rename = "destDir")]
-    pub(super) dest_dir: PathBuf,
+    dest_dir: PathBuf,
+
     #[validate(custom = "validate_unix_name")]
-    pub(super) user: String,
+    user: String,
+
     #[validate(custom = "validate_unix_name")]
-    pub(super) group: String,
-    pub(super) permissions: String,
+    group: String,
+
+    permissions: String,
 }
 
 impl Key {
-    pub(crate) async fn reader(&'_ self,) -> Result<Box<dyn AsyncRead + Send + Unpin + '_>, io::Error> {
-        if let Some(ref t) = self.text {
-            Ok(Box::new(Cursor::new(t)))
-        } else if let Some(ref p) = self.key_file {
-            Ok(Box::new(File::open(p).await?))
-        } else {
-            unreachable!("Neither `text` nor `keyFile` set. This should have been validated by Nix assertions.");
+    pub async fn reader(&'_ self,) -> Result<Box<dyn AsyncRead + Send + Unpin + '_>, io::Error> {
+        match &self.source {
+            KeySource::Text(content) => {
+                Ok(Box::new(Cursor::new(content)))
+            }
+            KeySource::Command(_command) => {
+                todo!("Implement keyCommand support")
+            }
+            KeySource::File(path) => {
+                Ok(Box::new(File::open(path).await?))
+            }
         }
     }
+
+    pub fn dest_dir(&self) -> &Path { &self.dest_dir }
+    pub fn user(&self) -> &str { &self.user }
+    pub fn group(&self) -> &str { &self.user }
+    pub fn permissions(&self) -> &str { &self.permissions }
 }
 
 fn validate_unix_name(name: &str) -> Result<(), ValidationError> {