From b0cf13ddfbfc5cce959dfc3b6ebce10e1b8990fb Mon Sep 17 00:00:00 2001 From: ghe0 Date: Wed, 11 Dec 2024 04:17:57 +0200 Subject: [PATCH] switch ports to port pairs --- src/state.rs | 72 ++++++++++++++++++++++++++++++++---------------- src/tcontract.rs | 2 +- 2 files changed, 50 insertions(+), 24 deletions(-) diff --git a/src/state.rs b/src/state.rs index 49bacff..2cbe1c9 100644 --- a/src/state.rs +++ b/src/state.rs @@ -4,7 +4,6 @@ use crate::constants::*; use anyhow::anyhow; use anyhow::Result; use sha2::{Digest, Sha256}; -use std::collections::HashMap; use std::collections::HashSet; use std::fs; use std::fs::remove_file; @@ -23,20 +22,21 @@ pub struct Resources { reserved_ips: HashSet, reserved_if_names: HashSet, // sha256sum -> absolute path - boot_files: HashMap, + boot_files: HashSet, } impl Resources { - fn reserve_ports(&mut self, nr: u16, config: &Config) -> Vec { + fn reserve_ports(&mut self, extra_ports: usize, config: &Config) -> Vec { use rand::Rng; - if config.public_port_range.len() < self.reserved_ports.len() + nr as usize { + let total_ports = extra_ports + 1; + if config.public_port_range.len() < self.reserved_ports.len() + total_ports as usize { return Vec::new(); } - if nr > config.max_ports_per_vm { + if total_ports > config.max_ports_per_vm as usize { return Vec::new(); } let mut published_ports = Vec::new(); - for _ in 0..nr { + for _ in 0..total_ports { for _ in 0..5 { let port = rand::thread_rng().gen_range(config.public_port_range.clone()); if self.reserved_ports.insert(port) { @@ -45,6 +45,7 @@ impl Resources { break; } } + published_ports.sort(); published_ports } @@ -153,8 +154,11 @@ impl Resources { if path.is_file() { match compute_sha256(&path) { Ok(hash) => { - self.boot_files - .insert(hash, path.to_string_lossy().to_string()); + if *hash == *entry.file_name() { + self.boot_files.insert(hash); + } else { + // TODO: rename file and insert + } } Err(e) => return Err(anyhow!("Error computing hash for {:?}: {}", path, e)), } @@ -164,10 +168,10 @@ impl Resources { } fn download_boot_file(&mut self, url: String, sha: String) -> Result<()> { - if !self.boot_files.contains_key(&sha) { + if !self.boot_files.contains(&sha) { download_and_check_sha(&url, &sha)?; } - self.boot_files.insert(sha, url); + self.boot_files.insert(sha); Ok(()) } } @@ -246,7 +250,7 @@ pub struct VM { uuid: String, hostname: String, admin_key: String, - fw_ports: Vec, + fw_ports: Vec<(u16, u16)>, nics: Vec, // currently hardcoded to EPYC-v4 // cpu_type: String, @@ -263,7 +267,7 @@ pub struct NewVMRequest { uuid: String, hostname: String, admin_key: String, - nr_of_fw_ports: u16, + extra_ports: Vec, public_ipv4: bool, public_ipv6: bool, disk_size: usize, @@ -293,7 +297,7 @@ impl VM { config: &Config, res: &mut Resources, ) -> Result { - if req.nr_of_fw_ports > 0 && req.public_ipv4 { + if req.extra_ports.len() > 0 && req.public_ipv4 { return Err(VMCreationErrors::NATandIPv4Conflict); } if config.max_cores_per_vm < req.vcpus { @@ -316,7 +320,7 @@ impl VM { if let Err(dtrfs_file_error) = res.download_boot_file(req.dtrfs_url, req.dtrfs_sha.clone()) { return Err(VMCreationErrors::BootFileError(format!( - "Could not get kernel: {dtrfs_file_error:?}" + "Could not get dtrfs: {dtrfs_file_error:?}" ))); }; @@ -353,14 +357,22 @@ impl VM { } } - let fw_ports = res.reserve_ports(req.nr_of_fw_ports, &config); - if fw_ports.len() == 0 { - for nic in vm_nics { - for ip in nic.ips { - res.reserved_ips.remove(&ip.address); + let mut host_ports: Vec = Vec::new(); + let mut port_pairs: Vec<(u16, u16)> = Vec::new(); + if !req.public_ipv4 { + host_ports.append(res.reserve_ports(req.extra_ports.len(), &config).as_mut()); + if host_ports.len() == 0 { + for nic in vm_nics { + for ip in nic.ips { + res.reserved_ips.remove(&ip.address); + } } + return Err(VMCreationErrors::NotEnoughPorts); + } + port_pairs.push((host_ports[0], 22)); + for i in 0..req.extra_ports.len() { + port_pairs.push((host_ports[i+1], req.extra_ports[i])); } - return Err(VMCreationErrors::NotEnoughPorts); } Ok(VM { @@ -373,7 +385,7 @@ impl VM { disk_size: req.disk_size, kernel_sha: req.kernel_sha, dtrfs_sha: req.dtrfs_sha, - fw_ports, + fw_ports: port_pairs, }) } @@ -433,8 +445,22 @@ impl VM { i += 1; } - vars += &format!("KERNEL={}\n", VM_BOOT_DIR.to_string() + "/" + &self.kernel_sha); - vars += &format!("INITRD={}\n", VM_BOOT_DIR.to_string() + "/" + &self.dtrfs_sha); + let mut ports = String::new(); + for port in self.fw_ports.iter() { + ports += &format!("{}:{} ", port.0, port.1); + } + if ports != "" { + vars += &format!("NAT_PORT_FW={}", ports.trim_end()); + } + + vars += &format!( + "KERNEL={}\n", + VM_BOOT_DIR.to_string() + "/" + &self.kernel_sha + ); + vars += &format!( + "INITRD={}\n", + VM_BOOT_DIR.to_string() + "/" + &self.dtrfs_sha + ); vars += &format!("PARAMS={}\n", self.kernel_params()); vars += &format!("CPU_TYPE={}\n", QEMU_VM_CPU_TYPE); vars += &format!("VCPUS={}\n", self.vcpus); diff --git a/src/tcontract.rs b/src/tcontract.rs index 1c5fea1..a2c0bf8 100644 --- a/src/tcontract.rs +++ b/src/tcontract.rs @@ -16,7 +16,7 @@ pub struct ResourceAllocation { pub vcpus: usize, pub memory: usize, pub storage: usize, - pub published_ports: Vec, + pub extra_ports: Vec, // storage tier: not part of MVP // pub storage_tier: usize, pub public_ipv4: Option,