diff --git a/src/utils.rs b/src/utils.rs index df140ba..eeaee30 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -6,6 +6,7 @@ pub mod healthcheck; pub mod httpclient; pub mod jwt; pub mod kuberconsul; +pub mod lazylock; pub mod metrics; pub mod parceyaml; pub mod state; diff --git a/src/utils/healthcheck.rs b/src/utils/healthcheck.rs index 427477f..5ba1d3e 100644 --- a/src/utils/healthcheck.rs +++ b/src/utils/healthcheck.rs @@ -1,3 +1,4 @@ +use crate::utils::lazylock::REVERSE_STORE; use crate::utils::structs::{InnerMap, UpstreamsDashMap, UpstreamsIdMap}; use crate::utils::tools::*; use dashmap::DashMap; @@ -20,6 +21,7 @@ pub async fn hc2(upslist: Arc, fullist: Arc, if !compare_dashmaps(&totest, &upslist) { clone_dashmap_into(&totest, &upslist); clone_idmap_into(&totest, &idlist); + REVERSE_STORE.clear(); } } } @@ -139,10 +141,7 @@ async fn detect_tls(ip: &str, port: &u16, client: &Client) -> (bool, Option { - // println!("{} => {:?} (HTTP)", http_url, response.version()); - (false, Some(response.version())) - } + Ok(response) => (false, Some(response.version())), Err(_) => { if ping_grpc(&http_url).await { (false, Some(Version::HTTP_2)) diff --git a/src/utils/lazylock.rs b/src/utils/lazylock.rs new file mode 100644 index 0000000..63676e2 --- /dev/null +++ b/src/utils/lazylock.rs @@ -0,0 +1,11 @@ +use dashmap::DashMap; +use moka::sync::Cache; +use pingora_limits::rate::Rate; +use std::net::IpAddr; +use std::sync::{Arc, LazyLock}; +use std::time::Duration; + +pub static REVERSE_STORE: LazyLock> = LazyLock::new(DashMap::new); +pub static RATE_LIMITER: LazyLock = LazyLock::new(|| Rate::new(Duration::from_secs(1))); +pub static REQUESTS_4XX: LazyLock> = LazyLock::new(|| Cache::builder().time_to_live(Duration::from_secs(1)).build()); +pub static LOCALHOST: LazyLock> = LazyLock::new(|| Arc::from("localhost")); diff --git a/src/utils/parceyaml.rs b/src/utils/parceyaml.rs index 1bd7d35..601cea8 100644 --- a/src/utils/parceyaml.rs +++ b/src/utils/parceyaml.rs @@ -1,4 +1,5 @@ use crate::utils::healthcheck; +use crate::utils::lazylock::REVERSE_STORE; use crate::utils::state::{is_first_run, mark_not_first_run}; use crate::utils::structs::*; use crate::utils::tools::{clone_dashmap, clone_dashmap_into, print_upstreams}; @@ -110,7 +111,6 @@ pub async fn load_configuration(d: &str, kind: &str) -> (Option, let mut parsed: Config = match serde_yml::from_str(&yaml_data) { Ok(cfg) => cfg, Err(e) => { - println!("================================================"); error!("Failed to parse upstreams file: {}", e); return (None, e.to_string()); } @@ -258,6 +258,7 @@ async fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) { clone_dashmap_into(&r, &config.upstreams); } info!("Upstream Config:"); + REVERSE_STORE.clear(); print_upstreams(&config.upstreams, &config.extraparams); } } diff --git a/src/utils/tools.rs b/src/utils/tools.rs index d7fb120..e80ebe3 100644 --- a/src/utils/tools.rs +++ b/src/utils/tools.rs @@ -1,6 +1,6 @@ use crate::tls::load; use crate::tls::load::CertificateConfig; -use crate::utils::structs::{InnerMap, InnerMapForJson, Extraparams, UpstreamSnapshotForJson, UpstreamsDashMap, UpstreamsIdMap}; +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}; @@ -159,8 +159,8 @@ pub fn clone_idmap_into(original: &UpstreamsDashMap, cloned: &UpstreamsIdMap) { x.port, x.is_http2, x.to_https, - x.rate_limit.unwrap_or_default(), x.x4xx_limit.unwrap_or_default(), + x.rate_limit.unwrap_or_default(), x.healthcheck.unwrap_or_default(), x.authorization ) diff --git a/src/web/bgservice.rs b/src/web/bgservice.rs index d9fafc7..f475a21 100644 --- a/src/web/bgservice.rs +++ b/src/web/bgservice.rs @@ -87,6 +87,7 @@ impl BackgroundService for LB { if let Some(ss) = val { clone_dashmap_into(&ss.upstreams, &self.ump_full); clone_dashmap_into(&ss.upstreams, &self.ump_upst); + clone_idmap_into(&ss.upstreams, &self.ump_byid); let current = self.extraparams.load_full(); let mut new = (*current).clone(); new.to_https = ss.extraparams.to_https; diff --git a/src/web/proxyhttp.rs b/src/web/proxyhttp.rs index d3f6237..ddea322 100644 --- a/src/web/proxyhttp.rs +++ b/src/web/proxyhttp.rs @@ -1,33 +1,29 @@ use crate::utils::auth::authenticate; +use crate::utils::lazylock::{LOCALHOST, RATE_LIMITER, REQUESTS_4XX, REVERSE_STORE}; use crate::utils::metrics::*; use crate::utils::structs::{AppConfig, Extraparams, Headers, InnerMap, UpstreamsDashMap, UpstreamsIdMap}; use crate::web::gethosts::{GetHost, GetHostsReturHeaders}; use arc_swap::ArcSwap; use async_trait::async_trait; use axum::body::Bytes; -use dashmap::DashMap; use log::{debug, error, warn}; -use moka::sync::Cache; use pingora::http::{RequestHeader, ResponseHeader, StatusCode}; use pingora::prelude::*; use pingora::ErrorSource::Upstream; use pingora_core::listeners::ALPN; use pingora_core::prelude::HttpPeer; -use pingora_limits::rate::Rate; use pingora_proxy::{ProxyHttp, Session}; use sha2::{Digest, Sha256}; use std::cell::RefCell; use std::fmt::Write; -use std::net::IpAddr; -use std::sync::{Arc, LazyLock}; -use std::time::Duration; +use std::sync::Arc; use tokio::time::Instant; -static REVERSE_STORE: LazyLock> = LazyLock::new(DashMap::new); thread_local! {static IP_BUFFER: RefCell = RefCell::new(String::with_capacity(50));} -pub static RATE_LIMITER: LazyLock = LazyLock::new(|| Rate::new(Duration::from_secs(1))); -pub static REQUESTS_4XX: LazyLock> = LazyLock::new(|| Cache::builder().time_to_live(Duration::from_secs(1)).build()); -pub static LOCALHOST: LazyLock> = LazyLock::new(|| Arc::from("localhost")); +// static REVERSE_STORE: LazyLock> = LazyLock::new(DashMap::new); +// pub static RATE_LIMITER: LazyLock = LazyLock::new(|| Rate::new(Duration::from_secs(1))); +// pub static REQUESTS_4XX: LazyLock> = LazyLock::new(|| Cache::builder().time_to_live(Duration::from_secs(1)).build()); +// pub static LOCALHOST: LazyLock> = LazyLock::new(|| Arc::from("localhost")); #[derive(Clone)] pub struct LB { @@ -185,12 +181,13 @@ impl ProxyHttp for LB { let mut s = String::with_capacity(64); write!( &mut s, - "{}:{}:{}:{}:{}:{}:{}:{:?}", + "{}:{}:{}:{}:{}:{}:{}:{}:{:?}", hostname, innermap.address, innermap.port, innermap.is_http2, innermap.to_https, + innermap.x4xx_limit.unwrap_or_default(), innermap.rate_limit.unwrap_or_default(), innermap.healthcheck.unwrap_or_default(), innermap.authorization @@ -278,7 +275,6 @@ impl ProxyHttp for LB { buf.push_str(&val.to_string()); buf.push_str("; HttpOnly; SameSite=Lax"); // buf.push_str("; Path=/; Max-Age=86400; HttpOnly; SameSite=Lax"); - // println!("{}", buf); let _ = _upstream_response.insert_header("set-cookie", buf.as_str()); } }