sopht 0.3.0

//! program implementing a sopht server/daemon. this server handles connections
//! on a single thread and completely sequentially, so if multiple connections
//! are made, later connections' commands will not be honored until the
//! first connection is closed
//!
use anyhow::{anyhow, Result};
use log::{debug, error, info, warn};
use sopht::{
    create_local_dirs, socket_file_path, Command, Connection, Error, Message,
    ProcessStatus, Response, State,
};
use std::os::unix::net::UnixListener;
use std::path::Path;
use std::sync::atomic::{self, AtomicBool};
use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender};
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use std::time::Duration;

/// type alias representing whether or not the server should shut down. this is
/// just to make the return type of [`handle_connection`] clearer
type Quit = bool;

/// information about the thread handling the process management and upkeep.
/// for more information on the manager thread, see [`spawn_manager_thread`]
struct ManagerThread {
    /// handle to the manager thread. this is joined when the server exits
    pub handle: JoinHandle<()>,
    /// the manager thread receives and executes commands sent over this sender
    pub command_sender: Sender<Command>,
    /// the manager thread sends exactly one response per command over this
    /// receiver
    pub response_receiver: Receiver<Response>,
}

fn main() {
    env_logger::init_from_env(
        env_logger::Env::new()
            .filter_or("SOPHT_LOG", "info")
            .write_style("SOPHT_LOG_STYLE"),
    );

    if let Err(e) = create_local_dirs() {
        error!("failed to create sopht local dirs. error: {}", e);
        return;
    }

    let socket_file = match socket_file_path() {
        Ok(fp) => fp,
        Err(e) => {
            error!("could not obtain path to socket file. error: {}", e);
            return;
        }
    };

    let interrupt = Arc::new(AtomicBool::new(false));

    signal_hook::flag::register_conditional_shutdown(
        signal_hook::consts::SIGINT,
        1,
        interrupt.clone(),
    )
    .expect("could not register conditional shutdown for SIGINT");
    signal_hook::flag::register(
        signal_hook::consts::SIGTERM,
        Arc::clone(&interrupt),
    )
    .expect("could not set SIGTERM handler");
    signal_hook::flag::register(
        signal_hook::consts::SIGINT,
        Arc::clone(&interrupt),
    )
    .expect("could not set SIGINT handler");

    if socket_file.exists() {
        error!("could not start server. socket already exists");
        return;
    }

    let state = State::new(1);
    let mut manager = spawn_manager_thread(state, interrupt.clone());

    if let Err(e) = listen(&mut manager, &socket_file, interrupt.clone()) {
        interrupt.store(true, atomic::Ordering::Relaxed);
        error!("failed to open listen server. error: {}", e);
    }

    manager
        .handle
        .join()
        .expect("failed to join manager thread. this might need some cleanup!");

    match std::fs::remove_file(socket_file) {
        Ok(_) => info!("shutdown OK"),
        Err(e) => error!("could not remove socket file. error: {}", e),
    }
}

/// consumes a [`State`] and spawns a thread responsible for two things:
///   * receiving and executing commands on that state, then sending responses
///     back to the calling thread
///   * calling [`State::upkeep`] on the that state periodically
fn spawn_manager_thread(
    mut state: State,
    interrupt: Arc<AtomicBool>,
) -> ManagerThread {
    let (cmdtx, cmdrx) = mpsc::channel::<Command>();
    let (rsptx, rsprx) = mpsc::channel::<Response>();
    // this is the interval at which the server will perform upkeep
    // in the future, this will be configurable!
    let timeout = Duration::from_secs(1);
    let handle = thread::spawn(move || loop {
        let cmd: Option<Command> = match cmdrx.recv_timeout(timeout) {
            Ok(cmd) => Some(cmd),
            Err(RecvTimeoutError::Timeout) => None,
            Err(e) => {
                panic!("manager thread receiver disconnected. error: {}", e)
            }
        };

        if let Some(cmd) = cmd {
            let resp = do_command(&mut state, &cmd);

            rsptx
                .send(resp)
                .expect("manager thread failed to send response");

            if let Command::Shutdown = cmd {
                info!("manager thread received shutdown command. exiting");
                return;
            }
        } else {
            match state.upkeep() {
                Ok(info) => {
                    for pid in info.restarted {
                        let new_pid = state.status_of(pid).restarted.unwrap();
                        warn!(
                            "process {} failed and was restarted. new PID: {}",
                            pid, new_pid,
                        );
                    }
                }
                Err(e) => warn!("encountered error during upkeep: {}", e),
            }
        }

        if interrupt.load(atomic::Ordering::Relaxed) {
            return;
        }
    });

    ManagerThread {
        handle,
        command_sender: cmdtx,
        response_receiver: rsprx,
    }
}

