diff --git a/src/config.rs b/src/config.rs index 65779fc..7168570 100644 --- a/src/config.rs +++ b/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")] - pub reserved_no_of_port: u32, + #[serde(with = "range_format")] + pub public_port_range: Range, 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, 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 { let content = std::fs::read_to_string(path)?; diff --git a/src/data.rs b/src/data.rs index 773e9ff..5414919 100644 --- a/src/data.rs +++ b/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, + 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::>() + } } #[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 = diff --git a/src/main.rs b/src/main.rs index 7d25e79..c2ad756 100644 --- a/src/main.rs +++ b/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 - - (1024 - + host_config.reserved_no_of_port - + host_resource.reserved_host_ports.len() as u32); + let avail_no_of_port = + (host_config.public_port_range.len() - 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; diff --git a/src/utils.rs b/src/utils.rs index 83019f3..8dc5569 100644 --- a/src/utils.rs +++ b/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) -> 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 { - 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 { +pub async fn is_port_available(port: u16) -> bool { TcpListener::bind(&format!("127.0.0.1:{}", port)) .await .is_ok()