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 {
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>,
}
pub fn init_state(mut config: AppConfig) -> Result<AppState, std::io::Error> {
if !config.admin_key.is_empty() {
config.no_admin_key = false;
} else if !config.no_admin_key {
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()));
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,
};
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."
);
}
}