hacker-challenge/src/main.rs

mod datastore;
mod grpc;
mod http_server;
mod persistence;
mod solana;

use crate::persistence::SealError;
use crate::{
    grpc::challenge::NodeUpdate, persistence::KeysFile, persistence::SealedFile, solana::SolClient,
};
use datastore::State;
use detee_sgx::{InstanceMeasurement, RaTlsConfig};
use std::{
    fs::File,
    io::Error,
    io::{BufRead, BufReader},
    sync::Arc,
};
use tokio::{
    sync::{broadcast, broadcast::Sender},
    task::JoinSet,
    time::{sleep, Duration},
};

const INIT_NODES_FILE: &str = "/host/detee_challenge_nodes";
const KEYS_FILE: &str = "/host/main/TRY_TO_HACK_THIS";
const MAX_CONNECTIONS: usize = 3;

#[cfg(feature = "test")]
async fn resolve_my_ip() -> Result<String, Error> {
    let node_ip = File::open("/host/detee_node_ip")?;
    let mut reader = BufReader::new(node_ip);
    let mut ip = String::new();
    reader.read_line(&mut ip)?;
    Ok(ip)
}

#[cfg(not(feature = "test"))]
async fn resolve_my_ip() -> Result<String, Error> {
    let err = "Can't resolve my external IP, try again";
    let ip = public_ip::addr_v4().await.ok_or(Error::new(std::io::ErrorKind::Other, err))?;
    Ok(format!("{}", ip))
}

pub async fn heartbeat_cron(my_ip: String, state: Arc<State>, tx: Sender<NodeUpdate>) {
    loop {
        sleep(Duration::from_secs(60)).await;
        let _ = tx.send((my_ip.clone(), state.get_my_info()).into());
        state.remove_inactive_nodes();
    }
}

async fn get_sol_client(state: Arc<State>, ratls_config: RaTlsConfig) -> SolClient {
    match KeysFile::read(KEYS_FILE) {
        Ok(keys_file) => {
            let sol_client = SolClient::try_from(keys_file)
                .map_err(|e| {
                    println!("Read malformed keys from the disk: {e}");
                    state.increase_net_attacks();
                })
                .unwrap();
            println!(
                "Found the following wallet saved to disk: {}",
                sol_client.get_wallet_pubkey()
            );
            println!("Loading token mint address {}", sol_client.get_token_address());
            return sol_client;
        }
        Err(SealError::Attack(e)) => {
            println!("The local keys file is corrupted: {}", e);
            state.increase_disk_attacks();
        }
        Err(e) => {
            println!("Could not read keys file: {e}");
        }
    };

    let init_nodes = match File::open(INIT_NODES_FILE) {
        Ok(init_nodes) => init_nodes,
        Err(_) => {
            println!("Can't initialize using init nodes from {INIT_NODES_FILE}");
            println!("Starting a new network with a new key...");
            return SolClient::new().await;
        }
    };

    let init_nodes_reader = BufReader::new(init_nodes);
    for init_node_ip in init_nodes_reader.lines().map(|l| l.unwrap()) {
        match grpc::client::key_grabber(state.clone(), init_node_ip, ratls_config.clone()).await {
            Ok(keys) => {
                let sol_client = SolClient::try_from(keys.clone())
                    .map_err(|e| {
                        println!("Received malformed keys from the network: {e}");
                        state.increase_net_attacks();
                    })
                    .unwrap();
                println!(
                    "Got keypair from the network. Joining the network using wallet {}",
                    sol_client.get_wallet_pubkey()
                );
                println!("The address of the Token is {}", sol_client.get_token_address());
                println!("Saving this data to disk in the file {KEYS_FILE}");
                if let Err(e) = sol_client.get_keys_file().write(KEYS_FILE) {
                    println!("Could not save data to disk: {e}");
                }
                return sol_client;
            }
            Err(e) => {
                println!("Could not get keypair: {e:?}");
            }
        };
    }
    panic!("could not get keypair.");
}

#[tokio::main]
async fn main() {
    env_logger::init_from_env(env_logger::Env::default().default_filter_or("warn"));
    let ratls_config = RaTlsConfig::new()
        .allow_instance_measurement(InstanceMeasurement::new().with_current_mrenclave().unwrap());
    let my_ip = resolve_my_ip().await.unwrap();
    println!("Starting on IP {}", my_ip);

    let state = Arc::new(State::new(my_ip.clone()));
    let sol_client = Arc::new(get_sol_client(state.clone(), ratls_config.clone()).await);

    let (tx, _) = broadcast::channel(500);

    let mut tasks = JoinSet::new();

    tasks.spawn(heartbeat_cron(my_ip.clone(), state.clone(), tx.clone()));
    tasks.spawn(http_server::init(state.clone(), sol_client.clone()));
    tasks.spawn(
        grpc::server::MyServer::init(
            state.clone(),
            sol_client.get_keys(),
            ratls_config.clone(),
            tx.clone(),
        )
        .start(),
    );

    if let Ok(input) = std::fs::read_to_string(INIT_NODES_FILE) {
        for line in input.lines() {
            tasks.spawn(
                grpc::client::ConnManager::init(
                    my_ip.clone(),
                    state.clone(),
                    ratls_config.clone(),
                    tx.clone(),
                )
                .start_with_node(line.to_string()),
            );
        }
    }

    for _ in 0..MAX_CONNECTIONS {
        tasks.spawn(
            grpc::client::ConnManager::init(
                my_ip.clone(),
                state.clone(),
                ratls_config.clone(),
                tx.clone(),
            )
            .start(),
        );
    }

    while let Some(Ok(_)) = tasks.join_next().await {}

    // task panicked
    println!("Shutting down...");
}