proto/src/types/shared.rs

use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use serde::{Deserialize, Serialize};

use crate::pb::shared as pb_shared;

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Config {
    pub filesystems: Vec<FileEntry>,
    pub environments: Vec<EnvironmentEntry>,
    pub child_processes: Vec<ChildProcess>,
    pub container: Option<Container>,
}

impl From<pb_shared::ManagerConfigPb> for Config {
    fn from(pb_val: pb_shared::ManagerConfigPb) -> Self {
        Config {
            filesystems: pb_val
                .filesystems
                .into_iter()
                .map(FileEntry::from)
                .collect(),
            environments: pb_val
                .environments
                .into_iter()
                .map(EnvironmentEntry::from)
                .collect(),
            child_processes: pb_val
                .child_processes
                .into_iter()
                .map(ChildProcess::from)
                .collect(),
            container: pb_val.container.map(Container::from),
        }
    }
}

impl From<Config> for pb_shared::ManagerConfigPb {
    fn from(val: Config) -> pb_shared::ManagerConfigPb {
        pb_shared::ManagerConfigPb {
            filesystems: val.filesystems.into_iter().map(Into::into).collect(),
            environments: val.environments.into_iter().map(Into::into).collect(),
            child_processes: val.child_processes.into_iter().map(Into::into).collect(),
            container: val.container.map(Into::into),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileEntry {
    pub path: String,
    pub content: Option<FileContent>,
}

impl From<pb_shared::FileEntry> for FileEntry {
    fn from(pb_val: pb_shared::FileEntry) -> Self {
        FileEntry {
            path: pb_val.path,
            content: Some(FileContent::Data(pb_val.content)),
        }
    }
}
impl From<FileEntry> for pb_shared::FileEntry {
    fn from(val: FileEntry) -> pb_shared::FileEntry {
        pb_shared::FileEntry {
            path: val.path,
            content: match val.content {
                Some(FileContent::Data(data)) => data,
                Some(FileContent::Path(path)) => path,
                None => String::new(),
            },
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum FileContent {
    #[serde(rename = "path")]
    Path(String),
    #[serde(rename = "data")]
    Data(String),
}

impl Default for FileContent {
    fn default() -> Self {
        FileContent::Data("".to_string())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnvironmentEntry {
    pub name: String,
    pub value: String,
}

impl From<pb_shared::EnvironmentEntry> for EnvironmentEntry {
    fn from(pb_val: pb_shared::EnvironmentEntry) -> Self {
        EnvironmentEntry {
            name: pb_val.name,
            value: pb_val.value,
        }
    }
}

impl From<EnvironmentEntry> for pb_shared::EnvironmentEntry {
    fn from(val: EnvironmentEntry) -> pb_shared::EnvironmentEntry {
        pb_shared::EnvironmentEntry {
            name: val.name,
            value: val.value,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChildProcess {
    pub path: String,
    pub arguments: Vec<String>,
    pub restart: Option<RestartPolicy>,
}

impl From<pb_shared::ChildProcess> for ChildProcess {
    fn from(pb_val: pb_shared::ChildProcess) -> Self {
        ChildProcess {
            path: pb_val.path,
            arguments: pb_val.arguments,
            restart: pb_val.restart.map(RestartPolicy::from),
        }
    }
}

impl From<ChildProcess> for pb_shared::ChildProcess {
    fn from(val: ChildProcess) -> pb_shared::ChildProcess {
        pb_shared::ChildProcess {
            path: val.path,
            arguments: val.arguments,
            restart: val.restart.map(Into::into),
        }
    }
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
pub struct RestartPolicy {
    pub max_retries: u32,
    pub delay_seconds: u32,
    pub policy: Option<RestartPolicyType>,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum RestartPolicyType {
    Always(bool),
    OnNonZeroExit(bool),
}

impl Default for RestartPolicyType {
    fn default() -> Self {
        RestartPolicyType::Always(true)
    }
}
impl From<pb_shared::RestartPolicy> for RestartPolicy {
    fn from(pb_val: pb_shared::RestartPolicy) -> Self {
        RestartPolicy {
            max_retries: pb_val.max_retries,
            delay_seconds: pb_val.delay_seconds,
            policy: match pb_val.policy_type {
                Some(pb_shared::restart_policy::PolicyType::Always(_)) => {
                    Some(RestartPolicyType::Always(true))
                }
                Some(pb_shared::restart_policy::PolicyType::OnNonZeroExit(_)) => {
                    Some(RestartPolicyType::OnNonZeroExit(true))
                }
                None => None,
            },
        }
    }
}

impl From<RestartPolicy> for pb_shared::RestartPolicy {
    fn from(val: RestartPolicy) -> pb_shared::RestartPolicy {
        pb_shared::RestartPolicy {
            max_retries: val.max_retries,
            delay_seconds: val.delay_seconds,
            policy_type: match val.policy {
                Some(RestartPolicyType::Always(_)) => {
                    Some(pb_shared::restart_policy::PolicyType::Always(true))
                }
                Some(RestartPolicyType::OnNonZeroExit(_)) => {
                    Some(pb_shared::restart_policy::PolicyType::OnNonZeroExit(true))
                }
                None => None,
            },
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Container {
    pub package_url: String,
    pub resource: Option<Resource>,
    #[serde(default)]
    pub uuid: String,
    #[serde(default)]
    pub admin_pubkey: String,
    pub node_pubkey: String,
}

impl From<pb_shared::Container> for Container {
    fn from(pb_val: pb_shared::Container) -> Self {
        Self {
            package_url: pb_val.package_url,
            resource: pb_val.resource.map(Resource::from),
            uuid: pb_val.uuid,
            admin_pubkey: pb_val.admin_pubkey,
            node_pubkey: pb_val.node_pubkey,
        }
    }
}

impl From<Container> for pb_shared::Container {
    fn from(val: Container) -> pb_shared::Container {
        pb_shared::Container {
            package_url: val.package_url,
            resource: val.resource.map(Into::into),
            // uuid: val.uuid.map(Into::into),
            uuid: val.uuid,
            admin_pubkey: val.admin_pubkey,
            node_pubkey: val.node_pubkey,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct Resource {
    pub memory_mb: u32,
    pub disk_mb: u32,
    pub vcpu: u32,
    pub port: Vec<u32>,
}

impl From<pb_shared::Resource> for Resource {
    fn from(pb_val: pb_shared::Resource) -> Self {
        Self {
            memory_mb: pb_val.memory_mb,
            disk_mb: pb_val.disk_mb,
            vcpu: pb_val.vcpu,
            port: pb_val.ports,
        }
    }
}
impl From<Resource> for pb_shared::Resource {
    fn from(val: Resource) -> pb_shared::Resource {
        pb_shared::Resource {
            memory_mb: val.memory_mb,
            disk_mb: val.disk_mb,
            vcpu: val.vcpu,
            ports: val.port,
        }
    }
}

impl From<(u16, u16)> for pb_shared::MappedPort {
    fn from(val: (u16, u16)) -> Self {
        Self {
            host_port: val.0 as u32,
            container_port: val.1 as u32,
        }
    }
}

impl From<pb_shared::MappedPort> for (u16, u16) {
    fn from(val: pb_shared::MappedPort) -> Self {
        (val.host_port as u16, val.container_port as u16)
    }
}

impl Config {
    pub fn from_path(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
        let config_str = std::fs::read_to_string(path)?;
        Ok(serde_yml::from_str(&config_str)?)
    }

    pub fn load_data(mut self) -> Result<Self, Box<dyn std::error::Error>> {
        self.filesystems.iter_mut().for_each(|x| {
            if let Some(FileContent::Path(path)) = &x.content {
                let content =
                    std::fs::read(path).unwrap_or_else(|_| panic!("Unable to read file {path}"));
                let encoded = BASE64.encode(content);
                x.content = Some(FileContent::Data(encoded));
            }
        });

        Ok(self)
    }
}