roughenough-server 2.2.0

//!
//! # Server Operating Model
//!
//! - **Multi-threaded**: Spawns one worker thread per CPU core. All threads bind to the same UDP
//!   port using `SO_REUSEPORT`, letting the kernel load-balance incoming packets across threads.
//!
//! - **Event-driven I/O**: Each thread runs an independent mio event loop (poll-based) to handle
//!   non-blocking UDP I/O. No async/await - just a simple synchronous event loop.
//!
//! - **Zero shared state**: Each thread is completely independent with its own online key,
//!   socket, and processing pipeline. No synchronization or sharing between threads.
//!
//! - **Request Batching**: Each thread collects up to 64 requests, builds a Merkle tree, signs
//!   the root once, then sends responses to each client. This amortizes the (relatively)
//!   expensive Ed25519 signature calculation across all requests in a batch.
//!
//! - **Processing Pipeline**: UDP Socket -> NetworkHandler -> RequestHandler -> BatchingResponder
//!   -> UDP Socket
//!
use std::io;
use std::net::{SocketAddr, UdpSocket as StdUdpSocket};
use std::path::Path;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering::{Acquire, Release};
use std::sync::mpsc::{SyncSender, channel, sync_channel};
use std::thread::JoinHandle;
use std::time::Duration;

use clap::Parser;
use mio::net::UdpSocket as MioUdpSocket;
use roughenough_keys::seed::{Seed, SeedBackend, try_choose_backend};
use roughenough_keys::storage::try_load_seed_sync;
use roughenough_protocol::util::ClockSource;
use roughenough_server::args::Args;
use roughenough_server::keysource::KeySource;
use roughenough_server::metrics::aggregator::{MetricsAggregator, WorkerMetrics};
use roughenough_server::metrics::snapshot::validate_metrics_directory;
use roughenough_server::responses::ResponseHandler;
use roughenough_server::seed_file;
use roughenough_server::worker::{ExitGuard, Worker};
use signal_hook::iterator::Signals;
use socket2::{Domain, Socket, Type};
use tracing::{debug, error, info, warn};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::{Layer, filter};

/// Global flag that will be set to `false` when all threads should exit.
static KEEP_RUNNING: AtomicBool = AtomicBool::new(true);

#[derive(Clone, Copy)]
struct WorkerConfig {
    socket_addr: SocketAddr,
    batch_size: usize,
    metrics_interval: Duration,
}

fn main() {
    set_signal_handler();

    let args = Args::parse();
    enable_logging(&args);
    debug!(
        listen_addr = %args.udp_socket_addr(),
        num_threads = args.num_threads,
        batch_size = args.batch_size,
        rotation_interval_hours = args.rotation_interval,
        metrics_interval_seconds = args.metrics_interval,
        seed_backend = %args.seed_backend,
        "Server configuration loaded"
    );

    let clock = choose_clock(&args);
    let seed = load_seed(&args);
    let key_source = KeySource::new(seed, clock.clone(), args.rotation_interval());
    let worker_config = WorkerConfig {
        socket_addr: args.udp_socket_addr(),
        batch_size: args.batch_size as usize,
        metrics_interval: Duration::from_secs(args.metrics_interval),
    };

    info!("Long term public key: {:?}", key_source.public_key());

    let (metrics_thread, metrics_chan_tx) =
        start_metrics_thread(&args, clock.clone(), &KEEP_RUNNING);

    let mut threads = Vec::new();
    threads.push(metrics_thread);

    let (exit_tx, exit_rx) = channel();

    for i in 0..args.num_threads {
        let clock = clock.clone();
        let ks = key_source.clone();
        let metrics_chan_tx = metrics_chan_tx.clone();
        let exit_tx = exit_tx.clone();

        let thread = std::thread::Builder::new()
            .name(format!("worker-{i}"))
            .spawn(move || {
                let _exit_guard = ExitGuard::new(i as usize, exit_tx);
                worker_task(i, ks, worker_config, clock, metrics_chan_tx)
            })
            .unwrap();

        threads.push(thread);
    }
    drop(exit_tx);

    // An unexpected worker exit probably means a panic. To be safe,
    // shutdown and let the server's parent restart it.
    let mut worker_died = false;
    for _ in 0..args.num_threads {
        match exit_rx.recv() {
            Ok(worker_id) => {
                if KEEP_RUNNING.load(Acquire) {
                    error!("worker-{worker_id} exited unexpectedly; shutting down");
                    KEEP_RUNNING.store(false, Release);
                    worker_died = true;
                }
            }
            Err(_) => break,
        }
    }

    for thread in threads {
        if thread.join().is_err() {
            worker_died = true;
        }
    }

    if worker_died {
        error!("Server exiting after unexpected worker death");
        std::process::exit(1);
    }

    info!("Server finished");
}

fn load_seed(args: &Args) -> Box<dyn SeedBackend> {
    let seed = match &args.seed_file {
        None => {
            warn!("Using all-zero seed because --insecure-zero-seed was specified");
            Seed::new(&[0u8; 32])
        }
        Some(path) => {
            let contents = match seed_file::read(path) {
                Ok(contents) => contents,
                Err(e) => {
                    error!("reading seed file '{}': {e}", path.display());
                    std::process::exit(1);
                }
            };
            let encoded_value = match contents.encoded_value() {
                Ok(value) => value,
                Err(e) => {
                    error!("reading seed file '{}': {e}", path.display());
                    std::process::exit(1);
                }
            };

            match try_load_seed_sync(encoded_value) {
                Ok(seed) => seed,
                Err(e) => {
                    error!("loading seed from '{}': {e}", path.display());
                    std::process::exit(1);
                }
            }
        }
    };

    let mut backend = match try_choose_backend(&args.seed_backend.to_string()) {
        Ok(backend) => backend,
        Err(e) => {
            error!("choosing seed backend: {e}");
            std::process::exit(1);
        }
    };

    info!(
        "Loaded {}-byte seed into '{}' backend",
        seed.len(),
        &args.seed_backend
    );

    if let Err(e) = backend.store_seed(seed) {
        error!("storing seed in '{}' backend: {e}", &args.seed_backend);
        std::process::exit(1);
    }
    backend
}

