trifid/tfclient/src/daemon.rs

use std::sync::mpsc;
use std::sync::mpsc::Sender;
use std::thread;
use log::{error, info};

use crate::apiworker::{apiworker_main, APIWorkerMessage};
use crate::config::load_config;

use crate::nebulaworker::{nebulaworker_main, NebulaWorkerMessage};
use crate::socketworker::{socketworker_main, SocketWorkerMessage};
use crate::timerworker::{timer_main, TimerWorkerMessage};
use crate::util::check_server_url;

pub fn daemon_main(name: String, server: String) {
    // Validate the `server`
    check_server_url(&server);

    info!("Loading config...");
    let config = match load_config(&name) {
        Ok(cfg) => cfg,
        Err(e) => {
            error!("Error loading configuration: {}", e);
            std::process::exit(1);
        }
    };

    info!("Creating transmitter");

    let (tx_api, rx_api) = mpsc::channel::<APIWorkerMessage>();
    let (tx_socket, rx_socket) = mpsc::channel::<SocketWorkerMessage>();
    let (tx_nebula, rx_nebula) = mpsc::channel::<NebulaWorkerMessage>();
    let (tx_timer, rx_timer) = mpsc::channel::<TimerWorkerMessage>();

    let transmitter = ThreadMessageSender {
        socket_thread: tx_socket,
        api_thread: tx_api,
        nebula_thread: tx_nebula,
        timer_thread: tx_timer,
    };

    let mainthread_transmitter = transmitter.clone();

    info!("Setting signal trap...");

    match ctrlc::set_handler(move || {
        info!("Ctrl-C detected. Stopping threads...");
        match mainthread_transmitter.nebula_thread.send(NebulaWorkerMessage::Shutdown) {
            Ok(_) => (),
            Err(e) => {
                error!("Error sending shutdown message to nebula worker thread: {}", e);
            }
        }
        match mainthread_transmitter.api_thread.send(APIWorkerMessage::Shutdown) {
            Ok(_) => (),
            Err(e) => {
                error!("Error sending shutdown message to api worker thread: {}", e);
            }
        }
        match mainthread_transmitter.socket_thread.send(SocketWorkerMessage::Shutdown) {
            Ok(_) => (),
            Err(e) => {
                error!("Error sending shutdown message to socket worker thread: {}", e);
            }
        }
        match mainthread_transmitter.timer_thread.send(TimerWorkerMessage::Shutdown) {
            Ok(_) => (),
            Err(e) => {
                error!("Error sending shutdown message to timer worker thread: {}", e);
            }
        }
    }) {
        Ok(_) => (),
        Err(e) => {
            error!("Unable to set sigtrap: {}", e);
            std::process::exit(1);
        }
    }

    info!("Starting API thread...");

    let config_api = config.clone();
    let transmitter_api = transmitter.clone();
    let name_api = name.clone();
    let server_api = server.clone();
    let api_thread = thread::spawn(move || {
        apiworker_main(config_api, name_api, server_api,transmitter_api, rx_api);
    });

    info!("Starting Nebula thread...");
    let config_nebula = config.clone();
    let transmitter_nebula = transmitter.clone();
    let name_nebula = name.clone();
    let nebula_thread = thread::spawn(move || {
        nebulaworker_main(config_nebula, name_nebula, transmitter_nebula, rx_nebula);
    });

    info!("Starting timer thread...");
    if !config.disable_automatic_config_updates {
        let timer_transmitter = transmitter.clone();
        let timer_thread = thread::spawn(move || {
            timer_main(timer_transmitter, rx_timer);
        });
        info!("Waiting for timer thread to exit...");
        match timer_thread.join() {
            Ok(_) => (),
            Err(_) => {
                error!("Error waiting for timer thread to exit.");
                std::process::exit(1);
            }
        }
        info!("Timer thread exited");
    } else {
        info!("automatic config updates have been disabled - not starting timer thread");
    }

    info!("Starting socket worker thread...");
    let name_socket = name.clone();
    let socket_thread = thread::spawn(move || {
        socketworker_main(config, name_socket, transmitter, rx_socket);
    });

    info!("Waiting for socket thread to exit...");
    match socket_thread.join() {
        Ok(_) => (),
        Err(_) => {
            error!("Error waiting for socket thread to exit.");
            std::process::exit(1);
        }
    }
    info!("Socket thread exited");

    info!("Waiting for API thread to exit...");
    match api_thread.join() {
        Ok(_) => (),
        Err(_) => {
            error!("Error waiting for api thread to exit.");
            std::process::exit(1);
        }
    }
    info!("API thread exited");

    info!("Waiting for Nebula thread to exit...");
    match nebula_thread.join() {
        Ok(_) => (),
        Err(_) => {
            error!("Error waiting for nebula thread to exit.");
            std::process::exit(1);
        }
    }
    info!("Nebula thread exited");

    info!("All threads exited");
}

#[derive(Clone)]
pub struct ThreadMessageSender {
    pub socket_thread: Sender<SocketWorkerMessage>,
    pub api_thread: Sender<APIWorkerMessage>,
    pub nebula_thread: Sender<NebulaWorkerMessage>,
    pub timer_thread: Sender<TimerWorkerMessage>
}