skubelb/src/main.rs

use std::sync::Mutex;
use std::process::Command;

use clap::Parser;

use skubelb::Rewriter;
use skubelb::Server;

use env_logger::Env;
use log::{info, warn};
use anyhow::Result;
use rouille::{router, Request, Response};

/// Implements a HTTP server which allows clients to 'register'
/// themselves. Their IP address will be used to replace a
/// needle in a set of config files. This is intended to be
/// used as a low-cost way of enabling Kubernetes ingress
/// using nginx running on a machine that has a public port.
/// 
/// The needle is expected to be a dummy IP address; something
/// fairly unique. The goal is to replace nginx files, where
/// we often repeat lines if we want nginx to load balance between
/// multiple destination IPs.
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// The needle that will be replaced. Anytime a line
    /// is encountered with this needle, the line is dropped
    /// and instead N lines (one per replacement) is added to
    /// the output.
    #[arg(short, long)]
    needle: String,

    /// The folder which contains the templates that
    /// will be be searched for the needle.
    #[arg(short, long)]
    template_dir: String,

    /// The symlink that should be updated each time the config changes.
    /// 
    /// Symlinks are used because file updates are not atomic.
    #[arg(short, long)]
    config_symlink: String,

    /// Where to actually store the generated configs.
    #[arg(short, long)]
    workspace_dir: String,

    /// Address to listen for http requests on.
    #[arg(short, long, default_value_t = String::from("0.0.0.0:8080"))]
    listen: String,

    /// Command to reload nginx.
    #[arg(short, long, default_value_t = String::from("sudo nginx -s reload"))]
    reload_cmd: String,
}

fn main() {
    // Log info and above by default
    env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
    let args = Args::parse();

    let rewriter = Rewriter::new(args.needle);
    let server_impl = Mutex::new(Server::new(rewriter, args.workspace_dir, args.template_dir, args.config_symlink));
    let reload_command = args.reload_cmd.leak();
    let reload_command: Vec<&str> = reload_command.split_ascii_whitespace().collect();

    rouille::start_server(args.listen, move |request| {
        info!("Processing request: {:?}", request);
        match handle(request, &server_impl) {
            Ok((resp, reload)) => {
                if reload && reload_command.len() > 0 {
                    let output = Command::new(reload_command[0])
                        .args(&reload_command[1..])
                        .output();
                    match output {
                        Ok(o) => {
                            info!("Ran {:?}; exit code: {}", reload_command, o.status);
                            info!("Ran {:?}; stdout: {}", reload_command, String::from_utf8_lossy(&o.stdout));
                            info!("Ran {:?}; stderr: {}", reload_command, String::from_utf8_lossy(&o.stderr));
                        },
                        Err(e) => {
                            warn!("Failed to run {:?}: {:?}", reload_command, e);
                        }
                    };
                }
                resp
            },
            Err(e) => {
                warn!("{:?}", e);
                Response{status_code: 500, ..Response::empty_400()}
            }
        }
    });
}

fn handle(request: &Request, server_impl: &Mutex<Server>) -> Result<(Response, bool)> {
    router!(request,
        (POST) (/register/{ip: String}) => {
            server_impl.lock().unwrap().register(request, &ip)?;
            Ok((Response{status_code: 200, ..Response::empty_204()}, true))
        },
        (DELETE) (/register/{ip: String}) => {
            server_impl.lock().unwrap().unregister(request, &ip)?;
            Ok((Response{status_code: 200, ..Response::empty_204()}, true))
        },
        _ => Ok((Response::empty_404(), false)),
    )
}