ctf-ad-mirror-tool-rust/src/main.rs

use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use mirror_tool::{auth::Auth, config::Config, control::Controller, web};
use std::future::IntoFuture;
use std::path::PathBuf;
use tokio_util::sync::CancellationToken;

#[derive(Parser)]
#[command(version, about = "Async TCP mirror and ForcAD-aware upstream router")]
struct Cli {
    #[arg(short, long, default_value = "mirror.toml", global = true)]
    config: PathBuf,
    #[command(subcommand)]
    command: Option<Command>,
}
#[derive(Subcommand)]
enum Command {
    /// Validate TOML and any existing saved topology without binding or network I/O.
    Check,
    /// Print a portable sample configuration (no secrets).
    Example,
}

fn main() -> Result<()> {
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()?;
    let result = runtime.block_on(run());
    // A filesystem worker can remain blocked in the OS after the async deadline.
    // Do not make runtime destruction an unbounded second shutdown phase.
    runtime.shutdown_timeout(std::time::Duration::from_millis(100));
    result
}
async fn run() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "mirror_tool=info".into()),
        )
        .init();
    let cli = Cli::parse();
    if matches!(cli.command, Some(Command::Example)) {
        print!("{}", include_str!("../mirror.example.toml"));
        return Ok(());
    }
    let config = Config::load(&cli.config)?;
    if matches!(cli.command, Some(Command::Check)) {
        let topology = mirror_tool::store::Topology::load(&config)?;
        println!(
            "Configuration valid: {} teams, {} services",
            topology.teams.len(),
            topology.services.len()
        );
        return Ok(());
    }
    let auth = Auth::from_env(&config.web)?;
    let listen = config.web.listen;
    // Reserve control listener before starting data listeners, so startup cannot half-succeed.
    let listener = tokio::net::TcpListener::bind(listen)
        .await
        .context("bind web listener")?;
    let controller = Controller::start(config).await?;
    let cancel = CancellationToken::new();
    let monitor = tokio::spawn(controller.clone().monitor(cancel.clone()));
    let board = controller
        .config
        .scoreboard
        .clone()
        .map(|config| tokio::spawn(controller.board.clone().run(config, cancel.clone())));
    let router = web::router(controller.clone(), auth, cancel.clone());
    tracing::info!(%listen,"Mirror Tool ready");
    let server = axum::serve(
        listener,
        router.into_make_service_with_connect_info::<std::net::SocketAddr>(),
    );
    let signal_cancel = cancel.clone();
    let server = server
        .with_graceful_shutdown(signal_cancel.cancelled_owned())
        .into_future();
    tokio::pin!(server);
    let result = tokio::select! {
        result = &mut server => result,
        _ = shutdown_signal() => {
            cancel.cancel();
            // HTTP body extraction has no idle deadline during normal operation.
            // Shutdown alone bounds it, while serialized control transactions finish first.
            let shutdown = async {
                controller.shutdown().await;
                let _ = monitor.await;
                if let Some(board) = board { let _ = board.await; }
                (&mut server).await
            };
            match tokio::time::timeout(
                std::time::Duration::from_millis(controller.config.network.drain_timeout_ms)
                    + std::time::Duration::from_secs(1), shutdown).await {
                Ok(result) => result,
                Err(_) => { tracing::warn!("shutdown deadline reached"); return Ok(()); }
            }
        }
    };
    cancel.cancel();
    controller.shutdown().await;
    result.context("HTTP server")?;
    Ok(())
}
async fn shutdown_signal() {
    #[cfg(unix)]
    {
        let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("install SIGTERM handler");
        tokio::select! {_=tokio::signal::ctrl_c()=>{},_=term.recv()=>{}}
    }
    #[cfg(not(unix))]
    {
        let _ = tokio::signal::ctrl_c().await;
    }
}