Refactor host port logic to port range

HostConfig to use public_port_range and update port mapping logic

src/config.rs

@@ -1,9 +1,10 @@
 use anyhow::Result;
-use serde::{Deserialize, Serialize};
+use serde::Deserialize;
+use std::ops::Range;
 
 use crate::global::IP_INFO;
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Deserialize)]
 pub struct HostConfig {
     pub brain_url: String,
     #[serde(default = "retrieve_node_ip")]
@@ -18,20 +19,35 @@ pub struct HostConfig {
     // price per unit per minute
     pub price: u64,
 
-    #[serde(default = "default_reserved_no_of_port")]
+    #[serde(with = "range_format")]
-    pub reserved_no_of_port: u32,
+    pub public_port_range: Range<u16>,
 
     pub delete_archive: bool,
 }
 
-fn default_reserved_no_of_port() -> u32 {
-    16
-}
-
 fn retrieve_node_ip() -> String {
     IP_INFO.ip.clone()
 }
 
+mod range_format {
+    use serde::{Deserialize, Deserializer, Serialize};
+    use std::ops::Range;
+
+    pub fn deserialize<'de, D>(deserializer: D) -> Result<Range<u16>, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let range_repr = RangeRepr::deserialize(deserializer)?;
+        Ok(range_repr.start..range_repr.end)
+    }
+
+    #[derive(Serialize, Deserialize)]
+    struct RangeRepr {
+        start: u16,
+        end: u16,
+    }
+}
+
 impl HostConfig {
     pub fn load_from_disk(path: &str) -> Result<Self> {
         let content = std::fs::read_to_string(path)?;

src/data.rs

@@ -1,16 +1,16 @@
 use anyhow::{anyhow, Result};
 use detee_shared::sgx::types::brain::AppDeployConfig;
 use detee_shared::sgx::types::brain::Resource as ResourceConfig;
+use rand::Rng;
+use serde::{Deserialize, Serialize};
 use std::collections::HashSet;
 use std::fs::File;
 use std::io::Write;
 
-use serde::{Deserialize, Serialize};
-
 use crate::container::delete_enclave;
 use crate::container::deploy_enclave;
 use crate::utils::handle_package;
-use crate::utils::prepare_port_map;
+use crate::utils::is_port_available;
 use crate::HostConfig;
 
 use crate::global::APP_NAME_PREFIX;
@@ -77,6 +77,36 @@ impl HostResources {
 
         self.save_to_disk()
     }
+    async fn prepare_port_map(
+        &self,
+        mut publishing_ports: Vec<u32>,
+        host_config: &HostConfig,
+    ) -> Vec<(u16, u16)> {
+        publishing_ports.insert(0, 34500);
+
+        let mut host_ports = vec![];
+
+        for _ in 0..publishing_ports.len() {
+            for _ in 0..10 {
+                let port = rand::rngs::OsRng.gen_range(host_config.public_port_range.clone());
+                if !self.reserved_host_ports.contains(&{ port }) && is_port_available(port).await {
+                    host_ports.push(port);
+                    break;
+                }
+            }
+        }
+
+        if host_ports.len() < publishing_ports.len() {
+            return vec![];
+        }
+
+        host_ports.sort();
+
+        host_ports
+            .into_iter()
+            .zip(publishing_ports.into_iter().map(|f| f as u16))
+            .collect::<Vec<(u16, u16)>>()
+    }
 }
 
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]
@@ -139,7 +169,12 @@ impl App {
         }
 
         let package_url = new_app_req.package_url.clone();
-        let mapped_ports = prepare_port_map(new_app_req.resource.port.clone()).await;
+        let mapped_ports = host_resource
+            .prepare_port_map(new_app_req.resource.port.clone(), host_config)
+            .await;
+        if mapped_ports.is_empty() {
+            return Err(anyhow!("not enough ports available"));
+        }
         let app_name = format!("{APP_NAME_PREFIX}-{app_uuid}");
 
         let package_path =

src/main.rs

@@ -158,10 +158,9 @@ impl AppHandler {
         let host_resource = self.host_resource.clone();
 
         let node_pubkey = PUBLIC_KEY.to_string();
-        let avail_no_of_port = 65535
+        let avail_no_of_port =
-            - (1024
+            (host_config.public_port_range.len() - host_resource.reserved_host_ports.len()) as u32;
-                + host_config.reserved_no_of_port
+
-                + host_resource.reserved_host_ports.len() as u32);
         let avail_vcpus = host_config.max_vcpu_reservation - host_resource.reserved_vcpus;
         let avail_memory_mb = host_config.max_mem_reservation_mb - host_resource.reserved_memory_mb;
         let avail_storage_mb = host_config.max_disk_reservation_mb - host_resource.reserved_disk_mb;

src/utils.rs

@@ -1,6 +1,5 @@
 use anyhow::{anyhow, Result};
 use flate2::read::GzDecoder;
-use rand::Rng;
 use reqwest::Client;
 use std::io::BufReader;
 use std::path::Path;
@@ -78,32 +77,7 @@ pub async fn download_file(url: &str, file_path: &Path) -> Result<(), Box<dyn st
     Ok(())
 }
 
-pub async fn prepare_port_map(mut publishing_ports: Vec<u32>) -> Vec<(u16, u16)> {
+pub async fn is_port_available(port: u16) -> bool {
-    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()