From 4a8888ed5918ec29e88c27c3a37f5ce572867486 Mon Sep 17 00:00:00 2001
From: Hannah Neary <hannah@evervault.com>
Date: Tue, 10 Dec 2024 11:41:56 +0000
Subject: [PATCH] clean up

---
 control-plane/src/orchestration.rs | 51 +++++++++++++++++++++++-------
 1 file changed, 40 insertions(+), 11 deletions(-)

diff --git a/control-plane/src/orchestration.rs b/control-plane/src/orchestration.rs
index d7a575c8..1d7ecafc 100644
--- a/control-plane/src/orchestration.rs
+++ b/control-plane/src/orchestration.rs
@@ -21,12 +21,37 @@ pub enum OrchestrationError {
 }
 
 static EMPTY_VEC: Vec<Value> = Vec::new();
+const ENCLAVE_CID: &str = "2021";
+const EIF_PATH: &str = "enclave.eif";
+const NITRO_CLI: &str = "nitro-cli";
+
+enum NitroCommand {
+    TerminateEnclave,
+    DescribeEnclaves,
+    RunEnclave,
+}
+
+impl NitroCommand {
+    pub fn as_str(&self) -> &str {
+        match self {
+            NitroCommand::TerminateEnclave => "terminate-enclave",
+            NitroCommand::DescribeEnclaves => "describe-enclaves",
+            NitroCommand::RunEnclave => "run-enclave",
+        }
+    }
+}
 
 pub struct Orchestration;
 
 impl Orchestration {
     pub async fn shutdown_all_enclaves() -> Result<String, OrchestrationError> {
-        let command = vec!["sh", "-c", "nitro-cli", "terminate-enclave", "--all"];
+        let command = vec![
+            "sh",
+            "-c",
+            NITRO_CLI,
+            NitroCommand::TerminateEnclave.as_str(),
+            "--all",
+        ];
         Self::run_command_capture_stdout(&command).await
     }
 
@@ -36,30 +61,31 @@ impl Orchestration {
         info!("[HOST] Checking for running enclaves...");
 
         let running_enclaves =
-            Self::run_command_capture_stdout(&["nitro-cli", "describe-enclaves"]).await?;
+            Self::run_command_capture_stdout(&[NITRO_CLI, NitroCommand::DescribeEnclaves.as_str()])
+                .await?;
         let enclaves: Value = serde_json::from_str(&running_enclaves)?;
         let enclaves_array = enclaves.as_array().unwrap_or(&EMPTY_VEC);
         if !enclaves_array.is_empty() {
             info!("There's an enclave already running on this host. Terminating it...");
             Self::shutdown_all_enclaves().await?;
             info!("Enclave terminated. Waiting 10s...");
-            std::thread::sleep(std::time::Duration::from_secs(10));
+            tokio::time::sleep(std::time::Duration::from_secs(10)).await;
         } else {
             info!("No enclaves currently running on this host.");
         }
 
         info!("Starting new enclave...");
         let mut run_command = vec![
-            "nitro-cli",
-            "run-enclave",
+            NITRO_CLI,
+            NitroCommand::RunEnclave.as_str(),
             "--cpu-count",
             &run_config.num_cpus,
             "--memory",
             &run_config.ram_size_mib,
             "--enclave-cid",
-            "2021",
+            ENCLAVE_CID,
             "--eif-path",
-            "enclave.eif",
+            EIF_PATH,
         ];
 
         if run_config.debug_mode == "true" {
@@ -72,18 +98,21 @@ impl Orchestration {
         Self::run_command_capture_stdout(&run_command).await?;
 
         info!("Enclave started... Waiting 5 seconds for warmup.");
-        std::thread::sleep(std::time::Duration::from_secs(10));
+        tokio::time::sleep(std::time::Duration::from_secs(5)).await;
 
         if run_config.debug_mode == "true" {
             info!("Attaching headless console for running enclaves...");
-            let running_enclaves =
-                Self::run_command_capture_stdout(&["nitro-cli", "describe-enclaves"]).await?;
+            let running_enclaves = Self::run_command_capture_stdout(&[
+                NITRO_CLI,
+                NitroCommand::DescribeEnclaves.as_str(),
+            ])
+            .await?;
             let enclaves: Value = serde_json::from_str(&running_enclaves)?;
             let enclaves_array = enclaves.as_array().unwrap_or(&EMPTY_VEC).clone();
             for enclave in enclaves_array {
                 let id = enclave["EnclaveID"].as_str().unwrap().to_string();
 
-                let mut child = Command::new("nitro-cli")
+                let mut child = Command::new(NITRO_CLI)
                     .args(["console", "--enclave-id", &id])
                     .stdout(Stdio::piped())
                     .stderr(Stdio::piped())