aralez/src/utils/parceyaml.rs

use crate::utils::deser::{send_to_worker, serialize_config};
use crate::utils::healthcheck;
use crate::utils::lazylock::{DOMAINS, GLOBAL_PARAMS, IM_A_MASTER, MASTER_KEY, NODE_STATUS, REVERSE_STORE, WORKERS};
use crate::utils::state::{is_first_run, mark_not_first_run};
use crate::utils::tools::{Die, clone_dashmap_into};
use crate::utils::types::*;
use dashmap::DashMap;
use log::{error, info, warn};
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::{env, fs};

pub async fn load_configuration(d: &str, kind: &str) -> (Option<Configuration>, String) {
    let mut conf_files = Vec::new();
    let yaml_data = match kind {
        "filepath" => {
            let mut data = String::new();
            let mut last_error = None;
            for _ in 0..5 {
                match fs::read_to_string(d) {
                    Ok(content) => {
                        if !content.trim().is_empty() {
                            data = content;
                            break;
                        }
                    }
                    Err(e) => {
                        error!("Config read failed, retrying...");
                        last_error = Some(e);
                    }
                }
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
            }
            if data.is_empty() {
                let err_msg = match last_error {
                    Some(e) => {
                        error!("Reading: {}: {:?}", d, e);
                        e.to_string()
                    }
                    None => {
                        error!("Reading: {}: File is empty after retries", d);
                        "File is empty".to_string()
                    }
                };
                warn!("Running with empty upstreams list, update it via API");
                return (None, err_msg);
            }

            let mut confdir = Path::new(d).parent().die("Invalid configuration path").to_path_buf();
            let mut autocfg = Path::new(d).parent().die("Invalid configuration path").to_path_buf();

            autocfg.push("autoconfigs");
            if fs::metadata(autocfg.clone()).is_err() {
                fs::create_dir_all(autocfg.clone()).ok();
            }
            autocfg.push("domains.json");
            if autocfg.exists() {
                let json: Option<Vec<String>> = fs::read_to_string(autocfg).ok().and_then(|s| serde_json::from_str(&s).ok());
                if let Some(domains) = json {
                    for domain in domains {
                        DOMAINS.insert(domain, true);
                    }
                }
            }

            confdir.push("conf.d");

            if let Ok(entries) = fs::read_dir(&confdir) {
                let mut paths: Vec<_> = entries
                    .flatten()
                    .map(|e| e.path())
                    .filter(|p| p.extension().and_then(|e| e.to_str()) == Some("yaml"))
                    .collect();
                paths.sort();

                for path in paths {
                    let content = fs::read_to_string(&path);
                    match content {
                        Ok(content) => {
                            conf_files.push(content);
                        }
                        Err(e) => {
                            error!("Reading: {}: {:?}", path.display(), e)
                        }
                    };
                }
            }

            info!("Reading upstreams from {}", d);
            data
        }
        "content" => {
            info!("Reading upstreams from API post body");
            d.to_string()
        }
        _ => {
            error!("Mismatched parameter, only filepath|content is allowed");
            return (None, "Mismatched parameter, only filepath|content is allowed".to_string());
        }
    };

    let mut parsed: Config = match noyalib::from_str(&yaml_data) {
        Ok(cfg) => cfg,
        Err(e) => {
            error!("Failed to parse upstreams file: {}", e);
            return (None, e.to_string());
        }
    };

    if let Some(ref mut upstreams) = parsed.upstreams {
        for uconf in conf_files {
            let p: HashMap<String, HostConfig> = match noyalib::from_str(&uconf) {
                Ok(ucfg) => ucfg,
                Err(e) => {
                    error!("Failed to parse upstreams file: {}", e);
                    return (None, e.to_string());
                }
            };
            upstreams.extend(p);
        }
    }

    let mut toreturn = Configuration::default();
    populate_headers_and_auth(&mut toreturn, &parsed).await;
    toreturn.typecfg = parsed.provider.clone();

    match parsed.provider.as_str() {
        "file" => {
            populate_file_upstreams(&mut toreturn, &parsed).await;
            (Some(toreturn), "Ok".to_string())
        }
        "consul" => {
            toreturn.consul = parsed.consul;
            (toreturn.consul.is_some().then_some(toreturn), "Ok".to_string())
        }
        "kubernetes" => {
            toreturn.kubernetes = parsed.kubernetes;
            (toreturn.kubernetes.is_some().then_some(toreturn), "Ok".to_string())
        }
        _ => {
            warn!("Unknown provider {}", parsed.provider);
            (None, "Unknown provider".to_string())
        }
    }
}

