test cp orchestration

control-plane/src/configuration.rs

@@ -35,12 +35,36 @@ pub fn get_aws_profile() -> String {
     std::env::var("AWS_PROFILE").unwrap_or_else(|_| "ev-local-customers".to_string())
 }
 
+
 pub fn get_aws_region() -> aws_types::region::Region {
     let region = std::env::var("AWS_REGION")
         .ok()
         .unwrap_or_else(|| "us-east-1".to_string());
     aws_types::region::Region::new(region)
 }
+
+#[derive(Clone)]
+pub struct EnclaveRunConfig {
+    pub num_cpus: String,
+    pub ram_size_mib: String,
+    pub debug_mode: String,
+}
+
+impl EnclaveRunConfig {
+    pub fn new(
+        num_cpus: String,
+        ram_size_mib: String,
+        debug_mode: String,
+    ) -> EnclaveRunConfig {
+        EnclaveRunConfig {
+            num_cpus,
+            ram_size_mib,
+            debug_mode,
+        }
+    }
+}
+
+
 #[derive(Clone)]
 pub struct EnclaveContext {
     pub uuid: String,
@@ -114,6 +138,13 @@ pub fn get_team_uuid() -> String {
     std::env::var("EV_TEAM_UUID").expect("EV_TEAM_UUID is not set in env")
 }
 
+pub fn get_enclave_run_config() -> EnclaveRunConfig {
+    let num_cpus = std::env::var("ENCLAVE_NUM_CPUS").unwrap_or_else(|_| "2".to_string());
+    let ram_size_mib = std::env::var("ENCLAVE_RAM_SIZE_MIB").unwrap_or_else(|_| "512".to_string());
+    let debug_mode = std::env::var("ENCLAVE_DEBUG_MODE").unwrap_or_else(|_| "false".to_string());
+    EnclaveRunConfig::new(num_cpus, ram_size_mib, debug_mode)
+}
+
 pub fn get_cert_provisoner_host() -> String {
     match get_rust_env() {
         Environment::Staging | Environment::Production => "provisioner.cages.internal".to_string(),

control-plane/src/lib.rs

@@ -14,6 +14,7 @@ pub mod health;
 pub mod stats_client;
 pub mod stats_proxy;
 pub mod tls_proxy;
+pub mod orchestration;
 
 #[cfg(test)]
 pub mod mocks;

control-plane/src/main.rs

@@ -1,9 +1,11 @@
 use control_plane::clients::{cert_provisioner, mtls_config};
 use control_plane::dns::{ExternalAsyncDnsResolver, InternalAsyncDnsResolver};
+use control_plane::orchestration::Orchestration;
 use control_plane::stats_client::StatsClient;
 use control_plane::stats_proxy::StatsProxy;
 use control_plane::{config_server, tls_proxy};
 use shared::{print_version, utils::pipe_streams, ENCLAVE_CONNECT_PORT};
+use tls_parser::nom::Or;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::process::Command;
 use storage_client_interface::s3;
@@ -251,11 +253,7 @@ fn listen_for_shutdown_signal() {
                 // Wait for 55 seconds before terminating enclave - ECS waits 55 seconds to kill the container
                 sleep(Duration::from_millis(55000)).await;
 
-                let output = Command::new("sh")
-                    .arg("-c")
-                    .arg("nitro-cli terminate-enclave --all")
-                    .output()
-                    .expect("failed to terminate enclave");
+                let output = Orchestration::shutdown_all_enclaves().expect("failed to terminate enclave");
 
                 log::info!(
                     "Terminated enclave: {}",

control-plane/src/orchestration.rs

@@ -0,0 +1,103 @@
+use std::{process::{Command, Output, Stdio}, vec};
+use log::info;
+use serde_json::Value;
+use thiserror::Error;
+
+use crate::configuration::get_enclave_run_config;
+
+#[derive(Debug, Error)]
+pub enum OrchestrationError {
+    #[error("An error occurred - {0:?}")]
+    CommandFailed(String),
+    #[error("Json Error — {0:?}")]
+    SerdeError(#[from] serde_json::Error),
+    #[error(transparent)]
+    Io(#[from] std::io::Error),
+}
+
+// trait Orchestration {
+//     fn shutdown_all_enclaves() -> Result<Output, OrchestrationError>;
+//     fn start_enclave() -> Result<(), OrchestrationError>;
+//     fn run_command_capture_stdout(args: &[&str]) -> Result<String, OrchestrationError>;
+// }
+
+pub struct Orchestration;
+
+impl Orchestration {
+    pub fn shutdown_all_enclaves() -> Result<Output, OrchestrationError> {
+        Command::new("sh")
+                .arg("-c")
+                    .arg("nitro-cli terminate-enclave --all")
+                    .output()
+                    .map_err(|e| OrchestrationError::Io(e))
+    }
+
+
+    pub fn start_enclave() -> Result<(), OrchestrationError> {
+        let run_config = get_enclave_run_config();
+
+        info!("[HOST] Checking for running enclaves...");
+
+        let running_enclaves = Self::run_command_capture_stdout(&["nitro-cli", "describe-enclaves"])?;
+        let enclaves: Value = serde_json::from_str(&running_enclaves)?;
+        let v = vec![];
+        let enclaves_array = enclaves.as_array().unwrap_or(&v);
+        if enclaves_array.len() > 0 {
+            info!("There's an enclave already running on this host. Terminating it...");
+            Self::shutdown_all_enclaves()?;
+            info!("Enclave terminated. Waiting 10s...");
+            std::thread::sleep(std::time::Duration::from_secs(10));
+        } else {
+            info!("No enclaves currently running on this host.");
+        }
+
+        println!("Starting new enclave...");
+        let mut run_command = vec![
+            "nitro-cli",
+            "run-enclave",
+            "--cpu-count", &run_config.num_cpus,
+            "--memory", &run_config.ram_size_mib,
+            "--enclave-cid", "2021",
+            "--eif-path", "enclave.eif",
+        ];
+
+        if run_config.debug_mode == "true" {
+            println!("Debug mode enabled...");
+            run_command.push("--debug-mode");
+        } else {
+            println!("Debug mode disabled...");
+        }
+
+        Self::run_command_capture_stdout(&run_command)?;
+
+        info!("Enclave started... Waiting 5 seconds for warmup.");
+        std::thread::sleep(std::time::Duration::from_secs(10));
+
+        if run_config.debug_mode == "true" {
+            println!("Attaching headless console for running enclaves...");
+            let running_enclaves = Self::run_command_capture_stdout(&["nitro-cli", "describe-enclaves"])?;
+            let enclaves: Value = serde_json::from_str(&running_enclaves)?;
+            let v = vec![];
+            let enclaves_array = enclaves.as_array().unwrap_or(&v);
+            for enclave in enclaves_array {
+                let id = enclave["EnclaveID"].as_str().unwrap();
+                Self::run_command_capture_stdout(&["nitro-cli", "console", "--enclave-id", id])?;
+            }
+        }
+        Ok(())
+    }
+
+    fn run_command_capture_stdout(args: &[&str]) -> Result<String, OrchestrationError> {
+        let output = Command::new(args[0])
+            .args(&args[1..])
+            .stderr(Stdio::inherit())
+            .output()?;
+
+        if !output.status.success() {
+            return Err(OrchestrationError::CommandFailed(format!("Command {:?} failed with exit status: {}", args, output.status).into()));
+        }
+
+        Ok(String::from_utf8_lossy(&output.stdout).to_string())
+    }
+}
+

scripts/start-cage.sh

@@ -4,45 +4,45 @@ INSTANCE_ID=$(curl http://169.254.169.254/latest/dynamic/instance-identity/docum
 
 export EC2_INSTANCE_ID=${INSTANCE_ID}
 
-describe_res=$(nitro-cli describe-enclaves)
-
-enclaves=($(echo "$describe_res" | jq -c '.[]'))
-
-if [[ ${#enclaves[@]} -gt 0 ]]; then
-  echo "[HOST] There is an enclave already running on this host. Terminating it..."
-  nitro-cli terminate-enclave --all
-  echo "[HOST] Enclave terminated. Waiting 10s..."
-  sleep 10
-else
-  echo "[HOST] No enclaves currently running on this host."
-fi
-
-
-# Provision new enclave using customer config
-echo "[HOST] Starting new enclave..."
-
-enclave_run_command="nitro-cli run-enclave --cpu-count $ENCLAVE_NUM_CPUS --memory $ENCLAVE_RAM_SIZE_MIB --enclave-cid 2021 --eif-path enclave.eif"
-
-if [ "$ENCLAVE_DEBUG_MODE" = "true" ] ; then
-  echo "[HOST] Debug mode enabled..."
-  eval "$enclave_run_command --debug-mode"
-else
-  echo "[HOST] Debug mode disabled..."
-  eval "$enclave_run_command"
-fi
-
-echo "[HOST] Enclave started... Waiting 5 seconds for warmup."
-sleep 5
-
-if [ "$ENCLAVE_DEBUG_MODE" = "true" ] ; then
-  # Create stdout streams for any running enclaves
-  echo "[HOST] Attaching headless console for running enclaves..."
-  for id in $(nitro-cli describe-enclaves | jq -r ".[] | .EnclaveID")
-  do
-      # Create console listener
-      nitro-cli console --enclave-id $id &
-  done
-fi
+# describe_res=$(nitro-cli describe-enclaves)
+
+# enclaves=($(echo "$describe_res" | jq -c '.[]'))
+
+# if [[ ${#enclaves[@]} -gt 0 ]]; then
+#   echo "[HOST] There is an enclave already running on this host. Terminating it..."
+#   nitro-cli terminate-enclave --all
+#   echo "[HOST] Enclave terminated. Waiting 10s..."
+#   sleep 10
+# else
+#   echo "[HOST] No enclaves currently running on this host."
+# fi
+
+
+# # Provision new enclave using customer config
+# echo "[HOST] Starting new enclave..."
+
+# enclave_run_command="nitro-cli run-enclave --cpu-count $ENCLAVE_NUM_CPUS --memory $ENCLAVE_RAM_SIZE_MIB --enclave-cid 2021 --eif-path enclave.eif"
+
+# if [ "$ENCLAVE_DEBUG_MODE" = "true" ] ; then
+#   echo "[HOST] Debug mode enabled..."
+#   eval "$enclave_run_command --debug-mode"
+# else
+#   echo "[HOST] Debug mode disabled..."
+#   eval "$enclave_run_command"
+# fi
+
+# echo "[HOST] Enclave started... Waiting 5 seconds for warmup."
+# sleep 5
+
+# if [ "$ENCLAVE_DEBUG_MODE" = "true" ] ; then
+#   # Create stdout streams for any running enclaves
+#   echo "[HOST] Attaching headless console for running enclaves..."
+#   for id in $(nitro-cli describe-enclaves | jq -r ".[] | .EnclaveID")
+#   do
+#       # Create console listener
+#       nitro-cli console --enclave-id $id &
+#   done
+# fi
 
 # Boot control plane
 echo "[HOST] Starting control plane..."