akas 2.4.18

AKAS: API Key Authorization Server
use crate::models::AppConfig;
use chrono::Utc;
use std::{
    collections::HashSet,
    io::ErrorKind,
    sync::{Arc, RwLock},
};
use tracing::Level;

pub const EMPTY_STRING_SHA256: &str =
    "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";

#[derive(Debug)]
pub struct AppState {
    /// Hashed keys set from input file
    pub hash_key_set: Arc<RwLock<HashSet<String>>>,
    pub admin_key: String,
    pub no_admin_key: bool,
    pub local: bool,
    pub enable_metrics: bool,
    pub log_level: Level,
    pub original_length: u16,
    pub metadata_length: u16,
    pub file_hash: Arc<RwLock<String>>,
    pub file_date: Arc<RwLock<String>>,
    pub file_key_count: Arc<RwLock<u32>>,
    pub key_length: u16,
    pub key_prefix: String,
    pub auth_counter: Arc<prometheus::IntCounterVec>,
}

/// Initializes the application state.
///
/// Creates a new `AppState` instance with the provided parameters.
/// This includes setting up the hash key set, admin key, log level, and file metadata placeholders.
///
/// # Arguments
///
/// * `admin_key` - The admin key for administrative operations.
/// * `no_admin_key` - A boolean flag indicating whether an admin key is required.
/// * `port` - The port number the server will listen on.
/// * `log_level` - The logging level for the application.
/// * `length` - The length of the keys.
/// * `prefix` - The prefix of the keys.
///
/// # Returns
///
/// * `Ok(AppState)` - A `Result` containing the initialized `AppState` on success.
/// * `Err(std::io::Error)` - A `Result` containing an error if initialization fails.
///
pub fn init_state(mut config: AppConfig) -> Result<AppState, std::io::Error> {
    // If an admin_key is provided (either via CLI or ENV), it overrides --no-admin-key
    if !config.admin_key.is_empty() {
        config.no_admin_key = false;
    } else if !config.no_admin_key {
        // Only error if no_admin_key is explicitly false (or not provided) AND admin_key is empty
        return Err(std::io::Error::new(
            ErrorKind::InvalidInput,
            "No admin key provided. Please provide an admin key using the --admin-key flag or use the --no-admin-key flag.".to_string(),
        ));
    }

    if config.port < 1 {
        return Err(std::io::Error::new(
            ErrorKind::InvalidInput,
            format!(
                "{} is not a valid port number. Valid range is 1-65535.",
                config.port
            ),
        ));
    }
    let hash_key_set = Arc::new(RwLock::new(HashSet::new()));

    // Set AppState.log_level from clap log_level or env AKAS_LOG_LEVEL
    // default: Level::INFO
    let log_level = match config.log_level.as_str() {
        "error" => Level::ERROR,
        "warn" => Level::WARN,
        "info" => Level::INFO,
        "debug" => Level::DEBUG,
        "trace" => Level::TRACE,
        _ => Level::INFO,
    };

    // Hash of an empty string.
    let file_hash = Arc::new(RwLock::new(EMPTY_STRING_SHA256.to_string()));
    let now = Utc::now();
    let iso_now = now.to_rfc3339();
    let file_date = Arc::new(RwLock::new(iso_now));
    let file_key_count = Arc::new(RwLock::new(0));

    Ok(AppState {
        hash_key_set,
        admin_key: config.admin_key.to_string(),
        no_admin_key: config.no_admin_key,
        local: config.local,
        enable_metrics: config.enable_metrics,
        log_level,
        original_length: config.original_length,
        metadata_length: config.metadata_length,
        file_hash,
        file_date,
        file_key_count,
        key_length: config.key_length,
        key_prefix: config.key_prefix.to_string(),
        auth_counter: config.auth_counter,
    })
}

#[cfg(test)]
mod init_tests {
    use super::*;
    use prometheus::{IntCounterVec, Opts};