async fn populate_headers_and_auth(config: &mut Configuration, parsed: &Config) {
    let mut ch: Vec<(String, Arc<str>)> = Vec::new();
    if let Some(headers) = &parsed.client_headers {
        for header in headers {
            if let Some((key, val)) = header.split_once(':') {
                ch.push((key.to_string(), Arc::from(val)));
            }
        }
    }

    let global_headers: DashMap<Arc<str>, Vec<(String, Arc<str>)>> = DashMap::new();
    global_headers.insert(Arc::from("/"), ch);
    config.client_headers.insert(Arc::from("GLOBAL_CLIENT_HEADERS"), global_headers);

    let mut sh: Vec<(String, Arc<str>)> = Vec::new();
    if let Some(headers) = &parsed.server_headers {
        for header in headers {
            if let Some((key, val)) = header.split_once(':') {
                sh.push((key.to_string(), Arc::from(val.trim())));
            }
        }
    }
    let server_global_headers: DashMap<Arc<str>, Vec<(String, Arc<str>)>> = DashMap::new();
    server_global_headers.insert(Arc::from("/"), sh);
    config.server_headers.insert(Arc::from("GLOBAL_SERVER_HEADERS"), server_global_headers);
    config.extraparams.to_https = parsed.to_https.unwrap_or(false);
    config.extraparams.sticky_sessions = parsed.sticky_sessions.unwrap_or(0);
    config.extraparams.rate_limit = parsed.rate_limit.unwrap_or(0);
    config.extraparams.x4xx_limit = parsed.x4xx_limit.unwrap_or(0);

    if let Some(rate) = &parsed.rate_limit {
        info!("Applied Global Rate Limit : {} request per second", rate);
    }

    if let Some(pa) = &parsed.authorization {
        let y: InnerAuth = InnerAuth {
            auth_type: Arc::from(pa.auth_type.clone()),
            auth_cred: Arc::from(pa.auth_cred.clone().unwrap_or_default()),
        };
        config.extraparams.authentication = Some(Arc::from(y));
    }
}

async fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
    let imtdashmap = UpstreamsDashMap::new();
    if let Some(upstreams) = &parsed.upstreams {
        for (hostname, host_config) in upstreams {
            let path_map = DashMap::new();
            let client_header_list = DashMap::new();
            let server_header_list = DashMap::new();
            for (path, path_config) in &host_config.paths {
                if let Some(rate) = &path_config.rate_limit {
                    info!("Applied Rate Limit for {} : {} request per second", hostname, rate);
                }
                let mut hl: Vec<(String, Arc<str>)> = Vec::new();
                let mut sl: Vec<(String, Arc<str>)> = Vec::new();
                build_headers(&path_config.client_headers, config, &mut hl);
                build_headers(&path_config.server_headers, config, &mut sl);
                client_header_list.insert(Arc::from(path.as_str()), hl);
                server_header_list.insert(Arc::from(path.as_str()), sl);
                let mut server_list = Vec::new();
                for server in &path_config.servers {
                    let mut path_auth: Option<Arc<InnerAuth>> = None;
                    if let Some(pa) = &path_config.authorization {
                        let y: InnerAuth = InnerAuth {
                            auth_type: Arc::from(pa.auth_type.clone()),
                            auth_cred: Arc::from(pa.auth_cred.clone().unwrap_or_default()),
                        };
                        path_auth = Some(Arc::from(y));
                    }

                    let redirect_link = path_config.redirect_to.as_ref().map(|www| Arc::from(www.as_str()));

                    if let Some((ip, port_str)) = server.rsplit_once(':') {
                        if let Ok(port) = port_str.parse::<u16>() {
                            server_list.push(Arc::from(InnerMap {
                                address: Arc::from(ip),
                                port,
                                is_ssl: path_config.is_ssl.unwrap_or(false),
                                is_http2: path_config.is_http2.unwrap_or(false),
                                to_https: path_config.to_https.unwrap_or(false),
                                rate_limit: path_config.rate_limit.unwrap_or(0),
                                x4xx_limit: path_config.x4xx_limit.unwrap_or(0),
                                healthcheck: path_config.healthcheck.unwrap_or(true),
                                redirect_to: redirect_link,
                                authorization: path_auth,
                            }));
                        }
                    }
                }
                path_map.insert(Arc::from(path.clone()), (server_list, AtomicUsize::new(0)));
            }
            config.client_headers.insert(Arc::from(hostname.clone()), client_header_list);
            config.server_headers.insert(Arc::from(hostname.clone()), server_header_list);

            imtdashmap.insert(Arc::from(hostname.clone()), path_map);
        }

        {
            let mut config_guard = GLOBAL_PARAMS.write().die("Global config lock poisoned");
            config_guard.client_headers = config.client_headers.clone();
            config_guard.server_headers = config.server_headers.clone();
            config_guard.extraparams = config.extraparams.clone();
            config_guard.kubernetes = config.kubernetes.clone();
            config_guard.consul = config.consul.clone();
            config_guard.typecfg = config.typecfg.clone();
        }

        if is_first_run() {
            clone_dashmap_into(&imtdashmap, &config.upstreams);
            mark_not_first_run();
        } else {
            let r = healthcheck::initiate_upstreams(&imtdashmap, "HEAD").await;
            clone_dashmap_into(&r, &config.upstreams);
        }
        info!("Upstream Config:");
        REVERSE_STORE.clear();

        if let Some(true) = IM_A_MASTER.get() {
            if !is_first_run() {
                let binary_output = serialize_config(&config.upstreams);
                let _ = send_to_worker(binary_output).await;
            }
        }
        // print_upstreams(&config.upstreams, &config.extraparams);
    }
}
pub fn parce_main_config(path: &str) -> AppConfig {
    let data = fs::read_to_string(path).die("Failed to read main config file");
    let mut cfo: AppConfig = noyalib::from_str(&data).die("Failed to parse main config file");
    if let Ok(jwt_key) = env::var("MASTER_KEY") {
        cfo.master_key = Some(jwt_key);
    } else if let Ok(jwt_key) = env::var("JWT_KEY") {
        cfo.master_key = Some(jwt_key);
    };

    // log_builder(&cfo, &cfo.log_file);
    cfo.hc_method = cfo.hc_method.to_uppercase();
    if let Some((ip, port_str)) = cfo.config_address.rsplit_once(':') {
        if let Ok(port) = port_str.parse::<u16>() {
            cfo.local_server = Option::from((ip.to_string(), port));
        }
    }
    if let Some(tlsport_cfg) = cfo.proxy_address_tls.clone() {
        if let Some((_, port_str)) = tlsport_cfg.rsplit_once(':') {
            cfo.proxy_port_tls = Some(port_str.to_string());
        }
    };

    if let Some((_, port_str)) = cfo.proxy_address_http.rsplit_once(':') {
        cfo.proxy_port = Some(port_str.to_string());
    }

    cfo.proxy_tls_grade = parce_tls_grades(cfo.proxy_tls_grade.clone());
    if let Some(member) = cfo.cluster_member.clone() {
        NODE_STATUS.set(Some(Arc::from(member.as_str()))).die("Config can only be initialized once");
        let mk = cfo.master_key.clone().die("In CLuster mode master key must be set");
        MASTER_KEY.set(Some(Arc::from(mk.as_str()))).die("Config can only be initialized once");
        if member == "master" {
            IM_A_MASTER.set(true).die("Master status can only be initialized once");
            let worker_str = cfo.worker_nodes.as_deref().die("At least one worker must be set in master mode");
            let mut wks = Vec::new();
            for worker in worker_str.split(',') {
                let clean = worker.replace("//", "");
                let wk_vec = clean.trim().split(':').collect::<Vec<&str>>();
                let wk = WorkerNodes {
                    address: Arc::from(wk_vec[1]),
                    port: wk_vec[2].parse::<u16>().die("Invalid worker port number"),
                    tls: match wk_vec[0] {
                        "https" => true,
                        "http" => false,
                        _ => panic!("invalid protocol string"),
                    },
                    key: cfo.master_key.clone().map(|s| Arc::from(s)),
                };
                wks.push(wk);
                // println!("{:?}", wk);
                // if let Some((host, port_str)) = worker.trim().split_once(':') {
                //     let port = port_str.parse::<u16>().die("Invalid worker port number");
                //     wks.push((host.to_string(), port));
                // } else {
                //     panic!("Invalid worker format, expected host:port");
                // }
            }
            WORKERS.set(wks).die("Workers list can only be initialized once");
        }
    }
    cfo
}

fn parce_tls_grades(what: Option<String>) -> Option<String> {
    match what {
        Some(g) => match g.to_ascii_lowercase().as_str() {
            "high" => {
                // info!("TLS grade set to: [ HIGH ]");
                Some("high".to_string())
            }
            "medium" => {
                // info!("TLS grade set to: [ MEDIUM ]");
                Some("medium".to_string())
            }
            "unsafe" => {
                // info!("TLS grade set to: [ UNSAFE ]");
                Some("unsafe".to_string())
            }
            _ => {
                warn!("Error parsing TLS grade, defaulting to: `medium`");
                Some("medium".to_string())
            }
        },
        None => {
            warn!("TLS grade not set, defaulting to: medium");
            Some("medium".to_string())
        }
    }
}

pub fn build_headers(path_config: &Option<Vec<String>>, _config: &Configuration, hl: &mut Vec<(String, Arc<str>)>) {
    if let Some(headers) = &path_config {
        for header in headers {
            if let Some((key, val)) = header.split_once(':') {
                hl.push((key.trim().to_string(), Arc::from(val.trim())));
            }
        }
    }
}