fn choose_clock(args: &Args) -> ClockSource {
    match args.fixed_offset {
        0 => ClockSource::System,
        offset => ClockSource::FixedOffset(offset),
    }
}

fn worker_task(
    idx: u16,
    key_source: KeySource,
    config: WorkerConfig,
    clock: ClockSource,
    metrics_channel: SyncSender<WorkerMetrics>,
) {
    let sock = bind_socket(config.socket_addr).expect("Failed to bind socket");
    let responder = ResponseHandler::new(config.batch_size as u8, key_source);
    let idx = idx as usize;

    let mut worker = Worker::new(
        idx,
        config.batch_size,
        responder,
        clock,
        metrics_channel,
        config.metrics_interval,
    );

    worker.run(sock, &KEEP_RUNNING);
}

// Bind to the server port using SO_REUSEPORT so the kernel will fairly
// balance traffic to each worker. https://lwn.net/Articles/542629/
fn bind_socket(sock_addr: SocketAddr) -> io::Result<MioUdpSocket> {
    let sock_domain = Domain::for_address(sock_addr);
    let socket = Socket::new(sock_domain, Type::DGRAM, None)?;
    socket.set_nonblocking(true)?;
    socket.set_reuse_port(true)?;

    #[cfg(target_os = "linux")]
    if let Err(e) = roughenough_server::filter::attach(&socket) {
        warn!("BPF request filter not attached, continuing without it: {e}");
    }

    socket.bind(&sock_addr.into())?;

    let std_socket: StdUdpSocket = socket.into();
    let mio_socket: MioUdpSocket = MioUdpSocket::from_std(std_socket);
    Ok(mio_socket)
}

fn set_signal_handler() {
    let mut signals = Signals::new([signal_hook::consts::SIGINT, signal_hook::consts::SIGTERM])
        .expect("Error setting signal handler");

    // detached on purpose: Signals::forever() never returns
    std::thread::Builder::new()
        .name("signal-handler".to_string())
        .spawn(move || {
            if signals.forever().next().is_some() {
                info!("Received shutdown signal, exiting...");
                KEEP_RUNNING.store(false, Release);
            }
        })
        .expect("Failed to spawn signal handler thread");
}

fn enable_logging(args: &Args) {
    // AWS, GCP, Rustls, Hyper, etc crates are quite verbose, "normal" level for them is WARN
    let cloud_sdk_verbosity = match args.verbose {
        _ if args.quiet => tracing::Level::ERROR,
        0 => tracing::Level::WARN,
        1 => tracing::Level::INFO,
        2 => tracing::Level::DEBUG,
        3.. => tracing::Level::TRACE,
    };

    let verbosity = match args.verbose {
        _ if args.quiet => tracing::Level::ERROR,
        0 => tracing::Level::INFO,
        1 => tracing::Level::DEBUG,
        2.. => tracing::Level::TRACE,
    };

    let filters = filter::Targets::new()
        .with_target("aws_config", cloud_sdk_verbosity)
        .with_target("aws_sdk_kms", cloud_sdk_verbosity)
        .with_target("aws_sdk_secretsmanager", cloud_sdk_verbosity)
        .with_target("google_cloud_kms_v1", cloud_sdk_verbosity)
        .with_target("google_cloud_secretmanager_v1", cloud_sdk_verbosity)
        .with_target("hyper_util", cloud_sdk_verbosity)
        .with_target("rustls", cloud_sdk_verbosity)
        .with_default(verbosity); // for all other targets

    let fmt_layer = tracing_subscriber::fmt::layer()
        .compact()
        .with_filter(filters);

    tracing_subscriber::registry().with(fmt_layer).init();
}

/// Creates metrics collection infrastructure and spawns collector thread
pub fn start_metrics_thread(
    args: &Args,
    clock: ClockSource,
    keep_running: &'static AtomicBool,
) -> (JoinHandle<()>, SyncSender<WorkerMetrics>) {
    let num_workers = args.num_threads as usize;
    let metrics_interval = Duration::from_secs(args.metrics_interval);

    let metrics_path = args.metrics_output.as_ref().map(|metrics_path_arg| {
        let path = Path::new(metrics_path_arg);
        if let Err(e) = validate_metrics_directory(path) {
            error!("metrics path: {e}");
            std::process::exit(1);
        }
        info!("Metrics will be written to: {}", path.display());
        path.to_path_buf()
    });

    // Each worker pushes at most once per interval, so num_workers * 2 should be plenty
    let (sender, receiver) = sync_channel(num_workers * 2);

    let collector = MetricsAggregator::new(
        receiver,
        num_workers,
        metrics_interval,
        keep_running,
        clock,
        metrics_path,
    );

    let thread = std::thread::Builder::new()
        .name("metrics-collector".to_string())
        .spawn(move || collector.run())
        .expect("Failed to spawn metrics collector thread");

    (thread, sender)
}