From f01f6ac197cf44255cd55a9be581bb783a2ad62e Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Thu, 1 Sep 2022 18:42:37 -0600
Subject: [PATCH] deployment: Move post-activation key upload after reboot

Ref: #113
---
 manual/src/release-notes.md |  1 +
 src/nix/deployment/mod.rs   | 66 ++++++++++++++++++-------------------
 2 files changed, 34 insertions(+), 33 deletions(-)

diff --git a/manual/src/release-notes.md b/manual/src/release-notes.md
index f512d38..57d39ed 100644
--- a/manual/src/release-notes.md
+++ b/manual/src/release-notes.md
@@ -13,6 +13,7 @@
 - The [`meta.nodeSpecialArgs`](./reference/meta.md#nodespecialargs) option has been added. It allows specifying node-specific `specialArgs` passed to NixOS modules ([#100](https://github.com/zhaofengli/colmena/pull/100)).
 - The [`repl`](./reference/cli.html#colmena-repl) subcommand has been added. It allows you to start an [interactive REPL](./features/eval.md#interactive-repl) with access to the complete node configurations.
 - The default goal for `colmena apply` is now `boot` if `--reboot` is specified, and `switch` otherwise ([#113](https://github.com/zhaofengli/colmena/issues/113)).
+- Post-activation keys are now uploaded after the reboot if `--reboot` is specified ([#113](https://github.com/zhaofengli/colmena/issues/113)).
 
 ## [Release 0.3.1](https://github.com/zhaofengli/colmena/releases/tag/v0.3.1) (2022/08/18)
 
diff --git a/src/nix/deployment/mod.rs b/src/nix/deployment/mod.rs
index 494829b..c0363ca 100644
--- a/src/nix/deployment/mod.rs
+++ b/src/nix/deployment/mod.rs
@@ -632,40 +632,8 @@ impl Deployment {
             Ok(target)
         }).await?;
 
-        // Upload post-activation keys
-        let mut target = if self.options.upload_keys {
-            let job = parent.create_job(JobType::UploadKeys, nodes.clone())?;
-            job.run_waiting(|job| async move {
-                let keys = target
-                    .config
-                    .keys
-                    .iter()
-                    .filter(|(_, v)| v.upload_at() == UploadKeyAt::PostActivation)
-                    .map(|(k, v)| (k.clone(), v.clone()))
-                    .collect::<HashMap<String, Key>>();
-
-                if keys.is_empty() {
-                    job.noop("No post-activation keys to upload".to_string())?;
-                    return Ok(target);
-                }
-
-                job.state(JobState::Running)?;
-                job.message("Uploading post-activation keys...".to_string())?;
-
-                let host = target.host.as_mut().unwrap();
-                host.set_job(Some(job.clone()));
-                host.upload_keys(&keys, true).await?;
-
-                job.success_with_message("Uploaded keys (post-activation)".to_string())?;
-                Ok(target)
-            })
-            .await?
-        } else {
-            target
-        };
-
         // Reboot
-        if self.options.reboot {
+        let mut target = if self.options.reboot {
             let job = parent.create_job(JobType::Reboot, nodes.clone())?;
             let arc_self = self.clone();
             job.run(|job| async move {
@@ -684,6 +652,38 @@ impl Deployment {
 
                 host.reboot(options).await?;
 
+                Ok(target)
+            })
+            .await?
+        } else {
+            target
+        };
+
+        // Upload post-activation keys
+        if self.options.upload_keys {
+            let job = parent.create_job(JobType::UploadKeys, nodes.clone())?;
+            job.run_waiting(|job| async move {
+                let keys = target
+                    .config
+                    .keys
+                    .iter()
+                    .filter(|(_, v)| v.upload_at() == UploadKeyAt::PostActivation)
+                    .map(|(k, v)| (k.clone(), v.clone()))
+                    .collect::<HashMap<String, Key>>();
+
+                if keys.is_empty() {
+                    job.noop("No post-activation keys to upload".to_string())?;
+                    return Ok(());
+                }
+
+                job.state(JobState::Running)?;
+                job.message("Uploading post-activation keys...".to_string())?;
+
+                let host = target.host.as_mut().unwrap();
+                host.set_job(Some(job.clone()));
+                host.upload_keys(&keys, true).await?;
+
+                job.success_with_message("Uploaded keys (post-activation)".to_string())?;
                 Ok(())
             })
             .await?;