akas 2.4.18

AKAS: API Key Authorization Server
mod cli;
mod handlers;
mod metrics;
mod models;
mod state;
mod utils;

use crate::{
    cli::{load_client, Cli, Commands},
    models::AppConfig,
    state::{init_state, AppState},
};
use actix_web::{web, App, HttpServer};
use clap::Parser;
use tracing::info;
use tracing_subscriber::{filter::LevelFilter, fmt, prelude::*};

#[derive(Debug, Copy, Clone)]
pub enum BindAddress {
    Local,
    Any,
}

impl BindAddress {
    pub fn as_str(&self) -> &'static str {
        match self {
            BindAddress::Local => "127.0.0.1",
            BindAddress::Any => "0.0.0.0",
        }
    }
}

/// The main entry point of the application.
///
/// Initializes the application state by reading a configuration file, creates an HTTP server,
/// and starts listening on the specified port.
///
/// # Errors
///
/// Returns a `std::io::Error` if there's an error initializing the application state,
/// binding the server, or running the server.
///
#[actix_web::main]
async fn main() -> std::io::Result<()> {
    let args = Cli::parse();

    match args.command {
        Commands::Serve {
            admin_key,
            no_admin_key,
            local,
            enable_metrics,
            port,
            log_level,
            original_length,
            metadata_length,
            key_length,
            key_prefix,
        } => {
            let (prometheus, auth_counter_arc) = metrics::initialize_metrics(enable_metrics);

            let app_state: AppState = match init_state(AppConfig {
                admin_key: admin_key.to_string(),
                no_admin_key,
                local,
                enable_metrics,
                port,
                log_level,
                original_length,
                metadata_length,
                key_length,
                key_prefix: key_prefix.to_string(),
                auth_counter: auth_counter_arc,
            }) {
                Ok(a) => a,
                Err(e) => {
                    eprintln!("Error: {}", e);
                    std::process::exit(1);
                }
            };

            tracing_subscriber::registry()
                .with(
                    fmt::layer()
                        .json() // Enable JSON output
                        .with_ansi(false) // Disable ANSI escape codes for cleaner JSON
                        .with_target(true) // Include the target field (e.g., "akas")
                        .with_level(true) // Include the log level
                        .with_thread_ids(false) // Disable thread IDs
                        .with_current_span(false) // Disable current span information
                        .with_span_events(fmt::format::FmtSpan::NONE) // Log span start/end events
                        .flatten_event(false) // Flatten event fields directly into the top-level JSON object
                        .with_thread_names(false),
                )
                .with(LevelFilter::from_level(app_state.log_level)) // Apply the configured log level
                .init();

            info!("Log level: {}", app_state.log_level);
            if app_state.no_admin_key {
                info!("No admin key provided.");
            } else {
                info!("Admin key provided.");
            }
            let bind_address: BindAddress = if app_state.local {
                info!("Server bound to localhost.");
                BindAddress::Local
            } else {
                info!("Server bound to all interfaces.");
                BindAddress::Any
            };
            if app_state.enable_metrics {
                info!("Metrics endpoint enabled at /metrics");
            } else {
                info!("Metrics endpoint disabled");
            }
            info!(
                "Max x-original headers length: {}",
                app_state.original_length.to_string()
            );
            info!(
                "Max metadata header length: {}",
                app_state.metadata_length.to_string()
            );
            if app_state.key_length > 0 {
                info!("Key length: {}", app_state.key_length);
            }
            if !app_state.key_prefix.is_empty() {
                info!("Key prefix: {}", app_state.key_prefix);
            }
            info!(
                "Starting server on http://{}:{} ...",
                bind_address.as_str(),
                port
            );

            let data = web::Data::new(app_state);

            if let Some(prometheus_middleware) = prometheus {
                HttpServer::new(move || {
                    App::new()
                        .app_data(data.clone())
                        .service(handlers::auth)
                        .service(handlers::auth_unauthorized)
                        .service(handlers::health)
                        .service(handlers::load)
                        .service(handlers::status)
                        .wrap(prometheus_middleware.clone())
                })
                .bind((bind_address.as_str(), port))?
                .run()
                .await
            } else {
                HttpServer::new(move || {
                    App::new()
                        .app_data(data.clone())
                        .service(handlers::auth)
                        .service(handlers::auth_unauthorized)
                        .service(handlers::health)
                        .service(handlers::load)
                        .service(handlers::status)
                })
                .bind((bind_address.as_str(), port))?
                .run()
                .await
            }
        }
        Commands::Load { host, file, format } => match load_client(&host, &file, &format) {
            Ok(m) => {
                println!("{}", m);
                Ok(())
            }
            Err(e) => {
                eprintln!("Error: {}", e);
                std::process::exit(1);
            }
        },
    }
}