/// begin listening and handling connections. this function handles all
/// received connections sequentially and blocks until the shutdown command
/// is received
fn listen(
    manager: &mut ManagerThread,
    socket_path: &Path,
    interrupt: Arc<AtomicBool>,
) -> Result<()> {
    let listener = match UnixListener::bind(socket_path) {
        Ok(l) => l,
        Err(e) => return Err(anyhow!(e)),
    };
    listener
        .set_nonblocking(true)
        .expect("failed to make listener nonblocking");

    info!(
        "sophtd v{} successfully started!",
        env!("CARGO_PKG_VERSION")
    );

    while !interrupt.load(atomic::Ordering::Relaxed) {
        let stream = listener.accept().map(|(s, _)| s);
        match stream {
            Ok(stream) => match Connection::new(stream) {
                Ok(conn) => match handle_connection(manager, conn) {
                    Ok(true) => {
                        info!("received shutdown command. quitting");
                        break;
                    }
                    Ok(_) => {}
                    Err(e) => {
                        warn!("connection terminated with error: {}", e)
                    }
                },
                Err(e) => {
                    warn!("failed to create connection. error: {}", e);
                }
            },
            Err(e) => {
                if e.kind() != std::io::ErrorKind::WouldBlock {
                    warn!("failed to connect to stream. error: {}", e);
                }
            }
        }
        std::thread::sleep(Duration::from_secs(1));
    }

    Ok(())
}

/// begin handling requests over a connection, blocking until the connection is
/// closed
fn handle_connection(
    manager: &mut ManagerThread,
    mut conn: Connection,
) -> Result<Quit> {
    loop {
        let cmd = match conn.recv_command() {
            Ok(cmd) => cmd,
            Err(Error::ConnectionClosed) => break,
            Err(e) => return Err(anyhow!(e)),
        };

        debug!("received command: {:?}", cmd);

        let shutdown = Command::Shutdown == cmd;

        manager.command_sender.send(cmd)?;
        let resp = manager.response_receiver.recv()?;

        debug!("generated response: {:?}", resp);

        match conn.send(&Message::Response(resp)) {
            Ok(_) => {}
            Err(Error::ConnectionClosed) => break,
            Err(e) => return Err(anyhow!(e)),
        }

        if shutdown {
            return Ok(true);
        }
    }

    Ok(false)
}

/// inspects a [`Command`], executes the command on the [`State`], then builds
/// a response and returns it
fn do_command(state: &mut State, cmd: &Command) -> Response {
    match cmd {
        Command::Hello => succeed("hello!"),
        Command::Shutdown => succeed("shutdown"),
        Command::Start(args) => match state.start(args) {
            Ok(p) => {
                info!("started \"{}\" with PID {}", args.prog, p.pid);
                succeed_status(
                    format!("started `{}` with PID {}", p.prog, p.pid),
                    &[p],
                )
            }
            Err(e) => {
                warn!("failed to start \"{}\". error: {}", args.prog, e);
                fail(format!("{}", e))
            }
        },
        Command::Output(pid) => match state.output(*pid) {
            Ok(data) => succeed(String::from_utf8_lossy(data)),
            Err(e) => fail(format!("{}", e)),
        },
        Command::Status(count) => Response {
            success: true,
            message: "".into(),
            status: Some(state.status(count)),
        },
        Command::Stop(pid) => match state.stop(*pid) {
            Ok(p) => {
                info!("stopped process {}", pid);
                succeed_status("successfully stopped process", &[p])
            }
            Err(e) => {
                warn!("failed to stop process {}", pid);
                fail(format!("{}", e))
            }
        },
        Command::Restart(pid) => match state.restart(*pid, false) {
            Ok(p) => {
                info!("restarted process {} with new PID {}", pid, p.pid);
                succeed_status(
                    format!(
                        "successfully restarted process. new PID: {}",
                        p.pid
                    ),
                    &[p],
                )
            }
            Err(e) => {
                warn!("failed to restart process {}. error: {}", pid, e);
                fail(format!("{}", e))
            }
        },
        Command::Send(args) => match state.send(args) {
            Ok(_) => succeed("successfully sent input to process"),
            Err(e) => fail(format!("could not send input. error: {}", e)),
        },
        Command::ChangeRestartPolicy(args) => {
            match state.change_restart_policy(args) {
                Ok(status) => {
                    succeed_status("altered restart policy", &[status])
                }
                Err(e) => {
                    fail(format!("could not change policy. error: {}", e))
                }
            }
        }
    }
}

/// helper function to generate a response with the success field set to true
fn succeed<S: Into<String>>(message: S) -> Response {
    Response {
        success: true,
        message: message.into(),
        status: None,
    }
}

/// helper function to generate a response with the success field set to false
fn fail<S: Into<String>>(message: S) -> Response {
    Response {
        success: false,
        message: message.into(),
        status: None,
    }
}

/// helper function that works like [`succeed`] but it takes a statuses slice
/// too and sets the status field of the response
fn succeed_status<S: Into<String>>(
    message: S,
    status: &[ProcessStatus],
) -> Response {
    Response {
        success: true,
        message: message.into(),
        status: Some(Vec::from(status)),
    }
}