detee-cli/src/snp/deploy.rs

use super::{
    grpc::{self, proto},
    injector, Distro, Dtrfs, Error, VmSshArgs, DEFAULT_ARCHLINUX, DEFAULT_DTRFS,
};
use crate::config::Config;
use crate::utils::block_on;
use log::{debug, info};
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize)]
pub enum IPv4Config {
    PublishPorts(Vec<u32>),
    PublicIPv4,
}

//  TODO: push this out of snp module
#[derive(Serialize, Deserialize, Default)]
pub struct Location {
    pub node_ip: Option<String>,
    pub country: Option<String>,
    pub region: Option<String>,
    pub city: Option<String>,
}

impl From<&str> for Location {
    fn from(s: &str) -> Self {
        match s {
            "Canada" => Self { country: Some("CA".to_string()), ..Default::default() },
            "Montreal" => Self { city: Some("Montréal".to_string()), ..Default::default() },
            "Vancouver" => Self { city: Some("Vancouver".to_string()), ..Default::default() },
            "US" => Self { country: Some("US".to_string()), ..Default::default() },
            "California" => Self { country: Some("US".to_string()), ..Default::default() },
            "France" => Self { country: Some("FR".to_string()), ..Default::default() },
            "GB" => Self { country: Some("GB".to_string()), ..Default::default() },
            "Random" => Self { ..Default::default() },
            "DE" => Self { country: Some("DE".to_string()), ..Default::default() },
            _ => Self { city: Some("Vancouver".to_string()), ..Default::default() },
        }
    }
}

#[derive(Serialize, Deserialize)]
pub struct Request {
    pub hostname: String,
    pub hours: u32,
    // price per unit per minute
    pub price: u64,
    pub location: Location,
    pub ipv4: IPv4Config,
    pub public_ipv6: bool,
    pub vcpus: u32,
    pub memory_mb: u32,
    pub disk_size_gb: u32,
    pub dtrfs: Option<Dtrfs>,
    pub distro: Option<Distro>,
}

impl Request {
    pub fn load_from_yaml(config_path: &str) -> Result<VmSshArgs, Error> {
        let new_vm_config = Self::load_from_file(config_path)?;
        new_vm_config.deploy()
    }

    fn load_from_file(config_path: &str) -> Result<Self, Error> {
        let content = std::fs::read_to_string(config_path)?;
        let new_vm_config: Self = serde_yaml::from_str(&content)?;
        Ok(new_vm_config)
    }

    pub fn deploy(&self) -> Result<VmSshArgs, Error> {
        let (node_ip, new_vm_resp) = self.send_vm_request()?;
        info!("Got confirmation from the node {node_ip} that VM started.");
        debug!("IPs and ports assigned by node are: {new_vm_resp:#?}");
        if !new_vm_resp.error.is_empty() {
            return Err(Error::Node(new_vm_resp.error));
        }
        let (kernel_sha, dtrfs_sha) = match self.dtrfs.clone() {
            Some(dtrfs) => (dtrfs.kernel_sha, dtrfs.dtrfs_sha),
            None => (DEFAULT_DTRFS.kernel_sha.clone(), DEFAULT_DTRFS.dtrfs_sha.clone()),
        };
        let args = new_vm_resp.args.ok_or(Error::NoMeasurement)?;
        let measurement_args = injector::Args {
            uuid: new_vm_resp.uuid.clone(),
            vcpus: self.vcpus,
            kernel: kernel_sha,
            initrd: dtrfs_sha,
            args: args.clone(),
        };
        let measurement = measurement_args.get_measurement()?;
        let (template_url, template_sha) = match &self.distro {
            Some(distro) => (distro.template_url.clone(), distro.template_sha.clone()),
            None => {
                (DEFAULT_ARCHLINUX.template_url.clone(), DEFAULT_ARCHLINUX.template_sha.clone())
            }
        };
        let mut ssh_args = injector::execute(
            measurement,
            args.dtrfs_api_endpoint.clone(),
            Some((&template_url, &template_sha)),
            &self.hostname,
        )?;
        ssh_args.uuid = new_vm_resp.uuid;
        ssh_args.hostname = self.hostname.clone();
        let _ = super::append_uuid_list(&ssh_args.uuid, &ssh_args.hostname);
        Ok(ssh_args)
    }

    // returns node IP and data regarding the new VM
    fn send_vm_request(&self) -> Result<(String, proto::NewVmResp), Error> {
        let admin_pubkey = Config::get_detee_wallet()?;
        let node = self.get_node()?;
        let (extra_ports, public_ipv4): (Vec<u32>, bool) = match &self.ipv4 {
            IPv4Config::PublishPorts(vec) => (vec.to_vec(), false),
            IPv4Config::PublicIPv4 => (Vec::new(), true),
        };
        let (kernel_url, kernel_sha, dtrfs_url, dtrfs_sha) = match self.dtrfs.clone() {
            Some(dtrfs) => (dtrfs.kernel_url, dtrfs.kernel_sha, dtrfs.dtrfs_url, dtrfs.dtrfs_sha),
            None => (
                DEFAULT_DTRFS.kernel_url.clone(),
                DEFAULT_DTRFS.kernel_sha.clone(),
                DEFAULT_DTRFS.dtrfs_url.clone(),
                DEFAULT_DTRFS.dtrfs_sha.clone(),
            ),
        };
        let locked_nano = super::calculate_nanolp(
            self.vcpus,
            self.memory_mb,
            self.disk_size_gb,
            public_ipv4,
            self.hours,
            self.price,
        );
        let brain_req = proto::NewVmReq {
            uuid: String::new(),
            hostname: self.hostname.clone(),
            admin_pubkey,
            node_pubkey: node.node_pubkey,
            extra_ports,
            public_ipv4,
            public_ipv6: self.public_ipv6,
            disk_size_gb: self.disk_size_gb,
            vcpus: self.vcpus,
            memory_mb: self.memory_mb,
            kernel_url,
            kernel_sha,
            dtrfs_url,
            dtrfs_sha,
            price_per_unit: self.price,
            locked_nano,
        };
        let new_vm_resp = block_on(grpc::create_vm(brain_req))?;
        if !new_vm_resp.error.is_empty() {
            return Err(Error::Node(new_vm_resp.error));
        }
        Ok((node.ip, new_vm_resp))
    }

    pub fn get_node(&self) -> Result<proto::VmNodeListResp, Error> {
        let (free_ports, offers_ipv4) = match &self.ipv4 {
            IPv4Config::PublishPorts(vec) => (vec.len() as u32, false),
            IPv4Config::PublicIPv4 => (0, true),
        };
        let filters = proto::VmNodeFilters {
            free_ports,
            offers_ipv4,
            offers_ipv6: self.public_ipv6,
            vcpus: self.vcpus,
            memory_mb: self.memory_mb,
            storage_gb: self.disk_size_gb,
            country: self.location.country.clone().unwrap_or_default(),
            region: self.location.region.clone().unwrap_or_default(),
            city: self.location.city.clone().unwrap_or_default(),
            ip: self.location.node_ip.clone().unwrap_or_default(),
            node_pubkey: String::new(),
        };
        match block_on(grpc::get_one_node(filters)) {
            Ok(node) => Ok(node),
            Err(e) => {
                log::error!("Coult not get node from brain: {e:?}");
                Err(Error::NoValidNodeFound)
            }
        }
    }
}