From 7be5c1abd7a2e02906e18fcb69fffe0bc5548b43 Mon Sep 17 00:00:00 2001
From: Noor <noormohammedb@protonmail.com>
Date: Tue, 21 Jan 2025 12:19:06 +0000
Subject: [PATCH] feat: add dynamic port mapping for enclave deployment

---
 Cargo.lock     |  2 +-
 Cargo.toml     |  1 +
 src/data.rs    |  4 +++-
 src/package.rs | 48 ++++++++++++++++++++++++++++++++++++++++++++----
 4 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 16f7f20..9860c2f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -305,6 +305,7 @@ dependencies = [
  "flate2",
  "prost",
  "prost-types",
+ "rand",
  "reqwest",
  "tar",
  "tokio",
@@ -315,7 +316,6 @@ dependencies = [
 [[package]]
 name = "detee-shared"
 version = "0.1.0"
-source = "git+ssh://git@gitea.detee.cloud/noormohammedb/detee-shared#358c2d84ca092cad5df69ac19cc6acbfb85ee8d7"
 dependencies = [
  "base64",
  "prost",
diff --git a/Cargo.toml b/Cargo.toml
index bd206ed..5e71c42 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,6 +17,7 @@ reqwest = "0.12.12"
 flate2 = "1.0.35"
 tar = "0.4.43"
 anyhow = "1.0.95"
+rand = "0.8.5"
 
 [build-dependencies]
 tonic-build = "0.12.3"
diff --git a/src/data.rs b/src/data.rs
index 4b9fc33..0f3fb82 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -31,7 +31,9 @@ impl DaemonState {
         req_data: ContainerConfig,
         unarchive_dir: String,
     ) -> Result<(), Box<dyn std::error::Error>> {
-        deploy_enclave(&unarchive_dir).await?;
+        let publishing_ports = req_data.resource.clone().unwrap().port;
+
+        deploy_enclave(&unarchive_dir, publishing_ports).await?;
 
         let container = Container {
             id: "123".to_string(),
diff --git a/src/package.rs b/src/package.rs
index 252956d..5717efa 100644
--- a/src/package.rs
+++ b/src/package.rs
@@ -1,12 +1,14 @@
 use anyhow::{anyhow, Result};
 use chrono::Utc;
 use flate2::read::GzDecoder;
+use rand::Rng;
 use reqwest::Client;
 use std::io::BufReader;
 use std::path::Path;
 use std::process::Command;
 use tar::Archive;
 use tokio::io::AsyncWriteExt;
+use tokio::net::TcpListener;
 use tokio::{fs, fs::File};
 
 pub async fn handle_package(package_url: String) -> Result<String> {
@@ -41,21 +43,28 @@ pub async fn handle_package(package_url: String) -> Result<String> {
 pub async fn deploy_enclave(
     enclave_path: &str,
     // enclave_name: String,
-    // publishing_ports: Vec<u32>,
+    publishing_ports: Vec<u32>,
     // ...
 ) -> Result<(), Box<dyn std::error::Error>> {
+    let port_map = prepare_port_map(publishing_ports).await;
+
+    let port_maping_string = port_map
+        .iter()
+        .map(|(host, container)| format!("-p {host}:{container}"))
+        .collect::<Vec<_>>()
+        .join(" ");
+
     println!("Deploying enclave: {:?}", enclave_path);
     let docker_deploy_str = format!(
-        "docker run -v {enclave_path}/enclave_packager:/enclave_packager --device /dev/sgx/enclave --device /dev/sgx/provision -p 34500:34500 noormohammedb/occlum-enclave:v1");
+        "docker run -v {enclave_path}/enclave_packager:/enclave_packager --device /dev/sgx/enclave --device /dev/sgx/provision {port_maping_string} noormohammedb/occlum-enclave:v1");
 
     println!("{}", &docker_deploy_str);
 
-    let child = Command::new("sh")
+    let _child = Command::new("sh")
         .arg("-c")
         .arg(docker_deploy_str)
         .spawn()?;
 
-    dbg!(child);
     Ok(())
 }
 
@@ -69,3 +78,34 @@ pub async fn download_file(url: &str, file_path: &Path) -> Result<(), Box<dyn st
 
     Ok(())
 }
+
+async fn prepare_port_map(mut publishing_ports: Vec<u32>) -> Vec<(u16, u16)> {
+    publishing_ports.insert(0, 34500);
+    let mut maped_ports = vec![];
+    for port in publishing_ports {
+        if is_port_available(port as u16).await {
+            maped_ports.push((port as u16, port as u16));
+        } else {
+            let host_port = get_random_available_port().await.unwrap();
+            maped_ports.push((host_port, port as u16));
+        }
+    }
+    maped_ports
+}
+
+pub async fn get_random_available_port() -> Option<u16> {
+    let mut rng = rand::rngs::OsRng;
+    for _ in 0..1000 {
+        let port = rng.gen_range(15000..45000);
+        if is_port_available(port).await {
+            return Some(port);
+        }
+    }
+    None
+}
+
+async fn is_port_available(port: u16) -> bool {
+    TcpListener::bind(&format!("127.0.0.1:{}", port))
+        .await
+        .is_ok()
+}