    #[test]
    fn test_init_state_with_admin_key_ok() {
        let auth_counter = Arc::new(
            IntCounterVec::new(Opts::new("test_auth", "Test auth counter"), &["status"]).unwrap(),
        );

        let result = init_state(AppConfig {
            admin_key: "my-admin-key".to_string(),
            no_admin_key: false,
            local: false,
            enable_metrics: false,
            port: 5001,
            log_level: "info".to_string(),
            original_length: 100,
            metadata_length: 0,
            key_length: 10,
            key_prefix: "test_".to_string(),
            auth_counter: auth_counter,
        });
        assert!(result.is_ok());

        let app_state = result.unwrap();
        assert_eq!(app_state.admin_key, "my-admin-key");
        assert_eq!(app_state.no_admin_key, false);
        assert_eq!(app_state.local, false);
        assert_eq!(app_state.enable_metrics, false);
        assert_eq!(app_state.log_level.as_str(), "INFO");
        assert_eq!(app_state.original_length, 100);
        assert_eq!(app_state.metadata_length, 0);
        assert_eq!(app_state.key_length, 10);
        assert_eq!(app_state.key_prefix, "test_");
    }

    #[test]
    fn test_init_state_with_no_admin_key() {
        let auth_counter = Arc::new(
            IntCounterVec::new(Opts::new("test_auth", "Test auth counter"), &["status"]).unwrap(),
        );

        let result = init_state(AppConfig {
            admin_key: "".to_string(),
            no_admin_key: true,
            local: false,
            enable_metrics: false,
            port: 5001,
            log_level: "info".to_string(),
            original_length: 15,
            metadata_length: 0,
            key_length: 10,
            key_prefix: "test_".to_string(),
            auth_counter: auth_counter,
        });
        assert!(result.is_ok());

        let app_state = result.unwrap();
        assert_eq!(app_state.admin_key, "");
        assert_eq!(app_state.no_admin_key, true);
        assert_eq!(app_state.local, false);
        assert_eq!(app_state.enable_metrics, false);
        assert_eq!(app_state.log_level.as_str(), "INFO");
        assert_eq!(app_state.original_length, 15);
        assert_eq!(app_state.metadata_length, 0);
        assert_eq!(app_state.key_length, 10);
        assert_eq!(app_state.key_prefix, "test_");
    }

    #[test]
    fn test_init_state_force_admin_key() {
        let auth_counter = Arc::new(
            IntCounterVec::new(Opts::new("test_auth", "Test auth counter"), &["status"]).unwrap(),
        );

        let result = init_state(AppConfig {
            admin_key: "my-admin-key".to_string(),
            no_admin_key: true,
            local: true,
            enable_metrics: true,
            port: 5001,
            log_level: "info".to_string(),
            original_length: 100,
            metadata_length: 0,
            key_length: 10,
            key_prefix: "test_".to_string(),
            auth_counter: auth_counter,
        });
        assert!(result.is_ok());

        let app_state = result.unwrap();
        assert_eq!(app_state.admin_key, "my-admin-key");
        assert_eq!(app_state.no_admin_key, false);
        assert_eq!(app_state.local, true);
        assert_eq!(app_state.enable_metrics, true);
    }

    #[test]
    fn test_init_state_error_admin_key() {
        let auth_counter = Arc::new(
            IntCounterVec::new(Opts::new("test_auth", "Test auth counter"), &["status"]).unwrap(),
        );

        let result = init_state(AppConfig {
            admin_key: "".to_string(),
            no_admin_key: false,
            local: false,
            enable_metrics: false,
            port: 0,
            log_level: "info".to_string(),
            original_length: 0,
            metadata_length: 100,
            key_length: 10,
            key_prefix: "test_".to_string(),
            auth_counter: auth_counter,
        });
        assert!(result.is_err());

        let error = result.unwrap_err();
        assert_eq!(
            error.to_string(),
            "No admin key provided. Please provide an admin key using the --admin-key flag or use the --no-admin-key flag."
        );
    }

    #[test]
    fn test_init_state_invalid_port() {
        let auth_counter = Arc::new(
            IntCounterVec::new(Opts::new("test_auth", "Test auth counter"), &["status"]).unwrap(),
        );

        let result = init_state(AppConfig {
            admin_key: "my-admin-key".to_string(),
            no_admin_key: false,
            local: false,
            enable_metrics: false,
            port: 0,
            log_level: "info".to_string(),
            original_length: 100,
            metadata_length: 0,
            key_length: 10,
            key_prefix: "test_".to_string(),
            auth_counter: auth_counter,
        });
        assert!(result.is_err());

        let error = result.unwrap_err();
        assert_eq!(
            error.to_string(),
            "0 is not a valid port number. Valid range is 1-65535."
        );
    }
}