aralez/src/utils/tools.rs

use crate::tls::load;
use crate::tls::load::CertificateConfig;
use crate::utils::structs::{Extraparams, InnerMap, InnerMapForJson, UpstreamSnapshotForJson, UpstreamsDashMap, UpstreamsIdMap};
use dashmap::DashMap;
use log::{error, info};
use notify::{event::ModifyKind, Config, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
use privdrop::PrivDrop;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::any::type_name;
use std::collections::{HashMap, HashSet};
use std::fmt::Write;
use std::net::SocketAddr;
use std::net::TcpListener;
use std::os::unix::fs::MetadataExt;
use std::str::FromStr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Sender};
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::{fs, process, thread, time};

pub fn print_upstreams(upstreams: &UpstreamsDashMap, extraparams: &Extraparams) {
    let mut out = String::new();
    for host_entry in upstreams.iter() {
        writeln!(out, "Hostname: {}", host_entry.key()).unwrap();
        for path_entry in host_entry.value().iter() {
            writeln!(out, "    Path: {}", path_entry.key()).unwrap();
            for f in path_entry.value().0.clone() {
                writeln!(
                    out,
                    "        IP: {}, Port: {}, SSL: {}, H2: {}, To HTTPS: {}, Rate Limit: {}, 4xx Limit: {}",
                    f.address,
                    f.port,
                    f.is_ssl,
                    f.is_http2,
                    f.to_https,
                    f.rate_limit.unwrap_or(extraparams.rate_limit.unwrap_or(0)),
                    f.x4xx_limit.unwrap_or(extraparams.x4xx_limit.unwrap_or(0))
                )
                .unwrap();
            }
        }
    }
    info!("\n{}", out.trim_end());
}
#[allow(dead_code)]
pub fn typeoff<T>(_: T) {
    let to = type_name::<T>();
    println!("{:?}", to);
}

#[allow(dead_code)]
pub fn string_to_bool(val: Option<&str>) -> Option<bool> {
    match val {
        Some(v) => match v {
            "yes" => Some(true),
            "true" => Some(true),
            _ => Some(false),
        },
        None => Some(false),
    }
}

pub fn clone_dashmap(original: &UpstreamsDashMap) -> UpstreamsDashMap {
    let new_map: UpstreamsDashMap = DashMap::new();

    for outer_entry in original.iter() {
        let hostname = outer_entry.key();
        let inner_map = outer_entry.value();

        let new_inner_map = DashMap::new();

        for inner_entry in inner_map.iter() {
            let path = inner_entry.key();
            let (vec, _) = inner_entry.value();
            let new_vec = vec.clone();
            let new_counter = AtomicUsize::new(0);
            new_inner_map.insert(path.clone(), (new_vec, new_counter));
        }
        new_map.insert(hostname.clone(), new_inner_map);
    }
    new_map
}

pub fn clone_dashmap_into(original: &UpstreamsDashMap, cloned: &UpstreamsDashMap) {
    cloned.clear();
    for outer_entry in original.iter() {
        let hostname = outer_entry.key();
        let inner_map = outer_entry.value();
        let new_inner_map = DashMap::new();
        for inner_entry in inner_map.iter() {
            let path = inner_entry.key();
            let (vec, _) = inner_entry.value();
            let new_vec = vec.clone();
            let new_counter = AtomicUsize::new(0);
            new_inner_map.insert(path.clone(), (new_vec, new_counter));
        }
        cloned.insert(hostname.clone(), new_inner_map);
    }
}

pub fn compare_dashmaps(map1: &UpstreamsDashMap, map2: &UpstreamsDashMap) -> bool {
    if map1.len() != map2.len() {
        return false;
    }
    for entry1 in map1.iter() {
        let Some(inner_map2) = map2.get(entry1.key()) else {
            return false;
        };
        let inner_map1 = entry1.value();
        if inner_map1.len() != inner_map2.len() {
            return false;
        }
        for path_entry in inner_map1.iter() {
            let Some(entry2) = inner_map2.get(path_entry.key()) else {
                return false;
            };
            let (vec1, _) = path_entry.value();
            let (vec2, _) = entry2.value();
            if vec1.len() != vec2.len() {
                return false;
            }
            let set1: HashSet<_> = vec1.iter().collect();
            let set2: HashSet<_> = vec2.iter().collect();
            if set1 != set2 {
                return false;
            }
        }
    }
    true
}

