sgx-daemon/src/data.rs

use anyhow::{anyhow, Result};
use detee_shared::types::brain::AppDeployConfig;
use detee_shared::types::brain::Resource as ResourceConfig;
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::HostConfig;

use crate::global::APP_CONFIG_DIR;
use crate::global::APP_NAME_PREFIX;
use crate::global::USED_RESOURCES;

#[derive(Debug, Default, Deserialize, Serialize)]
pub struct HostResources {
    pub existing_apps: HashSet<String>,
    pub reserved_vcpus: u32,
    pub reserved_memory_mb: u32,
    pub reserved_disk_mb: 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)?;
        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)?;
        let res: Self = serde_yml::from_str(&content)?;
        Ok(res)
    }

    pub fn reserve_resources(&mut self, app: &App) -> Result<()> {
        self.reserved_memory_mb += app.app_resource.memory_mb;
        self.reserved_vcpus += app.app_resource.vcpu;
        self.reserved_disk_mb += app.app_resource.disk_mb;

        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_mb -= app.app_resource.memory_mb;
        self.reserved_vcpus -= app.app_resource.vcpu;
        self.reserved_disk_mb -= app.app_resource.disk_mb;

        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()
    }
}

#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct App {
    pub uuid: String,
    pub name: String,
    pub package_path: String,
    pub status: String,
    pub admin: 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> {
        let app_uuid = new_app_req.uuid.clone();

        if host_config.max_cores_per_app < new_app_req.resource.vcpu {
            return Err(anyhow!("too many vcpus for app"));
        }
        if host_config.max_vcpu_reservation
            < host_resource
                .reserved_vcpus
                .saturating_add(new_app_req.resource.vcpu)
        {
            return Err(anyhow!("vcpus not available"));
        }
        if host_config.max_mem_reservation_mb
            < host_resource
                .reserved_memory_mb
                .saturating_add(new_app_req.resource.memory_mb)
        {
            return Err(anyhow!("not enough memory available"));
        }
        if new_app_req.resource.disk_mb < 128 {
            return Err(anyhow!("disk too small"));
        }

        let package_url = new_app_req.package_url.clone();
        let mapped_ports = prepare_port_map(new_app_req.resource.port.clone()).await;
        let app_name = format!("{APP_NAME_PREFIX}-{app_uuid}");

        let unarchive_dir =
            handle_package(package_url, app_uuid.clone(), host_config.delete_archive).await?;

        let exit_code =
            deploy_enclave(&unarchive_dir, app_name.clone(), mapped_ports.clone()).await?;

        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: unarchive_dir,
            status: "running".to_string(),
            admin: 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(&*APP_CONFIG_DIR)?;

        let mut file = File::create(format!("{}{}.yaml", *APP_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(())
    }
}