From 24339bcca74c4bd9e70dd049822e013440497418 Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Tue, 24 Aug 2021 23:25:46 -0700
Subject: [PATCH] Add `deployment.keys.<name>.uploadAt`

This mirrors the functionality recently added in morph and allows
for the uploading of keys after system profile activation.

Fixes #10.
---
 README.md             | 10 ++++++----
 src/nix/deployment.rs | 30 ++++++++++++++++++++++++++++--
 src/nix/eval.nix      | 12 ++++++++++++
 src/nix/key.rs        | 16 ++++++++++++++++
 4 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 3c7cafd..bb56f63 100644
--- a/README.md
+++ b/README.md
@@ -254,10 +254,12 @@ For example, to deploy DNS-01 credentials for use with `security.acme`:
       # may be specified.
       keyCommand = [ "vault" "read" "-field=env" "secret/dns01" ];
 
-      destDir = "/run/keys"; # Default: /run/keys
-      user = "acme";         # Default: root
-      group = "nginx";       # Default: root
-      permissions = "0640";  # Default: 0600
+      destDir = "/run/keys";       # Default: /run/keys
+      user = "acme";               # Default: root
+      group = "nginx";             # Default: root
+      permissions = "0640";        # Default: 0600
+
+      uploadAt = "pre-activation"; # Default: pre-activation, Alternative: post-activation
     };
     # Rest of configuration...
   };
diff --git a/src/nix/deployment.rs b/src/nix/deployment.rs
index 275a932..02f4f7d 100644
--- a/src/nix/deployment.rs
+++ b/src/nix/deployment.rs
@@ -6,6 +6,7 @@ use futures::future::join_all;
 use tokio::sync::{Mutex, Semaphore};
 
 use super::{Hive, Host, CopyOptions, NodeConfig, Profile, StoreDerivation, ProfileMap, host};
+use super::key::{Key, UploadAt};
 use crate::progress::{Progress, TaskProgress, OutputStyle};
 
 /// Amount of RAM reserved for the system, in MB.
@@ -523,10 +524,20 @@ impl Deployment {
             }
         }
 
-        if self.options.upload_keys && !target.config.keys.is_empty() {
+        let pre_activation_keys = target.config.keys.iter()
+            .filter(|(_, v)| v.upload_at() == UploadAt::PreActivation)
+            .map(|(k, v)| (k.clone(), v.clone()))
+            .collect::<HashMap<String, Key>>();
+
+        let post_activation_keys = target.config.keys.iter()
+            .filter(|(_, v)| v.upload_at() == UploadAt::PostActivation)
+            .map(|(k, v)| (k.clone(), v.clone()))
+            .collect::<HashMap<String, Key>>();
+
+        if self.options.upload_keys && !pre_activation_keys.is_empty() {
             bar.log("Uploading keys...");
 
-            if let Err(e) = target.host.upload_keys(&target.config.keys).await {
+            if let Err(e) = target.host.upload_keys(&pre_activation_keys).await {
                 bar.failure_err(&e);
 
                 let mut results = self.results.lock().await;
@@ -546,6 +557,21 @@ impl Deployment {
 
         match target.host.deploy(&profile, self.goal, copy_options).await {
             Ok(_) => {
+                // FIXME: This is ugly
+                if self.options.upload_keys && !post_activation_keys.is_empty() {
+                    bar.log("Uploading keys (post-activation)...");
+
+                    if let Err(e) = target.host.upload_keys(&post_activation_keys).await {
+                        bar.failure_err(&e);
+
+                        let mut results = self.results.lock().await;
+                        let stage = Stage::Apply(name.to_string());
+                        let logs = target.host.dump_logs().await.map(|s| s.to_string());
+                        results.push(DeploymentResult::failure(stage, logs));
+                        return;
+                    }
+                }
+
                 bar.success(self.goal.success_str().unwrap());
 
                 let mut results = self.results.lock().await;
diff --git a/src/nix/eval.nix b/src/nix/eval.nix
index 9ba3521..f32268c 100644
--- a/src/nix/eval.nix
+++ b/src/nix/eval.nix
@@ -241,6 +241,18 @@ let
         default = "0600";
         type = types.str;
       };
+      uploadAt = lib.mkOption {
+        description = ''
+          When to upload the keys.
+
+          - pre-activation (default): Upload the keys before activating the new system profile.
+          - post-activation: Upload the keys after successfully activating the new system profile.
+
+          For `colmena upload-keys`, all keys are uploaded at the same time regardless of the configuration here.
+        '';
+        default = "pre-activation";
+        type = types.enum [ "pre-activation" "post-activation" ];
+      };
     };
   };
 
diff --git a/src/nix/key.rs b/src/nix/key.rs
index a3e78a4..b1189a0 100644
--- a/src/nix/key.rs
+++ b/src/nix/key.rs
@@ -75,6 +75,18 @@ struct KeySources {
     file: Option<PathBuf>,
 }
 
+/// When to upload a given key.
+#[derive(Debug, Copy, Clone, PartialEq, Serialize, Deserialize)]
+pub enum UploadAt {
+    /// Before activating the system profile.
+    #[serde(rename = "pre-activation")]
+    PreActivation,
+
+    /// After successfully activating the system profile.
+    #[serde(rename = "post-activation")]
+    PostActivation,
+}
+
 #[derive(Debug, Clone, Validate, Serialize, Deserialize)]
 pub struct Key {
     #[serde(flatten)]
@@ -91,6 +103,9 @@ pub struct Key {
     group: String,
 
     permissions: String,
+
+    #[serde(rename = "uploadAt")]
+    upload_at: UploadAt,
 }
 
 impl Key {
@@ -133,6 +148,7 @@ impl Key {
     pub fn user(&self) -> &str { &self.user }
     pub fn group(&self) -> &str { &self.group }
     pub fn permissions(&self) -> &str { &self.permissions }
+    pub fn upload_at(&self) -> UploadAt { self.upload_at }
 }
 
 fn validate_unix_name(name: &str) -> Result<(), ValidationError> {