brain/src/grpc/types.rs

use crate::constants::{ACCOUNT, APP_NODE, ID_ALPHABET, NEW_APP_REQ, NEW_VM_REQ, VM_NODE};
use crate::db::prelude as db;
use detee_shared::app_proto::AppNodeListResp;
use detee_shared::common_proto::MappedPort;
use detee_shared::general_proto::{Account, AccountBalance, ListOperatorsResp};
use detee_shared::{app_proto::*, vm_proto::*};
use nanoid::nanoid;

use surrealdb::RecordId;

impl From<db::Account> for AccountBalance {
    fn from(account: db::Account) -> Self {
        AccountBalance { balance: account.balance, tmp_locked: account.tmp_locked }
    }
}

impl From<db::Account> for Account {
    fn from(account: db::Account) -> Self {
        Account {
            pubkey: account.id.to_string(),
            balance: account.balance,
            tmp_locked: account.tmp_locked,
        }
    }
}

impl From<NewVmReq> for db::NewVmReq {
    fn from(new_vm_req: NewVmReq) -> Self {
        Self {
            id: RecordId::from((NEW_VM_REQ, nanoid!(40, &ID_ALPHABET))),
            admin: RecordId::from((ACCOUNT, new_vm_req.admin_pubkey)),
            vm_node: RecordId::from((VM_NODE, new_vm_req.node_pubkey)),
            hostname: new_vm_req.hostname,
            extra_ports: new_vm_req.extra_ports,
            public_ipv4: new_vm_req.public_ipv4,
            public_ipv6: new_vm_req.public_ipv6,
            disk_size_gb: new_vm_req.disk_size_gb,
            vcpus: new_vm_req.vcpus,
            memory_mb: new_vm_req.memory_mb,
            kernel_url: new_vm_req.kernel_url,
            kernel_sha: new_vm_req.kernel_sha,
            dtrfs_url: new_vm_req.dtrfs_url,
            dtrfs_sha: new_vm_req.dtrfs_sha,
            price_per_unit: new_vm_req.price_per_unit,
            locked_nano: new_vm_req.locked_nano,
            created_at: surrealdb::sql::Datetime::default(),
            error: String::new(),
        }
    }
}

impl From<db::NewVmReq> for NewVmReq {
    fn from(new_vm_req: db::NewVmReq) -> Self {
        Self {
            uuid: new_vm_req.id.key().to_string(),
            hostname: new_vm_req.hostname,
            admin_pubkey: new_vm_req.admin.key().to_string(),
            node_pubkey: new_vm_req.vm_node.key().to_string(),
            extra_ports: new_vm_req.extra_ports,
            public_ipv4: new_vm_req.public_ipv4,
            public_ipv6: new_vm_req.public_ipv6,
            disk_size_gb: new_vm_req.disk_size_gb,
            vcpus: new_vm_req.vcpus,
            memory_mb: new_vm_req.memory_mb,
            kernel_url: new_vm_req.kernel_url,
            kernel_sha: new_vm_req.kernel_sha,
            dtrfs_url: new_vm_req.dtrfs_url,
            dtrfs_sha: new_vm_req.dtrfs_sha,
            price_per_unit: new_vm_req.price_per_unit,
            locked_nano: new_vm_req.locked_nano,
        }
    }
}

impl From<db::WrappedMeasurement> for NewVmResp {
    fn from(resp: db::WrappedMeasurement) -> Self {
        match resp {
            db::WrappedMeasurement::Args(uuid, args) => {
                Self { uuid, error: String::new(), args: Some(args) }
            }
            db::WrappedMeasurement::Error(uuid, error) => NewVmResp { uuid, error, args: None },
        }
    }
}

