trifid/tfclient/src/apiworker.rs

use std::sync::mpsc::{Receiver, TryRecvError};
use base64::Engine;
use chrono::Local;
use log::{error, info, warn};
use url::Url;
use trifid_pki::ca::NebulaCAPool;
use trifid_pki::cert::{serialize_ed25519_public, serialize_x25519_public};
use trifid_pki::ed25519_dalek::{SecretKey, SigningKey};
use trifid_pki::rand_core::OsRng;
use trifid_pki::x25519_dalek::StaticSecret;
use crate::api::{APIResponse, enroll, EnrollRequest};
use crate::config::{load_cdata, save_cdata, TFClientConfig};
use crate::daemon::ThreadMessageSender;

pub enum APIWorkerMessage {
    Shutdown,
    Enroll { code: String }
}

pub fn apiworker_main(config: TFClientConfig, instance: String, url: String, _transmitters: ThreadMessageSender, rx: Receiver<APIWorkerMessage>) {
    let server = Url::parse(&url).unwrap();

    // Generate dhPubkey and edPubkey if it doesn't exist
    // Load vardata
    let mut vdata = match load_cdata(&instance) {
        Ok(d) => d,
        Err(e) => {
            error!("Error loading vdata: {}", e);
            error!("APIWorker exiting with error");
            return;
        }
    };

    if vdata.ed_privkey.is_none() {
        info!("Generating ed25519 key");
        let mut csprng = OsRng;
        let key = SigningKey::generate(&mut csprng);
        let ed_key_bytes = key.to_bytes().to_vec();
        vdata.ed_privkey = Some(ed_key_bytes.try_into().unwrap());
    }
    if vdata.dh_privkey.is_none() {
        info!("Generating ecdh key");
        let mut csprng = OsRng;
        let key = StaticSecret::new(&mut csprng);
        let dh_key_bytes = key.to_bytes();
        vdata.dh_privkey = Some(dh_key_bytes);
    }

    info!("Loading keys");
    let ed_key = SigningKey::from_bytes(&SecretKey::from(vdata.ed_privkey.unwrap()));
    let dh_key = StaticSecret::from(vdata.dh_privkey.unwrap());
    info!("Keys loaded successfully");

    // Save vardata
    match save_cdata(&instance, vdata) {
        Ok(_) => (),
        Err(e) => {
            error!("Error saving vdata: {}", e);
            error!("APIWorker exiting with error");
            return;
        }
    }
    loop {
        match rx.try_recv() {
            Ok(msg) => {
                match msg {
                    APIWorkerMessage::Shutdown => {
                        info!("recv on command socket: shutdown, stopping");
                        break;
                    },
                    APIWorkerMessage::Enroll { code } => {
                        info!("recv on command socket: enroll {}", code);
                        let mut cdata = match load_cdata(&instance) {
                            Ok(c) => c,
                            Err(e) => {
                                error!("error in api worker thread: {}", e);
                                error!("APIWorker exiting with error");
                                return;
                            }
                        };
                        if cdata.host_id.is_some() {
                            warn!("enrollment failed: already enrolled");
                            continue;
                        }
                        let dh_encoded = base64::engine::general_purpose::STANDARD.encode(serialize_x25519_public(&dh_key.to_bytes()));
                        let ed_encoded = base64::engine::general_purpose::STANDARD.encode(serialize_ed25519_public(&ed_key.to_bytes()));
                        let req = EnrollRequest {
                            code,
                            dh_pubkey: dh_encoded,
                            ed_pubkey: ed_encoded,
                            timestamp: Local::now().format("%Y-%m-%dT%H:%M:%S.%f%:z").to_string(),
                        };
                        let res = match enroll(&server, &req) {
                            Ok(res) => res,
                            Err(e) => {
                                error!("error in api worker thread: {}", e);
                                error!("APIWorker exiting with error");
                                return;
                            }
                        };
                        let resp = match res {
                            APIResponse::Error(e) => {
                                error!("error with enrollment: {}: {}", e.errors[0].code, e.errors[0].message);
                                continue;
                            }
                            APIResponse::Success(resp) => resp
                        };

                        info!("Enrolled with server. Host-ID {} config count {}", resp.data.host_id, resp.data.counter);
                        info!("KeyPool {}, org {} {}", resp.data.trusted_keys, resp.data.organization.name, resp.data.organization.id);
                        info!("Config: {}", resp.data.config);

                        // Decode the CAPool and config
                        let key_pool_pem = match base64::engine::general_purpose::STANDARD.decode(resp.data.trusted_keys) {
                            Ok(p) => p,
                            Err(e) => {
                                error!("error with enrollment: {}", e);
                                continue;
                            }
                        };

                        let config = match base64::engine::general_purpose::STANDARD.decode(resp.data.config) {
                            Ok(p) => p,
                            Err(e) => {
                                error!("error with enrollment: {}", e);
                                continue;
                            }
                        };
                        let config_str = String::from_utf8(config).unwrap();


                        cdata.host_id = Some(resp.data.host_id);
                        cdata.counter = resp.data.counter as i32;
                        cdata.key_pool = Some(String::from_utf8(key_pool_pem).unwrap());
                        cdata.org_name = Some(resp.data.organization.name);
                        cdata.org_id = Some(resp.data.organization.id);
                        cdata.config = Some(config_str);

                        // Save vardata
                        match save_cdata(&instance, cdata) {
                            Ok(_) => (),
                            Err(e) => {
                                error!("Error saving cdata: {}", e);
                                error!("APIWorker exiting with error");
                                return;
                            }
                        }
                    }
                }
            },
            Err(e) => {
                match e {
                    TryRecvError::Empty => {}
                    TryRecvError::Disconnected => {
                        error!("apiworker command socket disconnected, shutting down to prevent orphaning");
                        break;
                    }
                }
            }
        }
    }
}