From f895ba4554a3dbf8a921c75f4ceb4d8983f58d49 Mon Sep 17 00:00:00 2001 From: Ara Sadoyan Date: Mon, 19 Oct 2026 19:13:24 +0200 Subject: [PATCH] Restructured cache. Added Redis|Valkey|Dragonfly support. --- Cargo.lock | 74 +++++++++++++- Cargo.toml | 1 + etc/main.yaml | 12 ++- src/cache/mod.rs | 2 + src/cache/redis.rs | 222 +++++++++++++++++++++++++++++++++++++++++ src/cache/types.rs | 41 ++++++++ src/core/bgservice.rs | 2 + src/core/proxyhttp.rs | 19 ++-- src/core/start.rs | 34 ++++--- src/lib.rs | 1 + src/utils/parceyaml.rs | 24 ++--- src/utils/tools.rs | 37 ++++++- src/utils/types.rs | 2 + 13 files changed, 426 insertions(+), 45 deletions(-) create mode 100644 src/cache/mod.rs create mode 100644 src/cache/redis.rs create mode 100644 src/cache/types.rs diff --git a/Cargo.lock b/Cargo.lock index ad6f55b..997815e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -173,6 +173,7 @@ dependencies = [ "prometheus", "rand 0.10.2", "rcgen", + "redis", "reqwest", "rustls 0.23.41", "rustls-pemfile", @@ -198,6 +199,12 @@ dependencies = [ "rustversion", ] +[[package]] +name = "arcstr" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" + [[package]] name = "arraydeque" version = "0.5.1" @@ -285,6 +292,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -1080,12 +1098,16 @@ checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" [[package]] name = "combine" -version = "4.6.7" +version = "4.6.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +checksum = "4356eba6976435498d05a311ee9fc486da3bc81c200d3f0d7dd9703b41913f92" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", ] [[package]] @@ -1481,7 +1503,7 @@ dependencies = [ "asn1-rs", "displaydoc", "nom", - "num-bigint", + "num-bigint 0.4.8", "num-traits", "rusticata-macros", ] @@ -3350,6 +3372,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-bigint-dig" version = "0.8.6" @@ -4224,6 +4256,34 @@ dependencies = [ "yasna", ] +[[package]] +name = "redis" +version = "1.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5b98ae99461e02895c3a96a25dbe4d852a3d5e5b815d31d5d9e4b5344d99e9" +dependencies = [ + "arc-swap", + "arcstr", + "async-lock", + "backon", + "bytes", + "cfg-if", + "combine", + "futures-channel", + "futures-util", + "itoa", + "num-bigint 0.5.1", + "percent-encoding", + "pin-project-lite", + "ryu", + "sha1_smol", + "socket2 0.6.4", + "tokio", + "tokio-util", + "url", + "xxhash-rust", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -4885,6 +4945,12 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sha2" version = "0.10.9" @@ -4971,7 +5037,7 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d585997b0ac10be3c5ee635f1bab02d512760d14b7c468801ac8a01d9ae5f1d" dependencies = [ - "num-bigint", + "num-bigint 0.4.8", "num-traits", "thiserror 2.0.18", "time", diff --git a/Cargo.toml b/Cargo.toml index 179f0c2..b7ea7a9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -72,3 +72,4 @@ reqwest = "0.12.28" hickory-resolver = { version = "0.26", features = ["tokio"] } aws-sdk-route53 = "1.127.0" aws-config = { version = "1.12.0", features = ["behavior-version-latest"] } +redis = { version = "1.7.1", features = ["tokio-comp", "connection-manager"] } diff --git a/etc/main.yaml b/etc/main.yaml index 6d5d6ff..9716122 100644 --- a/etc/main.yaml +++ b/etc/main.yaml @@ -4,7 +4,12 @@ runuser: aralez # Username for running aralez after dropping root privileges, re rungroup: aralez # Group for running aralez after dropping root privileges, requires program to start as root upstream_keepalive_pool_size: 500 # Pool size for upstream keepalive connections pid_file: /tmp/aralez.pid # Path to PID file +#error_log: /tmp/aralez_err.log # Path to error log upgrade_sock: /tmp/aralez.sock # Path to socket file +cache_enabled: true # Optional. Disabled by default +cache_storage: memory # memory | redis://host:port . Optional Defaults to memory works with Redis, Valkey, Dragonfly (requires redis:// prefix) +cache_size_mb: 50 # Optional. Set internal cache with size in megabytes of value. Defaults to 50Mb if cache_enabled: true +cache_ttl: 60 # Optional. Set Default TTL in seconds for cached items, works if cache control headers are not set via upstream. Defaults to 60 if cache_enabled: true config_api_enabled: true # Boolean to enable/disable remote config push capability. config_address: 127.0.0.1:3000 # HTTP API address for pushing upstreams.yaml from remote location proxy_address_http: 0.0.0.0:80 # Proxy HTTP bind address @@ -23,6 +28,7 @@ hc_interval: 2 #Interval for health checks in seconds tcp_keepalive_idle: 60 # Seconds of inactivity before the kernel starts sending keepalive probes to a downstream client tcp_keepalive_interval: 10 # Seconds between individual keepalive probes if the client does not respond tcp_keepalive_count: 5 # Number of unanswered probes before the kernel declares the connection dead and closes it -cache_size_mb: 50 # Optional. Enable internal cache with size in megabytes of value. Cache is entirely disabled when this line is commented out -cache_ttl: 120 # Optional. Set Default TTL in seconds for cached items, works if cache control headers are not set via upstream. Defaults to 60 if cache_size_mb is set - +#cluster_member: master | worker Optional. Enables cluster mode, Sets node type to master or worker. Disabled by default +#worker_nodes: http://127.0.0.1:8000, http://127.0.0.1:4000 Required is if cluster_member: master. Comma separated list of worker nodes http(s)://host:port +#log_structured: elasticsearch Optional. Enables structured logging and sets backend type. Only elasticsearch et the moment. +#acme_dns_provider: cloudflare Optional. Enable ACME DNS-01 challenge, sets DNS provider. ONly coudflare, route53, digitalocean at the moment. diff --git a/src/cache/mod.rs b/src/cache/mod.rs new file mode 100644 index 0000000..c7372e0 --- /dev/null +++ b/src/cache/mod.rs @@ -0,0 +1,2 @@ +pub mod redis; +pub mod types; diff --git a/src/cache/redis.rs b/src/cache/redis.rs new file mode 100644 index 0000000..6ace973 --- /dev/null +++ b/src/cache/redis.rs @@ -0,0 +1,222 @@ +use async_trait::async_trait; +use bytes::Bytes; +use log::warn; +use pingora::Result; +use pingora_cache::storage::{HandleHit, HandleMiss, HitHandler, MissFinishType, MissHandler, Storage}; +use pingora_cache::trace::SpanHandle; +use pingora_cache::{CacheKey, CacheMeta, PurgeOutcome, PurgeTarget, PurgeType}; +use redis::aio::ConnectionManager; +use redis::{AsyncCommands, Client, RedisResult}; +use std::any::Any; +use std::borrow::Cow; +use std::sync::Arc; + +#[derive(Clone)] +pub struct RedisCache { + pub client: Arc, + pub ttl: u64, +} + +pub struct RedisHitHandler { + body: Option, +} + +pub struct RedisMissHandler { + client: Arc, + body_key: String, + body_buffer: Vec, + ttl: u64, +} + +impl RedisCache { + pub async fn new(url: &str, ttl: u64) -> RedisResult { + let client = Client::open(url)?; + let manager = ConnectionManager::new(client).await?; + Ok(Self { client: Arc::new(manager), ttl }) + } + fn make_key(key: &CacheKey) -> Cow<'_, str> { + match key.primary_key_str() { + Some(s) => Cow::Borrowed(s), + None => Cow::Owned(key.primary_key().iter().map(|b| format!("{:02x}", b)).collect()), + } + } +} + +#[async_trait] +impl HandleHit for RedisHitHandler { + async fn read_body(&mut self) -> Result> { + Ok(self.body.take()) + } + + async fn finish(self: Box, _storage: &'static (dyn Storage + Sync), _key: &CacheKey, _trace: &SpanHandle) -> Result<()> { + Ok(()) + } + + fn can_seek(&self) -> bool { + false + } + + fn seek(&mut self, _start: usize, _end: Option) -> Result<()> { + Ok(()) + } + + fn should_count_access(&self) -> bool { + true + } + + fn get_eviction_weight(&self) -> usize { + self.body.as_ref().map(|b| b.len()).unwrap_or(0) + } + + fn entry_id(&self) -> Option { + None + } + + fn as_any(&self) -> &(dyn Any + Send + Sync) { + self + } + + fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync) { + self + } +} + +#[async_trait] +impl HandleMiss for RedisMissHandler { + async fn write_body(&mut self, data: Bytes, _eof: bool) -> Result<()> { + self.body_buffer.extend_from_slice(&data); + Ok(()) + } + + async fn finish(self: Box) -> Result { + let len = self.body_buffer.len(); + let mut conn = (*self.client).clone(); + if let Err(e) = conn.set_ex::<_, _, ()>(&self.body_key, &self.body_buffer, self.ttl).await { + warn!("Failed to write cache entry to Redis: {e}"); + } + Ok(MissFinishType::Appended(len, None)) + } + + fn entry_id(&self) -> Option { + None + } +} + +#[async_trait] +impl Storage for RedisCache { + async fn lookup(&'static self, key: &CacheKey, _trace: &SpanHandle) -> Result> { + let base = Self::make_key(key); + + // Pre-allocate exact capacities to prevent reallocations + let mut meta_key = String::with_capacity(base.len() + 5); + meta_key.push_str(&base); + meta_key.push_str(":meta"); + + let mut body_key = String::with_capacity(base.len() + 5); + body_key.push_str(&base); + body_key.push_str(":body"); + + let mut internal_key = String::with_capacity(meta_key.len() + 9); + internal_key.push_str(&meta_key); + internal_key.push_str(":internal"); + + let mut header_key = String::with_capacity(meta_key.len() + 7); + header_key.push_str(&meta_key); + header_key.push_str(":header"); + + let mut conn = (*self.client).clone(); + + let internal: Option> = conn.get(&internal_key).await.ok().flatten(); + let header: Option> = conn.get(&header_key).await.ok().flatten(); + let body: Option> = conn.get(&body_key).await.ok().flatten(); + + match (internal, header, body) { + (Some(internal_bytes), Some(header_bytes), Some(body_bytes)) => match CacheMeta::deserialize(&internal_bytes, &header_bytes) { + Ok(meta) => { + let handler: HitHandler = Box::new(RedisHitHandler { + body: Some(Bytes::from(body_bytes)), + }); + Ok(Some((meta, handler))) + } + Err(_) => { + let _: RedisResult<()> = conn.del(&[&internal_key, &header_key, &body_key]).await; + Ok(None) + } + }, + _ => Ok(None), + } + } + + async fn get_miss_handler(&'static self, key: &CacheKey, meta: &CacheMeta, _trace: &SpanHandle) -> Result { + let base = Self::make_key(key); + + let mut meta_key = String::with_capacity(base.len() + 5); + meta_key.push_str(&base); + meta_key.push_str(":meta"); + + let mut body_key = String::with_capacity(base.len() + 5); + body_key.push_str(&base); + body_key.push_str(":body"); + + let mut internal_key = String::with_capacity(meta_key.len() + 9); + internal_key.push_str(&meta_key); + internal_key.push_str(":internal"); + + let mut header_key = String::with_capacity(meta_key.len() + 7); + header_key.push_str(&meta_key); + header_key.push_str(":header"); + + let ttl = self.ttl; + let (internal_bytes, header_bytes) = meta.serialize()?; + let mut conn = (*self.client).clone(); + + let _: RedisResult<()> = conn.set_ex(&internal_key, &internal_bytes, ttl).await; + let _: RedisResult<()> = conn.set_ex(&header_key, &header_bytes, ttl).await; + + Ok(Box::new(RedisMissHandler { + client: self.client.clone(), + body_key, + body_buffer: Vec::with_capacity(4096), + ttl, + })) + } + + async fn purge(&'static self, _target: PurgeTarget<'_>, _purge_type: PurgeType, _trace: &SpanHandle) -> Result { + Ok(PurgeOutcome::NotFound) + } + + async fn update_meta(&'static self, key: &CacheKey, meta: &CacheMeta, _trace: &SpanHandle) -> Result { + let base = Self::make_key(key); + + let mut meta_key = String::with_capacity(base.len() + 5); + meta_key.push_str(&base); + meta_key.push_str(":meta"); + + let mut internal_key = String::with_capacity(meta_key.len() + 9); + internal_key.push_str(&meta_key); + internal_key.push_str(":internal"); + + let mut header_key = String::with_capacity(meta_key.len() + 7); + header_key.push_str(&meta_key); + header_key.push_str(":header"); + + let ttl = self.ttl; + match meta.serialize() { + Ok((internal_bytes, header_bytes)) => { + let mut conn = (*self.client).clone(); + let _: RedisResult<()> = conn.set_ex(&internal_key, &internal_bytes, ttl).await; + let _: RedisResult<()> = conn.set_ex(&header_key, &header_bytes, ttl).await; + Ok(true) + } + Err(_) => Ok(false), + } + } + + fn support_streaming_partial_write(&self) -> bool { + false + } + + fn as_any(&self) -> &(dyn Any + Send + Sync) { + self + } +} diff --git a/src/cache/types.rs b/src/cache/types.rs new file mode 100644 index 0000000..7cd4e80 --- /dev/null +++ b/src/cache/types.rs @@ -0,0 +1,41 @@ +use crate::cache::redis::RedisCache; +use crate::utils::lazylock::{CACHE_TTL, EVICTION}; +use crate::utils::types::AppConfig; +use log::{info, warn}; +use pingora_cache::eviction::simple_lru::Manager; +use pingora_cache::{MemCache, Storage}; +use std::sync::{Arc, OnceLock}; + +pub static GLOBAL_STORAGE: OnceLock> = OnceLock::new(); + +pub async fn init_cache(config: Arc) { + let cache_storage = config.cache_storage.clone().unwrap_or_else(|| { + info!("No cache_storage configured, defaulting to memory"); + "memory".to_string() + }); + let cache_type = cache_storage.splitn(2, ':').next().unwrap_or("memory"); + let ttl = config.cache_ttl.unwrap_or(60); + + match cache_type { + "memory" => { + CACHE_TTL.set(ttl).ok(); + if let Some(size) = config.cache_size_mb { + let cache_size = size * 1024 * 1024; + drop(EVICTION.set(Manager::new(cache_size))); + } + GLOBAL_STORAGE.set(Box::new(MemCache::new())).ok(); + info!("Cache initialized successfully. Storage: {}", cache_storage); + } + "redis" => match RedisCache::new(&cache_storage, ttl).await { + Ok(cache) => { + CACHE_TTL.set(ttl).ok(); + GLOBAL_STORAGE.set(Box::new(cache)).ok(); + info!("Cache initialized successfully. Storage: {}", cache_storage); + } + Err(e) => log::error!("Cache storage initialization failed: {}", e), + }, + _ => { + warn!("Unsupported storage, disabling caching: {}", cache_storage); + } + } +} diff --git a/src/core/bgservice.rs b/src/core/bgservice.rs index 4e4bd23..1bacc7f 100644 --- a/src/core/bgservice.rs +++ b/src/core/bgservice.rs @@ -1,3 +1,4 @@ +use crate::cache::types::init_cache; use crate::core::proxyhttp::LB; use crate::ingress::{APIUpstreamProvider, ConsulProvider, Discovery, FromFileProvider, KubernetesProvider}; use crate::logging::core::{init_access_log, init_access_logging, log_builder}; @@ -84,6 +85,7 @@ impl BackgroundService for LB { let acme_provider = self.config.acme_dns_provider.clone(); tokio::spawn(async move { refresh_order(certdir, confdir, acme_provider).await }); + init_cache(self.config.clone()).await; init_access_logging(self.config.access_log.clone()).await; loop { tokio::select! { diff --git a/src/core/proxyhttp.rs b/src/core/proxyhttp.rs index a656490..b68d368 100644 --- a/src/core/proxyhttp.rs +++ b/src/core/proxyhttp.rs @@ -1,7 +1,8 @@ use crate::auth::authenticate; +use crate::cache::types::GLOBAL_STORAGE; use crate::core::gethosts::{GetHost, GetHostsReturHeaders}; use crate::logging::core::access_log; -use crate::utils::lazylock::{CACHE_LOCK, CACHE_TTL, EVICTION, LOCALHOST, MEM_CACHE, RATE_LIMITER, REQUESTS_4XX, REVERSE_STORE}; +use crate::utils::lazylock::{CACHE_LOCK, CACHE_TTL, EVICTION, LOCALHOST, RATE_LIMITER, REQUESTS_4XX, REVERSE_STORE}; use crate::utils::metrics::*; use crate::utils::types::{AppConfig, Extraparams, Headers, InnerMap, UpstreamsDashMap, UpstreamsIdMap}; use arc_swap::ArcSwap; @@ -67,8 +68,7 @@ impl ProxyHttp for LB { fn cache_key_callback(&self, session: &Session, ctx: &mut Self::CTX) -> Result { let host = ctx.hostname.as_deref().unwrap_or("default"); - let path = session.req_header().uri.path(); - + let path = session.req_header().uri.path_and_query().map_or("/", |pq| pq.as_str()); let primary_key = KEY_BUF.with(|buf| { let mut b = buf.borrow_mut(); b.clear(); @@ -76,15 +76,14 @@ impl ProxyHttp for LB { b.push_str(path); b.clone() }); - - let secondary_key = session.req_header().uri.path_and_query().map(|pq| pq.as_str()).unwrap_or("/"); - Ok(CacheKey::new(primary_key, secondary_key)) + Ok(CacheKey::new(primary_key, host /* host is just for debugging */)) } - fn request_cache_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<()> { - if self.cache_enabled && session.req_header().method == Method::GET { - if let Some(eviction) = EVICTION.get() { - session.cache.enable(&*MEM_CACHE, Some(eviction), None, Some(&*CACHE_LOCK), None); + let has_auth = _ctx.upstream_peer.as_ref().map(|m| m.authorization.is_some()).unwrap_or(false) || _ctx.extraparams.authentication.is_some(); + if self.cache_enabled && session.req_header().method == Method::GET && !has_auth { + if let Some(storage) = GLOBAL_STORAGE.get() { + let eviction_mgr = EVICTION.get().map(|e| e as &(dyn pingora_cache::eviction::EvictionManager + Sync)); + session.cache.enable(storage.as_ref(), eviction_mgr, None, Some(&*CACHE_LOCK), None); } } Ok(()) diff --git a/src/core/start.rs b/src/core/start.rs index 6d552e6..b95b9c0 100644 --- a/src/core/start.rs +++ b/src/core/start.rs @@ -29,20 +29,24 @@ use std::time::Duration; use std::{fs, thread}; pub fn run() { - default_provider().install_default().expect("Failed to install RusTLS crypto provider"); + default_provider().install_default().die("Failed to install RusTLS crypto provider:"); let parameters = Opt::parse_args(); - let file = parameters.conf.clone().unwrap(); + let file = parameters.conf.clone().die("Failed to clone the config:"); 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 { - let cache_size = size * 1024 * 1024; - drop(EVICTION.set(Manager::new(cache_size))); - } - let mut server = Server::new(parameters).unwrap(); + let cache_enabled = maincfg.cache_enabled.unwrap_or(false); + + // ========================================================================== // + CACHE_TTL.set(maincfg.cache_ttl.unwrap_or(60)).die("Failed to set cache TTL"); + if cache_enabled { + let cache_size = maincfg.cache_size_mb.unwrap_or(50); + let _ = EVICTION.set(Manager::new(cache_size * 1024 * 1024)); + } + // ========================================================================== // + + let mut server = Server::new(parameters).die("Failed to start the server:"); server.bootstrap(); let uf_config = Arc::new(DashMap::new()); @@ -106,22 +110,22 @@ pub fn run() { if let Some(bind_address_tls) = bind_address_tls { check_priv(bind_address_tls.as_str()); let (tx, rx): (Sender>, Receiver>) = channel(); - let certs_path = cfg.proxy_configs.clone().unwrap() + "/certificates"; + let certs_path = cfg.proxy_configs.clone().die("Failed certificates path") + "/certificates"; if fs::metadata(certs_path.clone()).is_err() { - fs::create_dir_all(certs_path.clone()).unwrap(); + fs::create_dir_all(certs_path.clone()).die("Failed to create directories"); } thread::spawn(move || { - folder_watch(certs_path, tx).unwrap(); + folder_watch(certs_path, tx).die("Failed to spawn certificates channel"); }); - let certificate_configs = rx.recv().unwrap(); + let certificate_configs = rx.recv().die("Failed to spawn certificates channel"); let first_set = load::Certificates::new(&certificate_configs, grade.as_str()).unwrap_or_else(|| panic!("Unable to load initial certificate info")); let certificates = Arc::new(ArcSwap::from_pointee(first_set)); let certs_for_callback = certificates.clone(); let certs_for_watcher = certificates.clone(); let new_certs = load::Certificates::new(&certificate_configs, grade.as_str()); - certs_for_watcher.store(Arc::new(new_certs.unwrap())); + certs_for_watcher.store(Arc::new(new_certs.die("Failed to build the watcher"))); let mut tls_settings = TlsSettings::intermediate(&certs_for_callback.load().default_cert_path, &certs_for_callback.load().default_key_path).expect("unable to load or parse cert/key"); @@ -162,7 +166,7 @@ pub fn run() { if let Err(e) = write_pid_file(pf.as_str()) { panic!("Failed to write PID file: {} : {}", pf, e); } - let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT]).unwrap(); + let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT]).die("Failed to create signals"); if let Some(sig) = signals.forever().next() { match sig { SIGINT => info!("SIGINT received! Exiting..."), diff --git a/src/lib.rs b/src/lib.rs index 7bf3ad1..547c6f5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ pub mod auth; +pub mod cache; pub mod core; pub mod ingress; pub mod logging; diff --git a/src/utils/parceyaml.rs b/src/utils/parceyaml.rs index 6aa58a3..73bfdc6 100644 --- a/src/utils/parceyaml.rs +++ b/src/utils/parceyaml.rs @@ -2,7 +2,7 @@ use crate::utils::deser::{send_to_worker, serialize_config}; use crate::utils::healthcheck; 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_into; +use crate::utils::tools::{Die, clone_dashmap_into}; use crate::utils::types::*; use dashmap::DashMap; use log::{error, info, warn}; @@ -245,7 +245,7 @@ async fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) { } { - let mut config_guard = GLOBAL_PARAMS.write().expect("Global config lock poisoned"); + let mut config_guard = GLOBAL_PARAMS.write().die("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(); @@ -274,8 +274,8 @@ async fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) { } } 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"); + let data = fs::read_to_string(path).die("Failed to read main config file"); + let mut cfo: AppConfig = noyalib::from_str(&data).die("Failed to parse main config file"); 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") { @@ -301,19 +301,19 @@ 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"); + NODE_STATUS.set(Some(Arc::from(member.as_str()))).die("Config can only be initialized once"); + let mk = cfo.master_key.clone().die("In CLuster mode master key must be set"); + MASTER_KEY.set(Some(Arc::from(mk.as_str()))).die("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"); + IM_A_MASTER.set(true).die("Master status can only be initialized once"); + let worker_str = cfo.worker_nodes.as_deref().die("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"), + port: wk_vec[2].parse::().die("Invalid worker port number"), tls: match wk_vec[0] { "https" => true, "http" => false, @@ -324,13 +324,13 @@ pub fn parce_main_config(path: &str) -> AppConfig { 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"); + // let port = port_str.parse::().die("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"); + WORKERS.set(wks).die("Workers list can only be initialized once"); } } cfo diff --git a/src/utils/tools.rs b/src/utils/tools.rs index d9ef1ef..8bca645 100644 --- a/src/utils/tools.rs +++ b/src/utils/tools.rs @@ -7,7 +7,7 @@ use privdrop::PrivDrop; use serde_json::{Value, json}; use sha2::{Digest, Sha256}; use std::collections::{HashMap, HashSet}; -use std::fmt::Write; +use std::fmt::{Debug, Write}; use std::fs::OpenOptions; use std::io::Write as IoWrite; use std::net::SocketAddr; @@ -364,3 +364,38 @@ pub fn get_hostname() -> String { .filter(|s| !s.is_empty()) .unwrap_or_else(|| "unknown".to_string()) } + +// pub fn die(message: &str) -> ! { +// log::error!("{message}"); +// process::exit(1); +// } +// ========================================= // + +pub trait Die { + fn die(self, message: &str) -> T; +} + +impl Die for Result { + fn die(self, message: &str) -> T { + match self { + Ok(value) => value, + Err(e) => { + eprintln!("\x1b[1;31mERROR:\x1b[0m {message}: {e:?}"); + process::exit(1); + } + } + } +} + +impl Die for Option { + fn die(self, message: &str) -> T { + match self { + Some(value) => value, + None => { + eprintln!("\x1b[1;31mERROR:\x1b[0m {message}"); + process::exit(1); + } + } + } +} +// ========================================= // diff --git a/src/utils/types.rs b/src/utils/types.rs index 7a7b2f5..dedb611 100644 --- a/src/utils/types.rs +++ b/src/utils/types.rs @@ -112,6 +112,8 @@ pub struct AppConfig { pub hc_method: String, pub upstreams_conf: String, pub log_level: String, + pub cache_enabled: Option, + pub cache_storage: Option, pub cache_size_mb: Option, pub cache_ttl: Option, pub access_log: Option,