siguldry 0.8.0

// SPDX-License-Identifier: MIT
// Copyright (c) Microsoft Corporation.

use std::path::PathBuf;

use anyhow::Context;
use clap::Parser;
use siguldry::{
    bridge::{self, Config},
    config::load_config,
};
use tokio::signal::unix::{SignalKind, signal};
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use tracing_subscriber::{EnvFilter, fmt::format::FmtSpan, layer::SubscriberExt};

// The name of the default config file location. Since this is expected to run under systemd,
// this is looked for under the provided CONFIGURATION_DIRECTORY environment variable.
const DEFAULT_CONFIG: &str = "bridge.toml";

/// The siguldry bridge.
///
/// The bridge acts as a proxy between the client and the server. Both client and server
/// connect to the bridge and authenticate with TLS certificates. Traffic is then forwarded
/// between the client and server.
#[derive(Debug, Parser)]
#[command(version)]
struct Cli {
    /// The path to the bridge's configuration file.
    ///
    /// If no path is provided, the defaults are used. To view the service configuration,
    /// run the `config` subcommand.
    #[arg(long, short, env = "SIGULDRY_BRIDGE_CONFIG")]
    config: Option<PathBuf>,

    /// Emit logs when new tracing spans are created, and when they are closed.
    ///
    /// This is useful in debugging scenarios to trace functions and tasks, but can lead to rather
    /// verbose logs.
    #[arg(long)]
    pub span_events: bool,

    /// A set of one or more comma-separated directives to filter logs.
    ///
    /// The general format is "target_name[span_name{field=value}]=level" where level is
    /// one of TRACE, DEBUG, INFO, WARN, ERROR.
    ///
    /// Details: https://docs.rs/tracing-subscriber/0.3.19/tracing_subscriber/filter/struct.EnvFilter.html#directives
    #[arg(
        long,
        env = "SIGULDRY_BRIDGE_LOG",
        default_value = "WARN,siguldry=INFO"
    )]
    pub log_filter: String,
    #[command(subcommand)]
    pub command: Command,
}

#[derive(clap::Subcommand, Debug)]
enum Command {
    /// Run the service.
    Listen {
        /// The directory containing the service's secrets.
        ///
        /// Any file referenced in the configuration that are not absolute paths are
        /// expected to be in this directory.
        ///
        /// When run under systemd, providing a `ImportCredential=`,
        /// `LoadCredentialEncrypted=`, or `LoadCredential=` directive will
        /// set the environment variable automatically for you.
        #[arg(long, env = "CREDENTIALS_DIRECTORY")]
        credentials_directory: PathBuf,
    },

    /// See the current bridge configuration.
    Config {
        /// The directory containing the service's secrets.
        ///
        /// Any file referenced in the configuration that are not absolute paths are
        /// expected to be in this directory.
        ///
        /// When run under systemd, providing a `ImportCredential=`,
        /// `LoadCredentialEncrypted=`, or `LoadCredential=` directive will
        /// set the environment variable automatically for you.
        #[arg(
            long,
            env = "CREDENTIALS_DIRECTORY",
            default_value = "/etc/credstore.encrypted/"
        )]
        credentials_directory: PathBuf,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let opts = Cli::parse();

    // Unfortunately we can't use clap's value_parser since EnvFilter does not
    // implement Clone.
    let log_filter = EnvFilter::builder().parse(&opts.log_filter).context(
        "SIGULDRY_BRIDGE_LOG contains an invalid log directive; refer to \
            https://docs.rs/tracing-subscriber/0.3.19/tracing_subscriber/\
            filter/struct.EnvFilter.html#directives for format details.",
    )?;

    let registry = tracing_subscriber::registry();
    let stderr_layer = tracing_subscriber::fmt::layer()
        .without_time()
        .with_writer(std::io::stderr);
    let registry = if opts.span_events {
        registry.with(stderr_layer.with_span_events(FmtSpan::NEW | FmtSpan::CLOSE))
    } else {
        registry.with(stderr_layer)
    };
    let registry = registry.with(log_filter);
    tracing::subscriber::set_global_default(registry)
        .expect("Programming error: set_global_default should only be called once.");

    siguldry::raise_nofiles()?;

    let mut config = load_config::<Config>(opts.config, PathBuf::from(DEFAULT_CONFIG).as_path())?;
    tracing::info!("Loaded configuration");

    match opts.command {
        Command::Listen {
            credentials_directory,
        } => {
            config
                .credentials
                .with_credentials_dir(&credentials_directory)?;

            let root_span = tracing::info_span!("bridge");
            async move {
                let listener = bridge::listen(config).await?;
                tokio::spawn(signal_handler(listener.halt_token()));

                listener.wait_to_finish().await?;
                Ok::<_, anyhow::Error>(())
            }
            .instrument(root_span)
            .await?;
        }
        Command::Config {
            credentials_directory,
        } => {
            println!(
                "# This is the current configuration\n\n{config}\n# This concludes the configuration.\n"
            );
            _ = config.credentials.with_credentials_dir(&credentials_directory).inspect_err(|error|{
                eprintln!("The configuration format is valid, but the referenced credentials aren't valid: {error:?}");
            });
        }
    }

    Ok(())
}

/// Install and manage signal handlers for the process.
///
/// # SIGTERM and SIGINT
///
/// Sending SIGTERM or SIGINT to the process will cause it to stop accepting new
/// signing requests. Existing signing requests will be allowed to complete
/// before the process shuts down.
async fn signal_handler(halt_token: CancellationToken) -> Result<(), anyhow::Error> {
    let mut sigterm_stream = signal(SignalKind::terminate()).inspect_err(|error| {
        tracing::error!(?error, "Failed to register a SIGTERM signal handler");
    })?;
    let mut sigint_stream = signal(SignalKind::interrupt()).inspect_err(|error| {
        tracing::error!(?error, "Failed to register a SIGINT signal handler");
    })?;

    loop {
        tokio::select! {
            _ = sigterm_stream.recv() => {
                tracing::info!("SIGTERM received, beginning service shutdown");
                halt_token.cancel();
            }
            _ = sigint_stream.recv() => {
                tracing::info!("SIGINT received, beginning service shutdown");
                halt_token.cancel();
            }
        }
    }
}