pub fn merge_headers(target: &DashMap<Arc<str>, Vec<(String, Arc<str>)>>, source: &DashMap<Arc<str>, Vec<(String, Arc<str>)>>) {
    for entry in source.iter() {
        let global_key = entry.key().clone();
        let global_values = entry.value().clone();
        let mut target_entry = target.entry(global_key).or_default();
        target_entry.extend(global_values);
    }
}

pub fn clone_idmap_into(original: &UpstreamsDashMap, cloned: &UpstreamsIdMap) {
    cloned.clear();
    for outer_entry in original.iter() {
        let inner_map = outer_entry.value();
        let new_inner_map = DashMap::new();
        for inner_entry in inner_map.iter() {
            let path = inner_entry.key();
            let (vec, _) = inner_entry.value();
            let new_vec = vec.clone();
            for x in vec.iter() {
                let mut id = String::new();
                write!(
                    &mut id,
                    "{}:{}:{}:{}:{}:{}:{}:{}:{:?}",
                    outer_entry.key(),
                    x.address,
                    x.port,
                    x.is_http2,
                    x.to_https,
                    x.x4xx_limit.unwrap_or_default(),
                    x.rate_limit.unwrap_or_default(),
                    x.healthcheck.unwrap_or_default(),
                    x.authorization
                )
                .unwrap_or(());
                let mut hasher = Sha256::new();
                // address: "127.0.0.3", port: 8000, is_ssl: false, is_http2: false, to_https: false, rate_limit: Some(200), healthcheck: None, authorization: None } }

                hasher.update(id.clone().into_bytes());
                let hash = hasher.finalize();
                let hex_hash = base16ct::lower::encode_string(&hash);
                let hh = hex_hash[0..50].to_string();
                let to_add = InnerMap {
                    address: Arc::from("127.0.0.1"),
                    port: 0,
                    is_ssl: false,
                    is_http2: false,
                    to_https: false,
                    rate_limit: None,
                    x4xx_limit: None,
                    healthcheck: None,
                    redirect_to: None,
                    authorization: None,
                };
                cloned.insert(id, Arc::from(to_add));
                cloned.insert(hh, x.to_owned());
            }
            new_inner_map.insert(path.clone(), new_vec);
        }
    }
    info!("Upstreams are fully populated. Ready to server requests");
}

pub fn listdir(dir: String) -> Vec<load::CertificateConfig> {
    let mut f = HashMap::new();
    let mut certificate_configs: Vec<load::CertificateConfig> = vec![];
    let paths = fs::read_dir(dir).unwrap();
    for path in paths {
        let path_str = path.unwrap().path().to_str().unwrap().to_owned();
        if path_str.ends_with(".crt") {
            let name = path_str.replace(".crt", "");
            let mut inner = vec![];
            let domain = name.split("/").collect::<Vec<&str>>();
            inner.push(name.clone() + ".crt");
            inner.push(name.clone() + ".key");
            f.insert(domain[domain.len() - 1].to_owned(), inner);
            let y = CertificateConfig {
                cert_path: name.clone() + ".crt",
                key_path: name.clone() + ".key",
            };
            certificate_configs.push(y);
        }
    }
    // for (_, v) in f.iter() {
    //     let y = CertificateConfig {
    //         cert_path: v[0].clone(),
    //         key_path: v[1].clone(),
    //     };
    //     certificate_configs.push(y);
    // }
    certificate_configs
}

pub fn watch_folder(path: String, sender: Sender<Vec<CertificateConfig>>) -> notify::Result<()> {
    let (tx, rx) = channel();
    let mut watcher = RecommendedWatcher::new(tx, Config::default())?;
    watcher.watch(path.as_ref(), RecursiveMode::Recursive)?;
    info!("Watching for certificates in : {}", path);
    let certificate_configs = listdir(path.clone());
    sender.send(certificate_configs)?;
    let mut start = Instant::now();
    loop {
        match rx.recv_timeout(Duration::from_secs(1)) {
            Ok(Ok(event)) => match &event.kind {
                EventKind::Modify(ModifyKind::Data(_)) | EventKind::Create(_) | EventKind::Remove(_) => {
                    if start.elapsed() > Duration::from_secs(1) {
                        start = Instant::now();
                        let certificate_configs = listdir(path.clone());
                        sender.send(certificate_configs)?;
                        info!("Certificate changed: {:?}, {:?}", event.kind, event.paths);
                    }
                }
                _ => {}
            },
            Ok(Err(e)) => error!("Watch error: {:?}", e),
            Err(_) => {}
        }
    }
}

