switch ports to port pairs

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<String>,
     reserved_if_names: HashSet<String>,
     // sha256sum -> absolute path
-    boot_files: HashMap<String, String>,
+    boot_files: HashSet<String>,
 }
 
 impl Resources {
-    fn reserve_ports(&mut self, nr: u16, config: &Config) -> Vec<u16> {
+    fn reserve_ports(&mut self, extra_ports: usize, config: &Config) -> Vec<u16> {
         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
+                        if *hash == *entry.file_name() {
-                            .insert(hash, path.to_string_lossy().to_string());
+                            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<u16>,
+    fw_ports: Vec<(u16, u16)>,
     nics: Vec<VMNIC>,
     // 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<u16>,
     public_ipv4: bool,
     public_ipv6: bool,
     disk_size: usize,
@@ -293,7 +297,7 @@ impl VM {
         config: &Config,
         res: &mut Resources,
     ) -> Result<Self, VMCreationErrors> {
-        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,8 +357,11 @@ impl VM {
             }
         }
 
-        let fw_ports = res.reserve_ports(req.nr_of_fw_ports, &config);
+        let mut host_ports: Vec<u16> = Vec::new();
-        if fw_ports.len() == 0 {
+        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);
@@ -362,6 +369,11 @@ impl VM {
                 }
                 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]));
+            }
+        }
 
         Ok(VM {
             uuid: req.uuid,
@@ -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);
+        let mut ports = String::new();
-        vars += &format!("INITRD={}\n", VM_BOOT_DIR.to_string() + "/" + &self.dtrfs_sha);
+        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);

src/tcontract.rs

@@ -16,7 +16,7 @@ pub struct ResourceAllocation {
     pub vcpus: usize,
     pub memory: usize,
     pub storage: usize,
-    pub published_ports: Vec<u16>,
+    pub extra_ports: Vec<u16>,
     // storage tier: not part of MVP
     // pub storage_tier: usize,
     pub public_ipv4: Option<String>,