domani/src/main.rs

use clap::Parser;
use futures::stream::StreamExt;
use signal_hook_tokio::Signals;

use std::net::SocketAddr;
use std::str::FromStr;
use std::{path, sync};

#[derive(Parser, Debug)]
#[command(version)]
#[command(about = "A domiply to another dimension")]
struct Cli {
    #[arg(
        long,
        help = "OFF, ERROR, WARN, INFO, DEBUG, or TRACE",
        default_value_t = log::LevelFilter::Info,
        env = "DOMIPLY_LOG_LEVEL"
    )]
    log_level: log::LevelFilter,

    #[arg(long, default_value_t = false, env = "DOMIPLY_LOG_TIMESTAMP")]
    log_timestamp: bool,

    #[arg(long, required = true, env = "DOMIPLY_HTTP_DOMAIN")]
    http_domain: domiply::domain::Name,

    #[arg(long, default_value_t = SocketAddr::from_str("[::]:3030").unwrap(), env = "DOMIPLY_HTTP_LISTEN_ADDR")]
    http_listen_addr: SocketAddr,

    #[arg(
        long,
        help = "E.g. '[::]:443', if given then SSL certs will automatically be retrieved for all domains using LetsEncrypt",
        env = "DOMIPLY_HTTPS_LISTEN_ADDR",
        requires = "domain_acme_contact_email",
        requires = "domain_acme_store_dir_path"
    )]
    https_listen_addr: Option<SocketAddr>,

    #[arg(long, required = true, env = "DOMIPLY_PASSPHRASE")]
    passphrase: String,

    #[arg(long, required = true, env = "DOMIPLY_ORIGIN_STORE_GIT_DIR_PATH")]
    origin_store_git_dir_path: path::PathBuf,

    #[arg(long, required = true, env = "DOMIPLY_DOMAIN_CHECKER_TARGET_A")]
    domain_checker_target_a: std::net::Ipv4Addr,

    #[arg(long, default_value_t = String::from("1.1.1.1:53"), env = "DOMIPLY_DOMAIN_CHECKER_RESOLVER_ADDR")]
    domain_checker_resolver_addr: String,

    #[arg(long, required = true, env = "DOMIPLY_DOMAIN_CONFIG_STORE_DIR_PATH")]
    domain_config_store_dir_path: path::PathBuf,

    #[arg(long, env = "DOMIPLY_DOMAIN_ACME_STORE_DIR_PATH")]
    domain_acme_store_dir_path: Option<path::PathBuf>,

    #[arg(long, env = "DOMIPLY_DOMAIN_ACME_CONTACT_EMAIL")]
    domain_acme_contact_email: Option<String>,
}

#[derive(Clone)]
struct HTTPSParams {
    https_listen_addr: SocketAddr,
    domain_acme_store: sync::Arc<dyn domiply::domain::acme::store::Store>,
    domain_acme_manager: sync::Arc<dyn domiply::domain::acme::manager::Manager>,
}

#[tokio::main]
async fn main() {
    let config = Cli::parse();

    env_logger::Builder::new()
        .filter_level(config.log_level)
        .format_timestamp(
            config
                .log_timestamp
                .then_some(env_logger::TimestampPrecision::Micros),
        )
        .init();

    let canceller = tokio_util::sync::CancellationToken::new();

    {
        let canceller = canceller.clone();

        tokio::spawn(async move {
            let mut signals = Signals::new(signal_hook::consts::TERM_SIGNALS)
                .expect("initializing signals failed");

            if (signals.next().await).is_some() {
                log::info!("Gracefully shutting down...");
                canceller.cancel();
            }

            if (signals.next().await).is_some() {
                log::warn!("Forcefully shutting down");
                std::process::exit(1);
            };
        });
    }

    let origin_store = domiply::origin::store::git::new(config.origin_store_git_dir_path)
        .expect("git origin store initialization failed");

    let domain_checker = domiply::domain::checker::new(
        config.domain_checker_target_a,
        &config.domain_checker_resolver_addr,
    )
    .await
    .expect("domain checker initialization failed");

    let domain_config_store = domiply::domain::config::new(&config.domain_config_store_dir_path)
        .expect("domain config store initialization failed");

    let https_params = if let Some(https_listen_addr) = config.https_listen_addr {
        let domain_acme_store_dir_path = config.domain_acme_store_dir_path.unwrap();

        let domain_acme_store = domiply::domain::acme::store::new(&domain_acme_store_dir_path)
            .expect("domain acme store initialization failed");

        // if https_listen_addr is set then domain_acme_contact_email is required, see the Cli/clap
        // settings.
        let domain_acme_contact_email = config.domain_acme_contact_email.unwrap();

        let domain_acme_manager = domiply::domain::acme::manager::new(
            domain_acme_store.clone(),
            &domain_acme_contact_email,
        )
        .await
        .expect("domain acme manager initialization failed");

        Some(HTTPSParams {
            https_listen_addr,
            domain_acme_store,
            domain_acme_manager,
        })
    } else {
        None
    };

    let domain_manager = domiply::domain::manager::new(
        origin_store,
        domain_config_store,
        domain_checker,
        https_params.as_ref().map(|p| p.domain_acme_manager.clone()),
    );

    let domain_manager = sync::Arc::new(domain_manager);

    {
        let http_service = domiply::service::http::new(
            domain_manager.clone(),
            config.domain_checker_target_a,
            config.passphrase,
            config.http_listen_addr.clone(),
            config.http_domain.clone(),
            https_params.map(|p| domiply::service::http::HTTPSParams {
                listen_addr: p.https_listen_addr,
                cert_resolver: domiply::domain::acme::resolver::new(p.domain_acme_store),
            }),
        );

        canceller.cancelled().await;

        sync::Arc::into_inner(http_service).unwrap().stop().await;
    }

    sync::Arc::into_inner(domain_manager)
        .unwrap()
        .stop()
        .await
        .expect("domain manager failed to shutdown cleanly");

    log::info!("Graceful shutdown complete");
}