ipx/crates/directory/src/main.rs

//! `ipx-directory`: serve the directory, make a token, or check that a running one answers.

use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use ipx_directory::{Directory, config, store::Store};
use tracing_subscriber::prelude::*;

#[derive(Parser)]
#[command(version, about = "ipx's directory of public podcasts")]
struct Cli {
    /// The config file: where to listen, the database, and the tokens it accepts. Otherwise
    /// DIRECTORY_CONFIG, or else /config/directory.toml.
    #[arg(long)]
    config: Option<std::path::PathBuf>,
    #[command(subcommand)]
    cmd: Option<Cmd>,
}

#[derive(Subcommand)]
enum Cmd {
    /// Serve the directory (the default).
    Serve,
    /// Make a token, and print it with the lines that let it in.
    Token {
        #[command(subcommand)]
        cmd: TokenCmd,
    },
    /// Check that the directory on this machine answers and can read its database: the
    /// container's healthcheck.
    Health,
}

#[derive(Subcommand)]
enum TokenCmd {
    New {
        /// Who it is for, as the log and their snapshot will name them.
        name: String,
        #[arg(long, value_parser = ["ro", "rw", "rwa"])]
        access: String,
    },
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    let config_path = cli
        .config
        .or_else(|| std::env::var_os("DIRECTORY_CONFIG").map(Into::into))
        .unwrap_or_else(|| "/config/directory.toml".into());
    match cli.cmd.unwrap_or(Cmd::Serve) {
        Cmd::Serve => serve(&config_path).await,
        Cmd::Token { cmd: TokenCmd::New { name, access } } => {
            let token = config::new_token();
            println!("A token for {name}, to give them; it is not kept anywhere:\n\n    {token}\n");
            println!("To let it in, add this to {}:\n", config_path.display());
            println!("[[tokens]]\nname = {name:?}\naccess = {access:?}\ntoken_sha256 = {:?}", config::hash(&token));
            println!("\nThe directory reads the file again within a minute. Delete the entry to revoke it.");
            Ok(())
        }
        Cmd::Health => {
            let cfg = config::Config::load(&config_path)?;
            let port = cfg.bind.rsplit(':').next().unwrap_or("8100");
            let resp = reqwest::Client::new()
                .get(format!("http://127.0.0.1:{port}/healthz"))
                .timeout(std::time::Duration::from_secs(4))
                .send()
                .await
                .context("the directory does not answer")?;
            anyhow::ensure!(resp.status().is_success(), "the directory answers {}", resp.status());
            Ok(())
        }
    }
}

async fn serve(path: &std::path::Path) -> Result<()> {
    // One JSON object a line for Loki when DIRECTORY_LOG_FORMAT=json, as ipx does; text otherwise.
    let json = std::env::var("DIRECTORY_LOG_FORMAT").is_ok_and(|f| f.eq_ignore_ascii_case("json"));
    let filter = || tracing_subscriber::EnvFilter::try_from_env("DIRECTORY_LOG").unwrap_or_else(|_| "ipx_directory=info".into());
    // Traces to Tempo when OTEL_EXPORTER_OTLP_ENDPOINT is set, joining ipx's (#213).
    let otel = ipx_directory::otel::provider("ipx-directory")?;
    tracing_subscriber::registry()
        .with((!json).then(|| tracing_subscriber::fmt::layer().with_writer(std::io::stderr).with_filter(filter())))
        .with(json.then(|| tracing_subscriber::fmt::layer().json().flatten_event(true).with_writer(std::io::stderr).with_filter(filter())))
        .with(otel.as_ref().map(|p| {
            use opentelemetry::trace::TracerProvider;
            tracing_opentelemetry::layer().with_tracer(p.tracer("ipx-directory")).with_filter(tracing_subscriber::EnvFilter::new("ipx_directory=info"))
        }))
        .init();

    let cfg = config::Config::load(path)?;
    let tokens = cfg.tokens()?;
    let url = std::env::var(&cfg.database_url_env)
        .ok()
        .filter(|u| !u.trim().is_empty())
        .with_context(|| format!("{} is not set: the directory needs a database URL", cfg.database_url_env))?;
    let store = Store::open(&url).await.context("opening the database")?;
    let dir = Directory::new(store, vec![])?;
    let dropped = dir.set_tokens(tokens).await?;

    let listener = tokio::net::TcpListener::bind(&cfg.bind).await.with_context(|| format!("listening on {}", cfg.bind))?;
    tracing::info!(version = env!("CARGO_PKG_VERSION"), bind = %cfg.bind, tokens = cfg.tokens.len(), dropped, "directory started");
    tokio::spawn(ipx_directory::refresher(dir.clone()));
    tokio::spawn(ipx_directory::reloader(dir.clone(), path.to_owned()));
    axum::serve(listener, ipx_directory::router(dir))
        .with_graceful_shutdown(async {
            let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()).expect("SIGTERM handler");
            tokio::select! {
                _ = term.recv() => {}
                _ = tokio::signal::ctrl_c() => {}
            }
        })
        .await?;
    // The batch exporter holds the last few seconds of spans; without this they are lost.
    if let Some(p) = otel {
        let _ = p.shutdown();
    }
    Ok(())
}