diff --git a/Cargo.toml b/Cargo.toml index c6a41b5..ee2814d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,7 +27,8 @@ async-trait = "0.1.92" log = "0.4.34" notify = "9.0.0-rc.4" axum = { version = "0.8.9" } -serde = { version = "1.0.229", features = ["derive"] } +serde = { version = "1.0.229", features = ["derive", "rc"] } +#serde = { version = "1.0.229", features = ["derive"] } noyalib = { version = "0.0.44", features = ["compat-serde-yaml"] } rand = "0.10.2" base64 = "0.23.1" diff --git a/src/core/proxyhttp.rs b/src/core/proxyhttp.rs index 256b5b5..a656490 100644 --- a/src/core/proxyhttp.rs +++ b/src/core/proxyhttp.rs @@ -34,6 +34,7 @@ pub struct LB { pub config: Arc, pub extraparams: Arc>, pub cache_enabled: bool, + pub master_server: bool, } pub struct Context { @@ -105,6 +106,10 @@ impl ProxyHttp for LB { } async fn request_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result { + if self.master_server { + let _ = session.respond_error(400).await; + return Ok(true); + }; ACTIVE_SESSIONS.inc(); let hostname = return_header_host_from_upstream(session, &self.ump_upst); _ctx.hostname = hostname; diff --git a/src/core/start.rs b/src/core/start.rs index 6b7ddca..6d552e6 100644 --- a/src/core/start.rs +++ b/src/core/start.rs @@ -34,6 +34,7 @@ pub fn run() { let file = parameters.conf.clone().unwrap(); let maincfg = crate::utils::parceyaml::parce_main_config(file.as_str()); + let master_server = maincfg.cluster_member.as_deref() == Some("master"); let cache_enabled = maincfg.cache_size_mb.is_some(); CACHE_TTL.set(maincfg.cache_ttl.unwrap_or(60)).unwrap(); if let Some(size) = maincfg.cache_size_mb { @@ -70,6 +71,7 @@ pub fn run() { server_headers: sh_config, extraparams: ec_config, cache_enabled, + master_server, }; let grade = cfg.proxy_tls_grade.clone().unwrap_or("medium".to_string()); diff --git a/src/core/webserver.rs b/src/core/webserver.rs index 0a6ad5a..03a0dc5 100644 --- a/src/core/webserver.rs +++ b/src/core/webserver.rs @@ -1,6 +1,7 @@ use crate::auth::jwt::Claims; use crate::ingress::APIUpstreamProvider; use crate::tls::acme::http01::{acme_create, acme_order, http01_challenge}; +use crate::utils::lazylock::NODE_STATUS; use crate::utils::metrics::{MEMORY_USAGE, OPEN_FILES, get_memory_usage, get_open_files}; use crate::utils::tools::{upstreams_liveness_json, upstreams_to_json}; use crate::utils::types::{Config, Configuration, UpstreamsDashMap}; @@ -63,6 +64,7 @@ pub async fn run_server(config: &APIUpstreamProvider, mut to_return: mpsc::Sende .route("/acme_order/{*domain}", any(acme_order)) .route("/.well-known/acme-challenge/{*token}", any(http01_challenge)) .route("/conf", post(conf)) + .route("/workerconf", post(workerconf)) .route("/metrics", get(metrics)) .route("/status", get(status)) .with_state(app_state); @@ -84,7 +86,8 @@ pub async fn run_server(config: &APIUpstreamProvider, mut to_return: mpsc::Sende let listener = port_is_available("Config API", &config.address).await; info!("Starting the API server on: {}", config.address); - let api_server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + // let api_server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let api_server = tokio::spawn(async move { axum::serve(listener, app.into_make_service_with_connect_info::()).await.unwrap() }); let (tx, mut rx) = mpsc::channel(1); std::thread::spawn(move || { @@ -124,6 +127,54 @@ async fn conf(State(st): State, Query(params): Query, ConnectInfo(addr): ConnectInfo, headers: HeaderMap, content: String) -> impl IntoResponse { + let api_key = headers.get("API-KEY").and_then(|v| v.to_str().ok()); + if api_key != st.master_key.as_deref() { + warn!("Unauthorised access to workerconf endpoint from {}", addr.ip()); + return Response::builder() + .status(StatusCode::UNAUTHORIZED) + .body(Body::from("Unauthorized: Invalid API-KEY\n")) + .unwrap(); + } + if let Some(member) = NODE_STATUS.get() { + if let Some(mbr) = member.as_ref() { + match mbr.as_ref() { + "worker" => { + info!("Accepted new config from master server {}", addr.ip()); + let strcontent = content.as_str(); + let parsed = noyalib::from_str::(strcontent); + match parsed { + Ok(_) => { + tokio::spawn(async move { apply_config(content.as_str(), st, true).await }); + return Response::builder().status(StatusCode::OK).body(Body::from("Accepted! Applying in background\n")).unwrap(); + } + Err(err) => { + error!("Failed to parse upstreams file: {}", err); + return Response::builder().status(StatusCode::BAD_GATEWAY).body(Body::from(format!("Failed: {}\n", err))).unwrap(); + } + } + } + _ => { + return Response::builder() + .status(StatusCode::BAD_GATEWAY) + .body(Body::from("Node is not running on worker mode !\n")) + .unwrap(); + } + } + } + return Response::builder() + .status(StatusCode::BAD_GATEWAY) + .body(Body::from("Node is not running on worker mode !\n")) + .unwrap(); + } + Response::builder() + .status(StatusCode::BAD_GATEWAY) + .body(Body::from("Node is not running on worker mode !\n")) + .unwrap() +} + async fn apply_config(content: &str, st: AppState, save: bool) { let sl = crate::utils::parceyaml::load_configuration(content, "content").await; if let Some(serverlist) = sl.0 { diff --git a/src/utils/deser.rs b/src/utils/deser.rs new file mode 100644 index 0000000..31d8cd5 --- /dev/null +++ b/src/utils/deser.rs @@ -0,0 +1,297 @@ +use crate::utils::hcclient::{HttpClientRequest, httpclient}; +use crate::utils::lazylock::{GLOBAL_PARAMS, MASTER_KEY, WORKERS}; +use crate::utils::types::{Consul, Extraparams, Headers, Kubernetes, UpstreamsDashMap}; +use bytes::Bytes; +use log::{error, info, warn}; +use noyalib::to_string; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::ops::Deref; + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct UpstreamConfigDto { + pub paths: HashMap, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct PathConfigDto { + pub is_http2: bool, + pub healthcheck: bool, + pub is_ssl: bool, + pub to_https: bool, + pub rate_limit: isize, + pub x4xx_limit: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub redirect_to: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub authorization: Option>, + pub servers: Vec, + // #[serde(skip_serializing_if = "Option::is_none")] + // pub client_headers: Option>, + // #[serde(skip_serializing_if = "Option::is_none")] + // pub server_headers: Option>, +} + +#[derive(Clone, Debug)] +pub struct SendConfigDto { + pub upstreams: HashMap, + pub client_headers: Headers, + pub server_headers: Headers, + pub consul: Option, + pub kubernetes: Option, + pub typecfg: String, + pub extraparams: Extraparams, +} + +pub fn serialize_config(data: &UpstreamsDashMap) -> String { + let config = GLOBAL_PARAMS.read().unwrap(); + let mut snapshot: HashMap = HashMap::new(); + + for outer_entry in data.iter() { + let domain = outer_entry.key().to_string(); + let mut paths: HashMap = HashMap::new(); + + for inner_entry in outer_entry.value().iter() { + let path = inner_entry.key().to_string(); + let (vec_inner, _atomic_val) = inner_entry.value(); + + if let Some(first) = vec_inner.first() { + let servers = vec_inner.iter().map(|item| format!("{}:{}", item.address, item.port)).collect(); + let mut auth = None; + if let Some(ath) = first.authorization.clone() { + let mut a = HashMap::new(); + a.insert("type".to_string(), ath.auth_type.clone().to_string()); + a.insert("data".to_string(), ath.auth_cred.clone().to_string()); + auth = Some(a); + } + paths.insert( + path, + PathConfigDto { + is_http2: first.is_http2, + healthcheck: first.healthcheck, + is_ssl: first.is_ssl, + to_https: first.to_https, + rate_limit: first.rate_limit, + x4xx_limit: first.x4xx_limit, + redirect_to: first.redirect_to.clone().as_ref().map(|s| s.to_string()), + authorization: auth, + servers, + }, + ); + } + } + + snapshot.insert(domain, UpstreamConfigDto { paths }); + } + let payload = SendConfigDto { + upstreams: snapshot, + client_headers: config.client_headers.clone(), + server_headers: config.server_headers.clone(), + consul: config.consul.clone(), + kubernetes: config.kubernetes.clone(), + typecfg: config.typecfg.clone(), + extraparams: config.extraparams.clone(), + }; + to_yaml(&payload) +} + +#[derive(Serialize, Debug)] +struct YamlAuth { + #[serde(rename = "type")] + auth_type: String, + data: String, +} + +#[derive(Serialize)] +struct YamlPath { + servers: Vec, + #[serde(skip_serializing_if = "is_false")] + is_http2: bool, + #[serde(skip_serializing_if = "is_false")] + to_https: bool, + #[serde(skip_serializing_if = "is_false")] + is_ssl: bool, + #[serde(skip_serializing_if = "is_true")] // healthcheck true is default + healthcheck: bool, + #[serde(skip_serializing_if = "is_zero_isize")] + rate_limit: isize, + #[serde(skip_serializing_if = "is_zero_u32")] + x4xx_limit: u32, + #[serde(skip_serializing_if = "Option::is_none")] + redirect_to: Option, + #[serde(skip_serializing_if = "Option::is_none")] + authorization: Option, + #[serde(skip_serializing_if = "Option::is_none")] + client_headers: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + server_headers: Option>, +} + +#[derive(Serialize)] +struct YamlUpstream { + paths: HashMap, +} + +#[derive(Serialize)] +struct YamlConfig { + provider: String, + sticky_sessions: u64, + to_https: bool, + rate_limit: isize, + x4xx_limit: u32, + #[serde(skip_serializing_if = "Vec::is_empty")] + server_headers: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] + client_headers: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + authorization: Option, + upstreams: HashMap, +} + +fn is_false(v: &bool) -> bool { + !v +} +fn is_true(v: &bool) -> bool { + *v +} +fn is_zero_isize(v: &isize) -> bool { + *v == 0 +} +fn is_zero_u32(v: &u32) -> bool { + *v == 0 +} + +pub fn to_yaml(config: &SendConfigDto) -> String { + let server_headers: Vec = config + .server_headers + .get("GLOBAL_SERVER_HEADERS") + .map(|paths| { + paths + .value() + .iter() + .flat_map(|inner| inner.value().iter().map(|(k, v)| format!("{}:{}", k, v)).collect::>()) + .collect() + }) + .unwrap_or_default(); + + let client_headers: Vec = config + .client_headers + .get("GLOBAL_CLIENT_HEADERS") + .map(|paths| { + paths + .value() + .iter() + .flat_map(|inner| inner.value().iter().map(|(k, v)| format!("{}:{}", k, v)).collect::>()) + .collect() + }) + .unwrap_or_default(); + + let authorization = config.extraparams.authentication.as_ref().map(|auth| YamlAuth { + auth_type: auth.auth_type.to_string(), + data: auth.auth_cred.to_string(), + }); + + let upstreams: HashMap = config + .upstreams + .iter() + .map(|(domain, upstream)| { + let paths = upstream + .paths + .iter() + .map(|(path, pcfg)| { + let ch: Option> = { + if let Some(domain_ref) = config.client_headers.get(domain.as_str()) { + if let Some(path_ref) = domain_ref.get(path.as_str()) { + let collected: Vec = path_ref.value().iter().map(|(k, v)| format!("{}:{}", k, v)).collect(); + if !collected.is_empty() { Some(collected) } else { None } + } else { + None + } + } else { + None + } + }; + + let sh: Option> = { + if let Some(domain_ref) = config.server_headers.get(domain.as_str()) { + if let Some(path_ref) = domain_ref.get(path.as_str()) { + let collected: Vec = path_ref.value().iter().map(|(k, v)| format!("{}:{}", k, v)).collect(); + if !collected.is_empty() { Some(collected) } else { None } + } else { + None + } + } else { + None + } + }; + + let auth = pcfg.authorization.as_ref().map(|a| YamlAuth { + auth_type: a.get("type").cloned().unwrap_or_default(), + data: a.get("data").cloned().unwrap_or_default(), + }); + + ( + path.clone(), + YamlPath { + servers: pcfg.servers.clone(), + is_http2: pcfg.is_http2, + is_ssl: pcfg.is_ssl, + healthcheck: pcfg.healthcheck, + rate_limit: pcfg.rate_limit, + x4xx_limit: pcfg.x4xx_limit, + to_https: pcfg.to_https, + redirect_to: pcfg.redirect_to.clone(), + authorization: auth, + client_headers: ch, + server_headers: sh, + }, + ) + }) + .collect(); + (domain.clone(), YamlUpstream { paths }) + }) + .collect(); + + let yaml_config = YamlConfig { + provider: config.typecfg.clone(), + sticky_sessions: config.extraparams.sticky_sessions, + to_https: config.extraparams.to_https, + rate_limit: config.extraparams.rate_limit, + x4xx_limit: config.extraparams.x4xx_limit, + server_headers, + client_headers, + authorization, + upstreams, + }; + + to_string(&yaml_config).unwrap_or_default() +} + +pub async fn send_to_worker(data: String) { + let masterkey = MASTER_KEY.get().expect("Master key not set"); + if let Some(workers) = WORKERS.get() { + for worker in workers { + let prefix = if worker.tls { "https://" } else { "http://" }; + let worker_url = format!("{}{}:{}", prefix, worker.address, worker.port); + let status = httpclient(HttpClientRequest { + method: "POST", + tls: worker.tls, + host: worker.address.as_ref(), + path: "/workerconf", + address: worker.address.as_ref(), + port: worker.port, + url: worker_url.as_str(), + payload: Bytes::from(data.clone()), + headers: Some(vec![("API-KEY", masterkey.clone().expect("Master key must be set in the cluster mode").deref())]), + }) + .await; + if (200..500).contains(&status.2) { + info!("Successfully updated Worker {}, status {}", worker_url, status.2); + } else { + error!("Cannot connect to worker node, status {}", status.2); + } + } + } else { + warn!("Workers list has not been initialized yet."); + } +} diff --git a/src/utils/hcclient.rs b/src/utils/hcclient.rs index 351fd50..ae6949b 100644 --- a/src/utils/hcclient.rs +++ b/src/utils/hcclient.rs @@ -18,9 +18,10 @@ pub struct HttpClientRequest<'a> { pub port: u16, pub url: &'a str, pub payload: Bytes, + pub headers: Option>, // ← add this } -pub async fn httpclient(req: HttpClientRequest<'_>) -> (bool, bool) { +pub async fn httpclient(req: HttpClientRequest<'_>) -> (bool, bool, u16) { let HttpClientRequest { method, tls, @@ -30,6 +31,7 @@ pub async fn httpclient(req: HttpClientRequest<'_>) -> (bool, bool) { port, url, payload, + headers, } = req; let method = match method { "HEAD" => "HEAD", @@ -57,7 +59,7 @@ pub async fn httpclient(req: HttpClientRequest<'_>) -> (bool, bool) { } else { log::warn!("Fail to connect to addr: {}, tls {}, host {}", address, tls, host); } - return (false, false); + return (false, false, 500); } }; @@ -68,7 +70,7 @@ pub async fn httpclient(req: HttpClientRequest<'_>) -> (bool, bool) { } None => { if ping_grpc(url).await { - return (true, true); + return (true, true, 204); } } } @@ -77,7 +79,7 @@ pub async fn httpclient(req: HttpClientRequest<'_>) -> (bool, bool) { Ok(r) => r, Err(e) => { log::warn!("Failed to build request: {}", e); - return (false, false); + return (false, false, 500); } }; @@ -85,113 +87,33 @@ pub async fn httpclient(req: HttpClientRequest<'_>) -> (bool, bool) { let body_bytes = if method == "POST" { req.insert_header("Content-Type", "application/json").ok(); req.insert_header("Content-Length", payload.len().to_string()).ok(); + if let Some(headers) = headers { + for (k, v) in headers { + req.insert_header(k.to_owned(), v).ok(); + } + } payload } else { Bytes::new() }; if let Err(e) = http_session.write_request_header(Box::new(req)).await { log::warn!("Write header failed: {}", e); - return (false, false); + return (false, false, 500); } if let Err(e) = http_session.write_request_body(body_bytes, true).await { log::warn!("Write body failed: {}", e); - return (false, false); + return (false, false, 500); } let status = match http_session.read_response_header().await { Ok(_) => http_session.response_header().map(|r| r.status.as_u16()).unwrap_or(500), Err(e) => { if !is_h2 && ping_grpc(url).await { - return (true, true); + return (true, true, 204); } log::warn!("Health Check read failed ({}) : {} - {}", if is_h2 { "H2" } else { "H1" }, host, e); - return (false, false); + return (false, false, 500); } }; HC_CONNECTOR.release_http_session(http_session, &peer, None).await; - ((200..500).contains(&status), is_h2) + ((200..500).contains(&status), is_h2, status) } -/* -pub async fn httpclient(method: &str, tls: bool, host: &str, path: &str, address: &str, port: u16, url: String) -> (bool, bool) { - let method = match method { - "HEAD" => "HEAD", - "GET" => "GET", - "POST" => "POST", - _ => "GET", - }; - let mut peer = HttpPeer::new((address, port), tls, host.to_string()); - - if tls { - peer.options.verify_cert = false; - peer.options.verify_hostname = false; - } - - let mut is_h2 = false; - peer.options.total_connection_timeout = Option::from(Duration::from_secs(5)); - peer.options.alpn = ALPN::H2H1; - let (mut http_session, _) = match HC_CONNECTOR.get_http_session(&peer).await { - Ok(s) => s, - Err(e) => { - if let Some(msg) = e.context { - log::warn!("{}, type: {}", msg.as_str(), e.etype.as_str()); - } else { - log::warn!("Fail to connect to addr: {}, tls {}, host {}", address, tls, host); - } - return (false, false); - } - }; - - match http_session.as_http2() { - Some(_) => { - peer.options.alpn = ALPN::H2; - is_h2 = true; - } - None => { - if ping_grpc(url.as_str()).await { - return (true, true); - } - } - } - - let mut req = match RequestHeader::build(method, path.as_bytes(), None) { - Ok(r) => r, - Err(e) => { - log::warn!("Failed to build request: {}", e); - return (false, false); - } - }; - - req.insert_header("Host", host).ok(); - - if let Err(e) = http_session.write_request_header(Box::new(req)).await { - log::warn!("Write header failed: {}", e); - return (false, false); - } - - let status = if is_h2 { - if let Err(e) = http_session.write_request_body(Bytes::new(), true).await { - log::warn!("Write body failed: {}", e); - return (false, false); - } - match http_session.read_response_header().await { - Ok(_) => http_session.response_header().map(|r| r.status.as_u16()).unwrap_or(500), - Err(e) => { - if ping_grpc(url.as_str()).await { - return (true, true); - } - log::warn!("Health Check H2 read failed: {} - {}", host, e); - return (false, false); - } - } - } else { - match http_session.read_response_header().await { - Ok(_) => http_session.response_header().map(|r| r.status.as_u16()).unwrap_or(500), - Err(e) => { - log::warn!("Health Check H1 read failed: {} - {}", e, host); - return (false, false); - } - } - }; - HC_CONNECTOR.release_http_session(http_session, &peer, None).await; - ((200..500).contains(&status), is_h2) -} -*/ diff --git a/src/utils/healthcheck.rs b/src/utils/healthcheck.rs index a94f5e8..ae2e75f 100644 --- a/src/utils/healthcheck.rs +++ b/src/utils/healthcheck.rs @@ -1,38 +1,54 @@ +use crate::utils::deser::{send_to_worker, serialize_config}; use crate::utils::hcclient::{HttpClientRequest, httpclient}; -use crate::utils::lazylock::REVERSE_STORE; +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, fullist: Arc, idlist: Arc, 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 = build_upstreams(&fullist, params.0).await; + 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 async fn initiate_upstreams(fullist: UpstreamsDashMap) -> UpstreamsDashMap { - build_upstreams(&fullist, "HEAD").await -} - -async fn build_upstreams(fullist: &UpstreamsDashMap, method: &str) -> UpstreamsDashMap { +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() { @@ -77,6 +93,7 @@ async fn build_upstreams(fullist: &UpstreamsDashMap, method: &str) -> UpstreamsD port: upstream.port, url: link.as_str(), payload: EMPTY_PAYLOAD, + headers: None, }) .await; if resp.0 { diff --git a/src/utils/lazylock.rs b/src/utils/lazylock.rs index fa2838f..eb275bd 100644 --- a/src/utils/lazylock.rs +++ b/src/utils/lazylock.rs @@ -1,3 +1,4 @@ +use crate::utils::types::{Configuration, WorkerNodes}; use dashmap::DashMap; use moka::sync::Cache; use pingora_cache::MemCache; @@ -20,3 +21,8 @@ pub static EVICTION: OnceLock = OnceLock::new(); pub static CACHE_TTL: OnceLock = OnceLock::new(); pub static CHALLENGES: LazyLock>> = LazyLock::new(|| RwLock::new(HashMap::new())); pub static DOMAINS: LazyLock> = LazyLock::new(DashMap::new); +pub static NODE_STATUS: OnceLock>> = OnceLock::new(); +pub static MASTER_KEY: OnceLock>> = OnceLock::new(); +pub static GLOBAL_PARAMS: LazyLock> = LazyLock::new(|| RwLock::new(Configuration::default())); +pub static IM_A_MASTER: OnceLock = OnceLock::new(); +pub static WORKERS: OnceLock> = OnceLock::new(); diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 65f072f..431ffa0 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -4,6 +4,7 @@ pub mod hcclient; pub mod healthcheck; pub mod httpclient; // pub mod ingress; +pub mod deser; pub mod lazylock; pub mod metrics; pub mod parceyaml; @@ -11,4 +12,5 @@ pub mod state; pub mod tools; pub mod types; pub(crate) mod watch; + pub use watch::*; diff --git a/src/utils/parceyaml.rs b/src/utils/parceyaml.rs index ab7aa90..6aa58a3 100644 --- a/src/utils/parceyaml.rs +++ b/src/utils/parceyaml.rs @@ -1,7 +1,8 @@ +use crate::utils::deser::{send_to_worker, serialize_config}; use crate::utils::healthcheck; -use crate::utils::lazylock::{DOMAINS, REVERSE_STORE}; +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::{clone_dashmap, clone_dashmap_into, print_upstreams}; +use crate::utils::tools::clone_dashmap_into; use crate::utils::types::*; use dashmap::DashMap; use log::{error, info, warn}; @@ -239,28 +240,48 @@ async fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) { } 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().expect("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 y = clone_dashmap(&imtdashmap); - let r = healthcheck::initiate_upstreams(y).await; + let r = healthcheck::initiate_upstreams(&imtdashmap, "HEAD").await; clone_dashmap_into(&r, &config.upstreams); } info!("Upstream Config:"); REVERSE_STORE.clear(); - print_upstreams(&config.upstreams, &config.extraparams); + + 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).expect("Failed to read main config file"); let mut cfo: AppConfig = noyalib::from_str(&data).expect("Failed to parse main config file"); - if let Ok(jwt_key) = env::var("JWT_KEY") { + 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(':') { @@ -279,6 +300,39 @@ pub fn parce_main_config(path: &str) -> AppConfig { } 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()))).expect("Config can only be initialized once"); + let mk = cfo.master_key.clone().expect("In CLuster mode master key must be set"); + MASTER_KEY.set(Some(Arc::from(mk.as_str()))).expect("Config can only be initialized once"); + if member == "master" { + IM_A_MASTER.set(true).expect("Master status can only be initialized once"); + let worker_str = cfo.worker_nodes.as_deref().expect("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::>(); + let wk = WorkerNodes { + address: Arc::from(wk_vec[1]), + port: wk_vec[2].parse::().expect("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::().expect("Invalid worker port number"); + // wks.push((host.to_string(), port)); + // } else { + // panic!("Invalid worker format, expected host:port"); + // } + } + WORKERS.set(wks).expect("Workers list can only be initialized once"); + } + } cfo } diff --git a/src/utils/types.rs b/src/utils/types.rs index 2f3ba58..7a7b2f5 100644 --- a/src/utils/types.rs +++ b/src/utils/types.rs @@ -9,7 +9,7 @@ pub type UpstreamsDashMap = DashMap, DashMap, (Vec>; pub type Headers = DashMap, DashMap, Vec<(String, Arc)>>>; -#[derive(Clone, Debug, Default)] +#[derive(Clone, Debug, Default, Serialize)] pub struct Extraparams { pub sticky_sessions: u64, pub rate_limit: isize, @@ -140,9 +140,11 @@ pub struct AppConfig { pub tcp_keepalive_idle: Option, pub tcp_keepalive_interval: Option, pub tcp_keepalive_count: Option, + pub cluster_member: Option, + pub worker_nodes: Option, } -#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)] +#[derive(Debug, Default, Clone, PartialEq, Eq, Hash, Serialize)] pub struct InnerAuth { pub auth_type: Arc, pub auth_cred: Arc, @@ -203,3 +205,11 @@ pub struct UpstreamSnapshotForJson { pub backends: Vec, pub requests: usize, } + +#[derive(Debug)] +pub struct WorkerNodes { + pub address: Arc, + pub key: Option>, + pub port: u16, + pub tls: bool, +}