sgx-daemon/src/main.rs

pub mod container;
pub mod data;
pub mod grpc;
pub mod utils;

use std::time::Duration;

pub use data::DaemonState;
use detee_shared::pb::daemon::brain_message;
use detee_shared::pb::daemon::daemon_message;
use detee_shared::pb::daemon::BrainMessage;
use detee_shared::pb::daemon::DaemonMessage;
use detee_shared::pb::daemon::NewContainerRes;
use detee_shared::pb::shared::ContainerContracts;

use detee_shared::pb::shared::MappedPort;
use detee_shared::types::shared::Container as ContainerConfig;

use tokio::sync::mpsc::Receiver;
use tokio::sync::mpsc::Sender;
use tokio::time::sleep;
use utils::handle_package;

const NODE_PUBKEY: &str = "0xd0837609aedd53854651210327db90f5c2626188a00e940bbc9eea2c7e6838b7";
const ADMIN_PUBKEY: &str = "0x28a3a71197250b0fa4dd0f86288e07ec9cc78ce3338e21e2ebef84dd7780e3eb";

#[derive(Debug)]
pub struct Config {
    pub brain_url: String,
}

impl Default for Config {
    fn default() -> Self {
        let brain_url =
            std::env::var("BRAIN_URL").unwrap_or_else(|_| "http://127.0.0.1:31337".to_string());
        Self { brain_url }
    }
}

#[derive(Debug)]
pub struct ContainerHandler {
    pub receiver: Receiver<BrainMessage>,
    pub sender: Sender<DaemonMessage>,
    pub config: Config,
    pub data: DaemonState,
}

impl ContainerHandler {
    pub fn new(receiver: Receiver<BrainMessage>, sender: Sender<DaemonMessage>) -> Self {
        Self {
            receiver,
            sender,
            config: Config::default(),
            data: DaemonState::default(),
        }
    }

    fn handle_contracts(&mut self, contracts: Vec<ContainerContracts>) {
        dbg!(&contracts);
    }

    async fn run(mut self) {
        while let Some(brain_msg) = self.receiver.recv().await {
            match brain_msg.msg {
                Some(brain_message::Msg::NewContainerReq(msg)) => {
                    self.handle_new_container_req(msg.into()).await;
                }

                Some(brain_message::Msg::DeleteContainer(msg)) => {
                    dbg!(&msg);
                }

                Some(brain_message::Msg::ListContainer(msg)) => {
                    dbg!(&msg);
                }
                None => {
                    log::error!("Brain disconnected");
                    break;
                }
            }
        }
    }

    async fn handle_new_container_req(&mut self, new_container_req: ContainerConfig) {
        dbg!(&new_container_req);

        let container_uuid = new_container_req.uuid.clone();

        let unarchive_dir = match handle_package(new_container_req.package_url.clone()).await {
            Ok(unarchive_dir) => unarchive_dir,
            Err(e) => {
                let res = DaemonMessage {
                    msg: Some(daemon_message::Msg::NewContainerResp(NewContainerRes {
                        uuid: new_container_req.uuid,
                        status: "failed".to_string(),
                        error: e.to_string(),
                        ..Default::default()
                    })),
                };

                println!("sending response {:?}", res);
                let _ = self.sender.send(res).await;
                return;
            }
        };

        let mapped_ports = match self
            .data
            .create_new_container(new_container_req, unarchive_dir)
            .await
        {
            Ok(mapped_ports) => mapped_ports.into_iter().map(MappedPort::from).collect(),
            Err(e) => {
                let res = DaemonMessage {
                    msg: Some(daemon_message::Msg::NewContainerResp(NewContainerRes {
                        uuid: container_uuid,
                        status: "failed".to_string(),
                        error: e.to_string(),
                        ..Default::default()
                    })),
                };

                println!("sending response {:?}", res);
                let _ = self.sender.send(res).await;
                return;
            }
        };

        let res = DaemonMessage {
            msg: Some(daemon_message::Msg::NewContainerResp(NewContainerRes {
                uuid: container_uuid,
                status: "Success".to_string(),
                error: "".to_string(),
                ip_address: "".to_string(),
                mapped_ports,
            })),
        };

        println!("sending response {:?}", res);
        let _ = self.sender.send(res).await;
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    set_logging();
    log::info!("Detee daemon running");

    loop {
        let (brain_msg_tx, brain_msg_rx) = tokio::sync::mpsc::channel(6);
        let (daemon_msg_tx, daemon_msg_rx) = tokio::sync::mpsc::channel(6);

        let mut container_handler = ContainerHandler::new(brain_msg_rx, daemon_msg_tx.clone());
        let brain_url = container_handler.config.brain_url.clone();

        match grpc::register_node(&container_handler.config).await {
            Ok(container_contracts) => container_handler.handle_contracts(container_contracts),

            Err(e) => log::error!("Failed to connect to brain: {e}"),
        }

        tokio::spawn(async move {
            container_handler.run().await;
        });

        log::info!("Connecting to brain...");
        if let Err(e) = grpc::connect_and_run(grpc::ConnectionData {
            brain_url,
            brain_msg_tx,
            daemon_msg_rx,
            daemon_msg_tx,
        })
        .await
        {
            log::error!("The connection broke: {e}");
        }
        sleep(Duration::from_secs(3)).await;
    }
}

fn set_logging() {
    let log_level = match std::env::var("LOG_LEVEL") {
        Ok(val) => match val.as_str() {
            "DEBUG" => log::LevelFilter::Debug,
            "INFO" => log::LevelFilter::Info,
            _ => log::LevelFilter::Error,
        },
        _ => log::LevelFilter::Warn,
    };
    env_logger::builder()
        .filter_level(log_level)
        .format_timestamp(None)
        .init();
}