pub fn drop_priv(user: String, group: String, http_addr: String, tls_addr: Option<String>) {
    thread::sleep(time::Duration::from_millis(10));
    loop {
        thread::sleep(time::Duration::from_millis(10));
        if TcpListener::bind(&http_addr).is_err() {
            break;
        }
    }
    if let Some(tls_addr) = tls_addr {
        loop {
            thread::sleep(time::Duration::from_millis(10));
            if TcpListener::bind(&tls_addr).is_err() {
                break;
            }
        }
    }
    info!("Dropping ROOT privileges to: {}:{}", user, group);
    if let Err(e) = PrivDrop::default().user(user).group(group).apply() {
        error!("Failed to drop privileges: {}", e);
        process::exit(1)
    }
}

pub fn check_priv(addr: &str) {
    let port = SocketAddr::from_str(addr).map(|sa| sa.port()).expect("Failed to parse address port ");
    if port < 1024 {
        let meta = std::fs::metadata("/proc/self").map(|m| m.uid()).unwrap();
        if meta != 0 {
            error!("Running on privileged port requires to start as ROOT");
            process::exit(1)
        }
    }
}

#[allow(dead_code)]
pub fn upstreams_to_json(upstreams: &UpstreamsDashMap) -> serde_json::Result<String> {
    let mut outer = HashMap::new();

    for outer_entry in upstreams.iter() {
        let mut inner_map = HashMap::new();

        for inner_entry in outer_entry.value().iter() {
            let (backends, counter) = inner_entry.value();

            inner_map.insert(
                inner_entry.key().to_string(),
                UpstreamSnapshotForJson {
                    backends: backends
                        .iter()
                        .map(|a| InnerMapForJson {
                            address: a.address.to_string(),
                            port: a.port,
                            is_ssl: a.is_ssl,
                            is_http2: a.is_http2,
                            to_https: a.to_https,
                            rate_limit: a.rate_limit,
                            x4xx_limit: a.x4xx_limit,
                            healthcheck: a.healthcheck,
                        })
                        .collect(),
                    requests: counter.load(Ordering::Relaxed),
                },
            );
        }

        outer.insert(outer_entry.key().to_string(), inner_map);
    }

    // serde_json::to_string_pretty(&outer)
    serde_json::to_string(&outer)
}

pub fn upstreams_liveness_json(configured: &UpstreamsDashMap, current: &UpstreamsDashMap) -> Value {
    let mut result = serde_json::Map::new();

    for host_entry in configured.iter() {
        let hostname = host_entry.key().to_string();
        let configured_paths = host_entry.value();

        let mut paths_json = serde_json::Map::new();

        for path_entry in configured_paths.iter() {
            let path = path_entry.key().clone();
            let (configured_backends, _) = path_entry.value();
            let backends_json: Vec<Value> = configured_backends
                .iter()
                .map(|backend| {
                    let alive = if let Some(host_map) = current.get(&*hostname) {
                        if let Some(path_entry) = host_map.get(&*path) {
                            let list = &path_entry.value().0; // Vec<Arc<InnerMap>>
                            list.iter().any(|b| b.address == backend.address && b.port == backend.port)
                        } else {
                            false
                        }
                    } else {
                        false
                    };
                    json!({
                        "address": &*backend.address,
                        "port": backend.port,
                        "alive": alive
                    })
                })
                .collect();

            paths_json.insert(
                path.to_string(),
                json!({
                    "backends": backends_json
                }),
            );
        }

        result.insert(hostname, Value::Object(paths_json));
    }
    Value::Object(result)
}

#[allow(dead_code)]
pub fn prepend(prefix: &str, val: &Option<Arc<str>>, uri: &str, port: &str) -> Option<String> {
    val.as_ref().map(|s| {
        let mut buf = String::with_capacity(32);
        buf.push_str(prefix);
        buf.push_str(s);
        buf.push(':');
        buf.push_str(port);
        buf.push_str(uri);
        buf
    })
}