Github/aralez
1
0
Fork 0
mirror of https://github.com/sadoyan/aralez.git synced 2026-10-23 06:01:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Rewrote and optimized healtchecks

Browse Source
This commit is contained in:
Ara Sadoyan
2026-08-05 18:51:07 +02:00
parent 600cb602e3
commit b63d822577
10 changed files with 214 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/tls/acme/order.rs
+1 -4
View File
@@ -1,17 +1,14 @@
use crate::tls::acme::account::get_account;
use crate::utils::lazylock::CHALLENGES;
use crate::utils::parceyaml::DOMAINS;
use instant_acme::{ChallengeType, Identifier, NewOrder, RetryPolicy};
use log::{error, info};
use pingora::prelude::sleep;
use rcgen::{CertificateParams, DistinguishedName, KeyPair};
use std::collections::HashMap;
use std::fs;
use std::sync::{LazyLock, RwLock};
use std::time::Duration;
use x509_parser::prelude::*;
pub static CHALLENGES: LazyLock<RwLock<HashMap<String, String>>> = LazyLock::new(|| RwLock::new(HashMap::new()));
pub async fn refresh_order(certs_dir: String, autoconf_dir: String) {
let credsfile = autoconf_dir + "/acme_credentials.json";
loop {
src/utils.rs
+1
View File
@@ -3,6 +3,7 @@ pub mod discovery;
mod filewatch;
pub mod fordebug;
pub mod grpc;
pub mod hcclient;
pub mod healthcheck;
pub mod httpclient;
pub mod jwt;
src/utils/auth.rs
+20 -19
View File
@@ -1,5 +1,6 @@
use crate::utils::jwt::{check_jwt, JWT_TOKEN};
use crate::utils::structs::InnerAuth;
pub(crate) use crate::utils::tools::split_host_port;
use axum::http::StatusCode;
use base64::engine::general_purpose::STANDARD;
use base64::Engine;
@@ -224,22 +225,22 @@ pub fn get_query_param(session: &mut Session, key: &str) -> Option<String> {
params.get(key).and_then(|v| decode(v).ok()).map(|s| s.to_string())
}
#[allow(clippy::needless_return)]
fn split_host_port(addr: &str, tls: bool) -> Option<(&str, u16, bool, &str)> {
match addr.split_once(':') {
Some((h, p)) => match p.parse::<u16>() {
Ok(port) => return Some((h, port, tls, h)),
Err(_) => {
log::warn!("ForwardAuth: invalid port in {}", addr);
return None;
}
},
None => {
if tls {
return Some((addr, 443u16, tls, addr));
} else {
return Some((addr, 80u16, tls, addr));
}
}
};
}
// #[allow(clippy::needless_return)]
// pub fn split_host_port(addr: &str, tls: bool) -> Option<(&str, u16, bool, &str)> {
// match addr.split_once(':') {
// Some((h, p)) => match p.parse::<u16>() {
// Ok(port) => return Some((h, port, tls, h)),
// Err(_) => {
// log::warn!("ForwardAuth: invalid port in {}", addr);
// return None;
// }
// },
// None => {
// if tls {
// return Some((addr, 443u16, tls, addr));
// } else {
// return Some((addr, 80u16, tls, addr));
// }
// }
// };
// }
src/utils/hcclient.rs
+93
View File
@@ -0,0 +1,93 @@
use crate::utils::grpc::ping_grpc;
use bytes::Bytes;
use pingora_core::connectors::http::Connector;
use pingora_core::listeners::ALPN;
use pingora_core::prelude::HttpPeer;
use pingora_http::RequestHeader;
use std::sync::LazyLock;
pub static HC_CONNECTOR: LazyLock<Connector> = LazyLock::new(|| Connector::new(None));
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.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)
}
src/utils/healthcheck.rs
+16 -61
View File
@@ -1,10 +1,8 @@
use crate::utils::grpc::ping_grpc;
use crate::utils::hcclient::httpclient;
use crate::utils::lazylock::REVERSE_STORE;
use crate::utils::structs::{InnerMap, UpstreamsDashMap, UpstreamsIdMap};
use crate::utils::tools::*;
use dashmap::DashMap;
use log::{error, warn};
use reqwest::{Client, Version};
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use std::time::Duration;
@@ -14,12 +12,10 @@ use tokio_native_tls::native_tls::TlsConnector;
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));
let client = Client::builder().timeout(Duration::from_secs(params.1)).danger_accept_invalid_certs(true).build().unwrap();
loop {
tokio::select! {
_ = period.tick() => {
// populate_upstreams(&upslist, &fullist, &idlist, params, &client).await;
let totest = build_upstreams(&fullist, params.0, &client).await;
let totest = build_upstreams(&fullist, params.0).await;
if !compare_dashmaps(&totest, &upslist) {
clone_dashmap_into(&totest, &upslist);
clone_idmap_into(&totest, &idlist);
@@ -31,11 +27,10 @@ pub async fn hc2(upslist: Arc<UpstreamsDashMap>, fullist: Arc<UpstreamsDashMap>,
}
pub async fn initiate_upstreams(fullist: UpstreamsDashMap) -> UpstreamsDashMap {
let client = Client::builder().timeout(Duration::from_secs(2)).danger_accept_invalid_certs(true).build().unwrap();
build_upstreams(&fullist, "HEAD", &client).await
build_upstreams(&fullist, "HEAD").await
}
async fn build_upstreams(fullist: &UpstreamsDashMap, method: &str, client: &Client) -> UpstreamsDashMap {
async fn build_upstreams(fullist: &UpstreamsDashMap, method: &str) -> UpstreamsDashMap {
let totest: UpstreamsDashMap = DashMap::new();
let fclone = clone_dashmap(fullist);
for val in fclone.iter() {
@@ -48,15 +43,15 @@ async fn build_upstreams(fullist: &UpstreamsDashMap, method: &str, client: &Clie
for upstream in path_entry.value().0.iter() {
let tls = if upstream.healthcheck.unwrap_or(true) {
detect_tls(upstream.address.as_ref(), &upstream.port, client).await
detect_tls(upstream.address.as_ref(), &upstream.port, host).await
} else {
(false, None)
false
};
let mut scheme = InnerMap {
address: upstream.address.clone(),
port: upstream.port,
is_ssl: tls.0,
is_http2: matches!(tls.1, Some(Version::HTTP_2)),
is_ssl: tls,
is_http2: false,
to_https: upstream.to_https,
rate_limit: upstream.rate_limit,
x4xx_limit: upstream.x4xx_limit,
@@ -66,21 +61,17 @@ async fn build_upstreams(fullist: &UpstreamsDashMap, method: &str, client: &Clie
};
if scheme.healthcheck.unwrap_or(true) {
let link = if tls.0 {
let link = if tls {
format!("https://{}:{}{}", upstream.address, upstream.port, path)
} else {
format!("http://{}:{}{}", upstream.address, upstream.port, path)
};
let resp = http_request(&link, method, "", client).await;
let resp = httpclient(method, tls, host, path, upstream.address.as_ref(), upstream.port, link).await;
if resp.0 {
if resp.1 {
scheme.is_http2 = resp.1;
}
innervec.push(Arc::from(scheme));
} else {
warn!("Dead Upstream : {}", link);
}
} else {
innervec.push(Arc::from(scheme));
@@ -93,33 +84,12 @@ async fn build_upstreams(fullist: &UpstreamsDashMap, method: &str, client: &Clie
totest
}
async fn http_request(url: &str, method: &str, payload: &str, client: &Client) -> (bool, bool) {
if !["POST", "GET", "HEAD"].contains(&method) {
error!("Method {} not supported. Only GET|POST|HEAD are supported ", method);
return (false, false);
}
async fn send_request(client: &Client, method: &str, url: &str, payload: &str) -> Option<reqwest::Response> {
match method {
"POST" => client.post(url).body(payload.to_owned()).send().await.ok(),
"GET" => client.get(url).send().await.ok(),
"HEAD" => client.head(url).send().await.ok(),
_ => None,
}
}
async fn detect_tls(ip: &str, port: &u16, sni: &str) -> bool {
let addr = format!("{ip}:{port}");
match send_request(client, method, url, payload).await {
Some(response) => {
let status = response.status().as_u16();
((99..499).contains(&status), false)
}
None => (ping_grpc(url).await, true),
}
}
async fn detect_tls(ip: &str, port: &u16, client: &Client) -> (bool, Option<Version>) {
let addr = format!("{}:{}", ip, port);
let Ok(stream) = TcpStream::connect(&addr).await else {
return (false, Some(Version::HTTP_11));
let stream = match TcpStream::connect(&addr).await {
Ok(s) => s,
Err(_) => return false,
};
let connector = tokio_native_tls::TlsConnector::from(
@@ -129,20 +99,5 @@ async fn detect_tls(ip: &str, port: &u16, client: &Client) -> (bool, Option<Vers
.build()
.unwrap(),
);
let tls = connector.connect(ip, stream).await.is_ok();
if tls {
let vers = detect_version(ip, *port, true, client).await;
return (tls, vers);
}
(tls, Some(Version::HTTP_11))
}
async fn detect_version(ip: &str, port: u16, is_tls: bool, client: &Client) -> Option<Version> {
let scheme = if is_tls { "https" } else { "http" };
let url = format!("{}://{}:{}", scheme, ip, port);
match client.get(&url).send().await.ok().map(|r| r.version()) {
Some(version) => Some(version),
None => Some(Version::HTTP_11),
}
connector.connect(sni, stream).await.is_ok()
}
src/utils/lazylock.rs
+3 -1
View File
@@ -4,9 +4,10 @@ use pingora_cache::eviction::simple_lru::Manager;
use pingora_cache::lock::CacheLock;
use pingora_cache::MemCache;
use pingora_limits::rate::Rate;
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::OnceLock;
use std::sync::{Arc, LazyLock};
use std::sync::{OnceLock, RwLock};
use std::time::Duration;
pub static REVERSE_STORE: LazyLock<DashMap<String, String>> = LazyLock::new(DashMap::new);
@@ -17,3 +18,4 @@ pub static MEM_CACHE: LazyLock<MemCache> = LazyLock::new(MemCache::new);
pub static CACHE_LOCK: LazyLock<CacheLock> = LazyLock::new(|| CacheLock::new(Duration::from_secs(30)));
pub static EVICTION: OnceLock<Manager> = OnceLock::new();
pub static CACHE_TTL: OnceLock<u64> = OnceLock::new();
pub static CHALLENGES: LazyLock<RwLock<HashMap<String, String>>> = LazyLock::new(|| RwLock::new(HashMap::new()));
src/utils/tools.rs
+20
View File
@@ -377,3 +377,23 @@ pub fn write_pid_file(path: &str) -> std::io::Result<()> {
file.write_all(process::id().to_string().as_bytes())?;
Ok(())
}
#[allow(clippy::needless_return)]
pub fn split_host_port(addr: &str, tls: bool) -> Option<(&str, u16, bool, &str)> {
match addr.split_once(':') {
Some((h, p)) => match p.parse::<u16>() {
Ok(port) => return Some((h, port, tls, h)),
Err(_) => {
log::warn!("ForwardAuth: invalid port in {}", addr);
return None;
}
},
None => {
if tls {
return Some((addr, 443u16, tls, addr));
} else {
return Some((addr, 80u16, tls, addr));
}
}
};
}
src/web/acme.rs
+1 -5
View File
@@ -1,5 +1,5 @@
use crate::tls::acme::order::CHALLENGES;
use crate::tls::acme::{account, order};
use crate::utils::lazylock::CHALLENGES;
use axum::body::Body;
use axum::extract::State;
use axum::http::{Response, StatusCode};
@@ -44,10 +44,6 @@ pub async fn acme_order(State(state): State<crate::web::webserver::AppState>, ax
}
pub async fn http01_challenge(axum::extract::Path(token): axum::extract::Path<String>) -> impl IntoResponse {
if let Ok(challenges) = CHALLENGES.read() {
// for k in challenges.iter() {
// println!(" ==> {} : {}", k.0, k.1);
// }
if let Some(key_authorization) = challenges.get(&token) {
return Response::builder()
.status(StatusCode::OK)