// TODO: NewVmResp is identical to UpdateVmResp so we can actually remove it from proto
impl From<db::WrappedMeasurement> for UpdateVmResp {
    fn from(resp: db::WrappedMeasurement) -> Self {
        match resp {
            db::WrappedMeasurement::Args(uuid, args) => {
                Self { uuid, error: String::new(), args: Some(args) }
            }
            db::WrappedMeasurement::Error(uuid, error) => Self { uuid, error, args: None },
        }
    }
}

impl From<UpdateVmReq> for db::UpdateVmReq {
    fn from(new_vm_req: UpdateVmReq) -> Self {
        Self {
            id: RecordId::from((NEW_VM_REQ, new_vm_req.uuid)),
            admin: RecordId::from((ACCOUNT, new_vm_req.admin_pubkey)),
            // vm_node gets modified later, and only if the db::UpdateVmReq is required
            vm_node: RecordId::from((VM_NODE, String::new())),
            disk_size_gb: new_vm_req.disk_size_gb,
            vcpus: new_vm_req.vcpus,
            memory_mb: new_vm_req.memory_mb,
            kernel_url: new_vm_req.kernel_url,
            kernel_sha: new_vm_req.kernel_sha,
            dtrfs_url: new_vm_req.dtrfs_url,
            dtrfs_sha: new_vm_req.dtrfs_sha,
            created_at: surrealdb::sql::Datetime::default(),
            error: String::new(),
        }
    }
}

impl From<db::UpdateVmReq> for UpdateVmReq {
    fn from(update_vm_req: db::UpdateVmReq) -> Self {
        Self {
            uuid: update_vm_req.id.key().to_string(),
            // daemon does not care about VM hostname
            hostname: String::new(),
            admin_pubkey: update_vm_req.admin.key().to_string(),
            disk_size_gb: update_vm_req.disk_size_gb,
            vcpus: update_vm_req.vcpus,
            memory_mb: update_vm_req.memory_mb,
            kernel_url: update_vm_req.kernel_url,
            kernel_sha: update_vm_req.kernel_sha,
            dtrfs_url: update_vm_req.dtrfs_url,
            dtrfs_sha: update_vm_req.dtrfs_sha,
        }
    }
}

impl From<db::DeletedVm> for DeleteVmReq {
    fn from(delete_vm_req: db::DeletedVm) -> Self {
        Self {
            uuid: delete_vm_req.id.key().to_string(),
            admin_pubkey: delete_vm_req.admin.key().to_string(),
        }
    }
}

impl From<db::VmDaemonMsg> for BrainVmMessage {
    fn from(notification: db::VmDaemonMsg) -> Self {
        match notification {
            db::VmDaemonMsg::Create(new_vm_req) => {
                BrainVmMessage { msg: Some(brain_vm_message::Msg::NewVmReq(new_vm_req.into())) }
            }
            db::VmDaemonMsg::Update(update_vm_req) => BrainVmMessage {
                msg: Some(brain_vm_message::Msg::UpdateVmReq(update_vm_req.into())),
            },
            db::VmDaemonMsg::Delete(deleted_vm) => {
                BrainVmMessage { msg: Some(brain_vm_message::Msg::DeleteVm(deleted_vm.into())) }
            }
        }
    }
}

impl From<db::ActiveVmWithNode> for VmContract {
    fn from(db_c: db::ActiveVmWithNode) -> Self {
        let mut exposed_ports = Vec::new();
        for port in db_c.mapped_ports.iter() {
            exposed_ports.push(port.0);
        }
        VmContract {
            uuid: db_c.id.key().to_string(),
            hostname: db_c.hostname.clone(),
            admin_pubkey: db_c.admin.key().to_string(),
            node_pubkey: db_c.vm_node.id.key().to_string(),
            node_ip: db_c.vm_node.ip.clone(),
            location: format!(
                "{}, {}, {}",
                db_c.vm_node.city, db_c.vm_node.region, db_c.vm_node.country
            ),
            memory_mb: db_c.memory_mb,
            vcpus: db_c.vcpus,
            disk_size_gb: db_c.disk_size_gb,
            mapped_ports: db_c
                .mapped_ports
                .iter()
                .map(|(h, g)| MappedPort { host_port: *h, guest_port: *g })
                .collect(),
            vm_public_ipv6: db_c.public_ipv6.clone(),
            vm_public_ipv4: db_c.public_ipv4.clone(),
            locked_nano: db_c.locked_nano,
            dtrfs_sha: db_c.dtrfs_sha.clone(),
            kernel_sha: db_c.kernel_sha.clone(),
            nano_per_minute: db_c.price_per_minute(),
            created_at: db_c.created_at.to_rfc3339(),
            // TODO: remove updated_at from the proto
            // This will get moved to VM history (users will be able to
            // query old contracts, which also shows updates of existing contracts).
            updated_at: db_c.created_at.to_rfc3339(),
            collected_at: db_c.collected_at.to_rfc3339(),
        }
    }
}

