Skip to content

Commit

Permalink
clean up
Browse files Browse the repository at this point in the history
  • Loading branch information
hanneary committed Dec 10, 2024
1 parent 277cb4a commit 4a8888e
Showing 1 changed file with 40 additions and 11 deletions.
51 changes: 40 additions & 11 deletions control-plane/src/orchestration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand All @@ -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" {
Expand All @@ -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())
Expand Down

0 comments on commit 4a8888e

Please sign in to comment.