aralez/src/utils/healthcheck.rs

use crate::utils::deser::{send_to_worker, serialize_config};
use crate::utils::hcclient::{HttpClientRequest, httpclient};
use crate::utils::lazylock::{IM_A_MASTER, NODE_STATUS, REVERSE_STORE};
use crate::utils::tools::*;
use crate::utils::types::{InnerMap, UpstreamsDashMap, UpstreamsIdMap};
use bytes::Bytes;
use dashmap::DashMap;
use log::info;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::time::{interval, timeout};
use tokio_native_tls::native_tls::TlsConnector;

const EMPTY_PAYLOAD: Bytes = Bytes::from_static(b"");

pub async fn hc2(upslist: Arc<UpstreamsDashMap>, fullist: Arc<UpstreamsDashMap>, idlist: Arc<UpstreamsIdMap>, params: (&str, u64)) {
    let mut period = interval(Duration::from_secs(params.1));
    if let Some(member) = NODE_STATUS.get() {
        if let Some(mbr) = member.as_ref() {
            match mbr.as_ref() {
                "worker" => {
                    info!("Disabling health-checks. Running in cluster mode as {}", mbr);
                    return;
                }
                "master" => {
                    info!("Disabling client/server routing. Running in cluster mode as {}", mbr);
                }
                _ => {}
            }
        }
    }
    loop {
        tokio::select! {
            _ = period.tick() => {
                let totest = initiate_upstreams(&fullist, params.0).await;
                if !compare_dashmaps(&totest, &upslist) {
                    clone_dashmap_into(&totest, &upslist);
                    clone_idmap_into(&totest, &idlist);
                    if let Some(true) = IM_A_MASTER.get(){
                        let binary_output = serialize_config(&totest);
                        let _ = send_to_worker(binary_output).await;
                    }
                    REVERSE_STORE.clear();
                }
            }
        }
    }
}
pub(crate) async fn initiate_upstreams(fullist: &UpstreamsDashMap, method: &str) -> UpstreamsDashMap {
    let totest: UpstreamsDashMap = DashMap::new();
    let fclone = clone_dashmap(fullist);
    for val in fclone.iter() {
        let host = val.key();
        let inner = DashMap::new();

        for path_entry in val.value().iter() {
            let path = path_entry.key();
            let mut innervec = Vec::new();

            for upstream in path_entry.value().0.iter() {
                let tls = if upstream.healthcheck {
                    detect_tls(upstream.address.as_ref(), &upstream.port, host).await
                } else {
                    upstream.is_ssl
                };
                let mut scheme = InnerMap {
                    address: upstream.address.clone(),
                    port: upstream.port,
                    is_ssl: tls,
                    is_http2: upstream.is_http2,
                    to_https: upstream.to_https,
                    rate_limit: upstream.rate_limit,
                    x4xx_limit: upstream.x4xx_limit,
                    healthcheck: upstream.healthcheck,
                    redirect_to: upstream.redirect_to.clone(),
                    authorization: upstream.authorization.clone(),
                };

                if scheme.healthcheck {
                    let link = if tls {
                        format!("https://{}:{}{}", upstream.address, upstream.port, path)
                    } else {
                        format!("http://{}:{}{}", upstream.address, upstream.port, path)
                    };
                    let resp = httpclient(HttpClientRequest {
                        method,
                        tls,
                        host,
                        path,
                        address: upstream.address.as_ref(),
                        port: upstream.port,
                        url: link.as_str(),
                        payload: EMPTY_PAYLOAD,
                        headers: None,
                    })
                    .await;
                    if resp.0 {
                        if resp.1 {
                            scheme.is_http2 = resp.1;
                        }
                        innervec.push(Arc::from(scheme));
                    }
                } else {
                    innervec.push(Arc::from(scheme));
                }
            }
            inner.insert(path.clone(), (innervec, AtomicUsize::new(0)));
        }
        totest.insert(host.clone(), inner);
    }
    totest
}

async fn detect_tls(ip: &str, port: &u16, sni: &str) -> bool {
    let addr = format!("{ip}:{port}");
    let timeout_duration = Duration::from_secs(2); // THINK !!!
    let stream = match timeout(timeout_duration, TcpStream::connect(&addr)).await {
        Ok(Ok(s)) => s,
        Ok(Err(_)) => return false,
        Err(_) => return false,
    };
    let connector = tokio_native_tls::TlsConnector::from(
        TlsConnector::builder()
            .danger_accept_invalid_certs(true)
            .danger_accept_invalid_hostnames(true)
            .build()
            .die("Failed to create TLS connector"),
    );
    connector.connect(sni, stream).await.is_ok()
}