impl From<db::Error> for tonic::Status {
    fn from(e: db::Error) -> Self {
        Self::internal(format!("Internal error: {e}"))
    }
}

impl From<db::Operator> for ListOperatorsResp {
    fn from(db_o: db::Operator) -> Self {
        ListOperatorsResp {
            pubkey: db_o.account.key().to_string(),
            escrow: db_o.escrow,
            email: db_o.email,
            app_nodes: db_o.app_nodes,
            vm_nodes: db_o.vm_nodes,
            reports: db_o.reports,
        }
    }
}

impl From<db::VmNodeWithReports> for VmNodeListResp {
    fn from(vm_node: db::VmNodeWithReports) -> Self {
        Self {
            operator: vm_node.operator.key().to_string(),
            node_pubkey: vm_node.id.key().to_string(),
            country: vm_node.country,
            region: vm_node.region,
            city: vm_node.city,
            ip: vm_node.ip,
            reports: vm_node.reports.iter().map(|n| n.reason.clone()).collect(),
            price: vm_node.price,
        }
    }
}

impl From<db::AppNodeWithReports> for AppNodeListResp {
    fn from(app_node: db::AppNodeWithReports) -> Self {
        Self {
            operator: app_node.operator.key().to_string(),
            node_pubkey: app_node.id.key().to_string(),
            country: app_node.country,
            region: app_node.region,
            city: app_node.city,
            ip: app_node.ip,
            reports: app_node.reports.iter().map(|n| n.reason.clone()).collect(),
            price: app_node.price,
        }
    }
}

impl From<VmNodeResources> for db::VmNodeResources {
    fn from(res: VmNodeResources) -> Self {
        Self {
            avail_mem_mb: res.avail_memory_mb,
            avail_vcpus: res.avail_vcpus,
            avail_storage_gbs: res.avail_storage_gb,
            avail_ipv4: res.avail_ipv4,
            avail_ipv6: res.avail_ipv6,
            avail_ports: res.avail_ports,
            max_ports_per_vm: res.max_ports_per_vm,
        }
    }
}

impl From<db::ActiveAppWithNode> for AppContract {
    fn from(value: db::ActiveAppWithNode) -> Self {
        let public_package_mr_enclave =
            Some(hex::decode(value.mr_enclave.clone()).unwrap_or_default());

        AppContract {
            uuid: value.id.key().to_string(),
            package_url: value.package_url,
            admin_pubkey: value.admin.key().to_string(),
            node_pubkey: value.app_node.id.key().to_string(),
            public_ipv4: value.host_ipv4,
            resource: Some(AppResource {
                memory_mb: value.memory_mb,
                disk_size_gb: value.disk_size_gb,
                vcpus: value.vcpus,
                ports: value.mapped_ports.iter().map(|(_, g)| *g).collect(),
            }),
            mapped_ports: value
                .mapped_ports
                .iter()
                .map(|(h, g)| MappedPort { host_port: *h, guest_port: *g })
                .collect(),

            created_at: value.created_at.to_rfc3339(),
            updated_at: value.created_at.to_rfc3339(),
            nano_per_minute: value.price_per_unit,
            locked_nano: value.locked_nano,
            collected_at: value.collected_at.to_rfc3339(),
            hratls_pubkey: value.mr_enclave,
            public_package_mr_enclave,
            app_name: value.app_name,
        }
    }
}

