sgx-daemon/src/data.rs

// SPDX-License-Identifier: Apache-2.0

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 crate::container::delete_enclave;
use crate::container::deploy_enclave;
use crate::utils::handle_package;
use crate::utils::is_port_available;
use crate::HostConfig;

use crate::global::APP_NAME_PREFIX;
use crate::global::DEPLOYED_APPS_CONFIG_DIR;
use crate::global::USED_RESOURCES_PATH;

#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct HostResources {
    pub existing_apps: HashSet<String>,
    pub reserved_vcpus: u32,
    pub reserved_memory_mib: u32,
    pub reserved_disk_mib: u32,
    pub reserved_host_ports: HashSet<u16>,
}

impl HostResources {
    pub fn new() -> Self {
        // TODO: compute host resources
        Self::default()
    }

    // TODO: implement load and save
    fn save_to_disk(&self) -> Result<()> {
        let mut file = File::create(&*USED_RESOURCES_PATH)?;
        file.write_all(serde_yml::to_string(self)?.as_bytes())?;
        Ok(())
    }

    pub fn load_from_disk() -> Result<Self> {
        let content = std::fs::read_to_string(&*USED_RESOURCES_PATH).unwrap_or_else(|_| {
            let host_resource = Self::new();
            host_resource.save_to_disk().unwrap();
            serde_yml::to_string(&host_resource).unwrap()
        });
        let res: Self = serde_yml::from_str(&content)?;
        Ok(res)
    }

    pub fn reserve_resources(&mut self, app: &App) -> Result<()> {
        self.reserved_memory_mib += app.app_resource.memory_mib;
        self.reserved_vcpus += app.app_resource.vcpus;
        self.reserved_disk_mib += app.app_resource.disk_size_mib;

        for (port, _) in app.mapped_ports.iter() {
            self.reserved_host_ports.insert(*port);
        }
        self.existing_apps.insert(app.uuid.clone());

        self.save_to_disk()
    }

    pub fn un_reserve_resources(&mut self, app: &App) -> Result<()> {
        self.reserved_memory_mib -= app.app_resource.memory_mib;
        self.reserved_vcpus -= app.app_resource.vcpus;
        self.reserved_disk_mib -= app.app_resource.disk_size_mib;

        for (port, _) in app.mapped_ports.iter() {
            self.reserved_host_ports.remove(port);
        }
        self.existing_apps.take(&app.uuid).ok_or_else(|| {
            log::error!("App \"{}\" not found", app.uuid);
            anyhow!("App \"{}\" not found", app.uuid)
        })?;

        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)]
pub struct App {
    pub uuid: String,
    pub name: String,
    pub package_path: String,
    pub status: String,
    pub admin_pubkey: String,
    pub app_resource: ResourceConfig,
    pub mapped_ports: Vec<(u16, u16)>,
}

impl App {
    pub async fn new(
        new_app_req: AppDeployConfig,
        host_config: &HostConfig,
        host_resource: &mut HostResources,
    ) -> Result<Self> {
        if new_app_req.node_unit_price < host_config.price {
            return Err(anyhow!("price is too low"));
        }

        if host_resource.existing_apps.contains(&new_app_req.uuid) {
            let content = std::fs::read_to_string(format!(
                "{}/{}.yaml",
                *DEPLOYED_APPS_CONFIG_DIR, &new_app_req.uuid
            ))?;
            let app: App = serde_yml::from_str(&content)?;
            return Err(anyhow!("app already exists\n{:?}", app));
        }

        let app_uuid = new_app_req.uuid.clone();

        if host_config.max_cores_per_app < new_app_req.resource.vcpus {
            return Err(anyhow!("too many vcpus for app"));
        }

        if host_config.max_memory_mib_per_app < new_app_req.resource.memory_mib {
            return Err(anyhow!("too much memory for app"));
        }

        // TODO: revert state if any error
        if host_config.max_vcpu_reservation
            < host_resource
                .reserved_vcpus
                .saturating_add(new_app_req.resource.vcpus)
        {
            return Err(anyhow!("vcpus not available"));
        }
        if host_config.max_mem_reservation_mib
            < host_resource
                .reserved_memory_mib
                .saturating_add(new_app_req.resource.memory_mib)
        {
            return Err(anyhow!("not enough memory available"));
        }
        if new_app_req.resource.disk_size_mib < 1 {
            return Err(anyhow!("disk too small"));
        }

        let package_url = new_app_req.package_url.clone();
        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 =
            handle_package(package_url, app_uuid.clone(), host_config.delete_archive).await?;

        let exit_code = deploy_enclave(
            &package_path,
            app_name.clone(),
            mapped_ports.clone(),
            new_app_req.hratls_pubkey,
            new_app_req.resource.clone(),
        )?;

        if exit_code != 0 {
            // TODO: cleanup unarchive_dir
            return Err(anyhow!("Failed to deploy enclave"));
        }

        let app_instance = Self {
            uuid: app_uuid,
            name: app_name,
            package_path,
            status: "running".to_string(),
            admin_pubkey: new_app_req.admin_pubkey,
            app_resource: new_app_req.resource,
            mapped_ports,
        };
        app_instance.write_config()?;
        host_resource
            .existing_apps
            .insert(app_instance.uuid.clone());

        host_resource.reserve_resources(&app_instance)?;

        Ok(app_instance)
    }

    fn write_config(&self) -> Result<()> {
        std::fs::create_dir_all(&*DEPLOYED_APPS_CONFIG_DIR)?;

        let mut file = File::create(format!("{}/{}.yaml", *DEPLOYED_APPS_CONFIG_DIR, &self.uuid))?;
        file.write_all(serde_yml::to_string(self)?.as_bytes())?;
        Ok(())
    }

    pub async fn delete_app(&self, host_resource: &mut HostResources) -> Result<()> {
        let container_name = format!("{APP_NAME_PREFIX}-{}", &self.uuid);
        delete_enclave(container_name)?;
        host_resource.un_reserve_resources(self)?;

        Ok(())
    }
}