impl From<NewAppReq> for db::NewAppReq {
    fn from(val: NewAppReq) -> Self {
        let resource = val.resource.unwrap_or_default();

        let mr_enclave = val
            .public_package_mr_enclave
            .unwrap_or_default()
            .iter()
            .fold(String::new(), |acc, x| acc + &format!("{x:02x?}"));

        Self {
            id: RecordId::from((NEW_APP_REQ, nanoid!(40, &ID_ALPHABET))),
            admin: RecordId::from((ACCOUNT, val.admin_pubkey)),
            app_node: RecordId::from((APP_NODE, val.node_pubkey)),
            app_name: val.app_name,
            package_url: val.package_url,
            mr_enclave,
            hratls_pubkey: val.hratls_pubkey,
            ports: resource.ports,
            memory_mb: resource.memory_mb,
            vcpus: resource.vcpus,
            disk_size_gb: resource.disk_size_gb,
            locked_nano: val.locked_nano,
            price_per_unit: val.price_per_unit,
            error: String::new(),
            created_at: surrealdb::sql::Datetime::default(),
        }
    }
}

impl From<db::NewAppReq> for NewAppReq {
    fn from(value: db::NewAppReq) -> Self {
        let resource = AppResource {
            vcpus: value.vcpus,
            memory_mb: value.memory_mb,
            disk_size_gb: value.disk_size_gb,
            ports: value.ports,
        };
        let mr_enclave = Some(hex::decode(value.mr_enclave).unwrap_or_default());

        Self {
            package_url: value.package_url,
            node_pubkey: value.app_node.key().to_string(),
            resource: Some(resource),
            uuid: value.id.key().to_string(),
            admin_pubkey: value.admin.key().to_string(),
            price_per_unit: value.price_per_unit,
            locked_nano: value.locked_nano,
            hratls_pubkey: value.hratls_pubkey,
            public_package_mr_enclave: mr_enclave,
            app_name: value.app_name,
        }
    }
}

impl From<db::DeletedApp> for DelAppReq {
    fn from(value: db::DeletedApp) -> Self {
        Self { uuid: value.id.key().to_string(), admin_pubkey: value.admin.key().to_string() }
    }
}

impl From<db::AppDaemonMsg> for BrainMessageApp {
    fn from(value: db::AppDaemonMsg) -> Self {
        match value {
            db::AppDaemonMsg::Create(new_app_req) => {
                BrainMessageApp { msg: Some(brain_message_app::Msg::NewAppReq(new_app_req.into())) }
            }
            db::AppDaemonMsg::Delete(del_app_req) => BrainMessageApp {
                msg: Some(brain_message_app::Msg::DeleteAppReq(del_app_req.into())),
            },
        }
    }
}

impl From<AppNodeResources> for db::AppNodeResources {
    fn from(value: AppNodeResources) -> Self {
        Self {
            avail_no_of_port: value.avail_no_of_port,
            avail_vcpus: value.avail_vcpus,
            avail_memory_mb: value.avail_memory_mb,
            avail_storage_mb: value.avail_storage_mb,
            max_ports_per_app: value.max_ports_per_app,
        }
    }
}

impl From<db::ActiveApp> for NewAppRes {
    fn from(val: db::ActiveApp) -> Self {
        let mapped_ports = val
            .mapped_ports
            .iter()
            .map(|(h, g)| MappedPort { host_port: *h, guest_port: *g })
            .collect();
        Self {
            uuid: val.id.key().to_string(),
            ip_address: val.host_ipv4,
            mapped_ports,
            error: String::new(),
        }
    }
}