use super::{
Error, get_bool_conf, get_hash_key, get_str_conf, get_str_slice_conf,
};
use async_trait::async_trait;
use bstr::ByteSlice;
use bytes::Bytes;
use bytesize::ByteSize;
use fancy_regex::Regex;
use http::{Method, StatusCode};
use humantime::parse_duration;
use pingap_cache::{HttpCache, new_cache_backend};
use pingap_config::{PluginCategory, PluginConf};
use pingap_core::{
Ctx, HttpResponse, Plugin, PluginStep, RequestPluginResult,
ensure_client_ip, get_cache_key,
};
use pingap_util::IpRules;
use pingora::cache::CacheOptionOverrides;
use pingora::cache::eviction::EvictionManager;
use pingora::cache::eviction::simple_lru::Manager;
use pingora::cache::key::CacheHashKey;
use pingora::cache::lock::{CacheKeyLock, CacheLock};
use pingora::cache::predictor::{CacheablePredictor, Predictor};
use pingora::proxy::Session;
use std::borrow::Cow;
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::LazyLock;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::time::Duration;
use tracing::{debug, error};
type Result<T> = std::result::Result<T, Error>;
static PREDICTOR: OnceLock<Predictor<32>> = OnceLock::new();
static EVICTION_MANAGER: OnceLock<Manager> = OnceLock::new();
static CACHE_LOCKS: LazyLock<
Mutex<HashMap<Duration, &'static (dyn CacheKeyLock + Send + Sync)>>,
> = LazyLock::new(|| Mutex::new(HashMap::new()));
pub struct Cache {
plugin_step: PluginStep,
eviction: Option<&'static (dyn EvictionManager + Sync)>,
predictor: Option<&'static (dyn CacheablePredictor + Sync)>,
lock: Option<&'static (dyn CacheKeyLock + Send + Sync)>,
lock_retries: Option<usize>,
http_cache: &'static HttpCache,
max_file_size: usize,
max_ttl: Option<Duration>,
namespace: Option<String>,
headers: Option<Vec<String>>,
vary_headers: Option<Arc<Vec<String>>>,
check_cache_control: bool,
purge_ip_rules: IpRules,
skip: Option<Regex>,
hash_value: String,
}
fn get_eviction_manager(cache_max_size: u64) -> &'static Manager {
EVICTION_MANAGER.get_or_init(|| Manager::new(cache_max_size as usize))
}
fn get_cache_lock(
lock: Duration,
) -> Option<&'static (dyn CacheKeyLock + Send + Sync)> {
if lock.is_zero() {
return None;
}
let mut locks = CACHE_LOCKS.lock().ok()?;
if let Some(cache_lock) = locks.get(&lock) {
return Some(*cache_lock);
}
let cache_lock: &'static (dyn CacheKeyLock + Send + Sync) =
Box::leak(CacheLock::new_boxed(lock));
locks.insert(lock, cache_lock);
Some(cache_lock)
}
fn get_predictor() -> &'static (dyn CacheablePredictor + Sync) {
PREDICTOR.get_or_init(|| Predictor::new(128, None))
}
impl TryFrom<&PluginConf> for Cache {
type Error = Error;
fn try_from(value: &PluginConf) -> Result<Self> {
let hash_value = get_hash_key(value);
let directory = get_str_conf(value, "directory");
let cache = new_cache_backend(directory.as_str()).map_err(|e| {
Error::Invalid {
category: "cache".to_string(),
message: e.to_string(),
}
})?;
let cache_max_size = cache.max_size;
let eviction = if value.contains_key("eviction") {
if cache_max_size > 0 {
let eviction = get_eviction_manager(cache_max_size);
Some(eviction as &'static (dyn EvictionManager + Sync))
} else {
error!(
directory,
"eviction is only supported by the memory cache backend, ignoring it"
);
None
}
} else {
None
};
let lock = get_str_conf(value, "lock");
let lock = if !lock.is_empty() {
parse_duration(&lock).map_err(|e| Error::Invalid {
category: PluginCategory::Cache.to_string(),
message: e.to_string(),
})?
} else {
Duration::from_secs(1)
};
let lock_retries = if value.get("lock_retries").is_some() {
let retries =
usize::try_from(crate::get_int_conf(value, "lock_retries"))
.map_err(|_| Error::Invalid {
category: PluginCategory::Cache.to_string(),
message:
"lock_retries should be a non-negative integer"
.to_string(),
})?;
Some(retries)
} else {
None
};
let max_ttl = get_str_conf(value, "max_ttl");
let max_ttl = if !max_ttl.is_empty() {
Some(parse_duration(&max_ttl).map_err(|e| Error::Invalid {
category: PluginCategory::Cache.to_string(),
message: e.to_string(),
})?)
} else {
None
};
let max_file_size = get_str_conf(value, "max_file_size");
let max_file_size = if !max_file_size.is_empty() {
ByteSize::from_str(&max_file_size).map_err(|e| Error::Invalid {
category: PluginCategory::Cache.to_string(),
message: e.to_string(),
})?
} else {
ByteSize::mb(1)
};
let namespace = get_str_conf(value, "namespace");
if !namespace.is_empty() && cache.directory.is_some() {
let path = format!(
"{}/{namespace}",
cache.directory.clone().unwrap_or_default()
);
if let Err(e) = std::fs::create_dir_all(&path) {
error!(
error = e.to_string(),
path, "create directory of cache fail"
);
}
}
let namespace = if namespace.is_empty() {
None
} else {
Some(namespace)
};
let headers = get_str_slice_conf(value, "headers");
let headers = if headers.is_empty() {
None
} else {
Some(headers)
};
let vary_headers = get_str_slice_conf(value, "vary_headers");
let vary_headers = if vary_headers.is_empty() {
None
} else {
Some(Arc::new(
vary_headers
.iter()
.map(|name| name.trim().to_ascii_lowercase())
.collect(),
))
};
let predictor = if value.contains_key("predictor") {
Some(get_predictor())
} else {
None
};
let purge_ip_rules =
IpRules::try_new(&get_str_slice_conf(value, "purge_ip_list"))
.map_err(|e| Error::Invalid {
category: PluginCategory::Cache.to_string(),
message: e.to_string(),
})?;
let skip_value = get_str_conf(value, "skip");
let skip = if skip_value.is_empty() {
None
} else {
Some(Regex::new(&skip_value).map_err(|e| Error::Regex {
category: "cache".to_string(),
source: Box::new(e),
})?)
};
let params = Self {
hash_value,
http_cache: cache,
plugin_step: PluginStep::Request,
eviction,
predictor,
lock: get_cache_lock(lock),
lock_retries,
max_ttl,
max_file_size: max_file_size.as_u64() as usize,
namespace,
headers,
vary_headers,
purge_ip_rules,
check_cache_control: get_bool_conf(value, "check_cache_control"),
skip,
};
Ok(params)
}
}
impl Cache {
fn cache_option_overrides(&self) -> Option<CacheOptionOverrides> {
let retries = self.lock_retries?;
let mut overrides = CacheOptionOverrides::default();
overrides.max_lock_retries = Some(retries);
Some(overrides)
}
pub fn new(params: &PluginConf) -> Result<Self> {
debug!(params = params.to_string(), "new http cache plugin");
Self::try_from(params)
}
}
static METHOD_PURGE: LazyLock<Method> = LazyLock::new(|| {
Method::from_bytes(b"PURGE").expect("Failed to create PURGE method")
});
#[async_trait]
impl Plugin for Cache {
#[inline]
fn config_key(&self) -> Cow<'_, str> {
Cow::Borrowed(&self.hash_value)
}
#[inline]
async fn handle_request(
&self,
step: PluginStep,
session: &mut Session,
ctx: &mut Ctx,
) -> pingora::Result<RequestPluginResult> {
if step != self.plugin_step {
return Ok(RequestPluginResult::Skipped);
}
let req_header = session.req_header();
let method = &req_header.method;
if ![&Method::GET, &Method::HEAD, &*METHOD_PURGE].contains(&method) {
return Ok(RequestPluginResult::Skipped);
}
if let Some(skip) = &self.skip
&& let Some(value) = req_header.uri.path_and_query()
&& skip.is_match(value.as_str()).unwrap_or_default()
{
return Ok(RequestPluginResult::Skipped);
}
let mut keys = Vec::with_capacity(4);
{
let cache_info = ctx.cache.get_or_insert_default();
cache_info.namespace = self.namespace.clone();
}
if let Some(headers) = &self.headers {
for key in headers.iter() {
let buf = session.get_header_bytes(key).to_str_lossy();
if !buf.is_empty() {
keys.push(buf.to_string());
}
}
}
if !keys.is_empty() {
ctx.extend_cache_keys(keys);
if let Some(cache_info) = &ctx.cache {
debug!("Cache keys: {:?}", cache_info.keys);
}
}
if method == *METHOD_PURGE {
let ip = ensure_client_ip(session, ctx);
let found = match self.purge_ip_rules.is_match(ip) {
Ok(matched) => matched,
Err(e) => {
return Ok(RequestPluginResult::Respond(
HttpResponse::bad_request(e.to_string()),
));
},
};
if !found {
return Ok(RequestPluginResult::Respond(HttpResponse {
status: StatusCode::FORBIDDEN,
body: Bytes::from_static(b"Forbidden, ip is not allowed"),
..Default::default()
}));
}
if session.req_header().uri.path() == "/*" {
let Some(namespace) = &self.namespace else {
return Ok(RequestPluginResult::Respond(HttpResponse {
status: StatusCode::NOT_IMPLEMENTED,
body: Bytes::from_static(
b"namespace purge requires the namespace option",
),
..Default::default()
}));
};
let result =
self.http_cache.cache.purge_namespace(namespace).await?;
let Some(stats) = result else {
return Ok(RequestPluginResult::Respond(HttpResponse {
status: StatusCode::NOT_IMPLEMENTED,
body: Bytes::from_static(
b"namespace purge is not supported by the memory cache backend",
),
..Default::default()
}));
};
return Ok(RequestPluginResult::Respond(HttpResponse::text(
format!("purged: {}, fail: {}", stats.success, stats.fail),
)));
}
for method in [Method::GET, Method::HEAD] {
let key = get_cache_key(
ctx,
method.as_ref(),
&session.req_header().uri,
);
self.http_cache
.cache
.remove(&key.combined(), key.user_tag().as_bytes())
.await?;
}
return Ok(
RequestPluginResult::Respond(HttpResponse::no_content()),
);
}
if let Some(cache_info) = &mut ctx.cache {
cache_info.max_ttl = self.max_ttl;
cache_info.check_cache_control = self.check_cache_control;
cache_info.vary_headers = self.vary_headers.clone();
}
session.cache.enable(
self.http_cache,
self.eviction,
self.predictor,
self.lock,
self.cache_option_overrides(),
);
if self.max_file_size > 0 {
session.cache.set_max_file_size_bytes(self.max_file_size);
}
if let Some(stats) = self.http_cache.stats()
&& let Some(cache_info) = ctx.cache.as_mut()
{
cache_info.reading_count = Some(stats.reading);
cache_info.writing_count = Some(stats.writing);
}
Ok(RequestPluginResult::Continue)
}
}
register_plugin!("cache", Cache);
#[cfg(test)]
mod tests {
use super::*;
use pingap_config::PluginConf;
use pingap_core::{Ctx, PluginStep};
use pingora::proxy::Session;
use pretty_assertions::assert_eq;
use tokio_test::io::Builder;
#[test]
fn test_cache_params() {
let params = Cache::try_from(
&toml::from_str::<PluginConf>(
r###"
eviction = true
headers = ["Accept-Encoding"]
lock = "2s"
lock_retries = 5
max_file_size = "100kb"
predictor = true
max_ttl = "1m"
vary_headers = ["Accept-Encoding", " accept "]
"###,
)
.unwrap(),
)
.unwrap();
assert_eq!(
Some(vec!["accept-encoding".to_string(), "accept".to_string()]),
params.vary_headers.as_deref().cloned()
);
assert_eq!(Some(5), params.lock_retries);
assert_eq!(
Some(5),
params.cache_option_overrides().unwrap().max_lock_retries
);
assert_eq!(true, params.eviction.is_some());
assert_eq!(
r#"Some(["Accept-Encoding"])"#,
format!("{:?}", params.headers)
);
assert_eq!(true, params.lock.is_some());
assert_eq!(100 * 1000, params.max_file_size);
assert_eq!(60, params.max_ttl.unwrap().as_secs());
assert_eq!(true, params.predictor.is_some());
let params = Cache::try_from(
&toml::from_str::<PluginConf>(r###"lock = "1s""###).unwrap(),
)
.unwrap();
assert_eq!(None, params.lock_retries);
assert_eq!(true, params.cache_option_overrides().is_none());
let err = Cache::try_from(
&toml::from_str::<PluginConf>("lock_retries = -1").unwrap(),
)
.err()
.unwrap()
.to_string();
assert_eq!(true, err.contains("lock_retries"), "{err}");
}
#[test]
fn test_cache_lock_any_duration() {
let lock_of = |lock: &str| {
Cache::try_from(
&toml::from_str::<PluginConf>(&format!(
r###"lock = "{lock}""###
))
.unwrap(),
)
.unwrap()
.lock
.is_some()
};
for lock in ["1s", "2s", "3s", "5s", "30s", "500ms", "1m"] {
assert_eq!(true, lock_of(lock), "lock({lock}) was disabled");
}
assert_eq!(false, lock_of("0s"));
}
#[tokio::test]
async fn test_cache() {
let cache = Cache::try_from(
&toml::from_str::<PluginConf>(
r###"
namespace = "pingap"
eviction = true
headers = ["Accept-Encoding"]
purge_ip_list = ["127.0.0.1"]
lock = "2s"
max_file_size = "100kb"
predictor = true
max_ttl = "1m"
"###,
)
.unwrap(),
)
.unwrap();
let headers = ["Accept-Encoding: gzip"].join("\r\n");
let input_header =
format!("GET /vicanso/pingap?size=1 HTTP/1.1\r\n{headers}\r\n\r\n");
let mock_io = Builder::new().read(input_header.as_bytes()).build();
let mut session = Session::new_h1(Box::new(mock_io));
session.read_request().await.unwrap();
let mut ctx = Ctx::default();
cache
.handle_request(PluginStep::Request, &mut session, &mut ctx)
.await
.unwrap();
assert_eq!(
"pingap",
ctx.cache.as_ref().unwrap().namespace.as_ref().unwrap()
);
assert_eq!(
"gzip",
ctx.cache.as_ref().unwrap().keys.as_ref().unwrap().join(":")
);
assert_eq!(true, session.cache.enabled());
assert_eq!(100 * 1000, cache.max_file_size);
}
async fn purge(cache: &Cache, path: &str) -> HttpResponse {
let input_header = format!("PURGE {path} HTTP/1.1\r\n\r\n");
let mock_io = Builder::new().read(input_header.as_bytes()).build();
let mut session = Session::new_h1(Box::new(mock_io));
session.read_request().await.unwrap();
let mut ctx = Ctx::default();
ctx.conn.client_ip = Some("127.0.0.1".to_string());
let result = cache
.handle_request(PluginStep::Request, &mut session, &mut ctx)
.await
.unwrap();
let RequestPluginResult::Respond(resp) = result else {
panic!("purge must respond, got a pass-through");
};
resp
}
#[tokio::test]
async fn test_purge_namespace_memory_backend_unsupported() {
let cache = Cache::try_from(
&toml::from_str::<PluginConf>(
r###"
namespace = "purge-mem"
purge_ip_list = ["127.0.0.1"]
"###,
)
.unwrap(),
)
.unwrap();
let resp = purge(&cache, "/*").await;
assert_eq!(StatusCode::NOT_IMPLEMENTED, resp.status);
}
#[tokio::test]
async fn test_purge_namespace_file_backend() {
let dir = tempfile::TempDir::new().unwrap();
let cache = Cache::try_from(
&toml::from_str::<PluginConf>(&format!(
r###"
directory = "{}"
namespace = "purge-ns"
purge_ip_list = ["127.0.0.1"]
"###,
dir.path().to_string_lossy()
))
.unwrap(),
)
.unwrap();
let obj = pingap_cache::CacheObject {
meta: (Bytes::from_static(b"meta0"), Bytes::from_static(b"meta1")),
body: Bytes::from_static(b"cached body"),
};
cache
.http_cache
.cache
.put("purge-key", b"purge-ns", obj.clone())
.await
.unwrap();
cache
.http_cache
.cache
.put("other-key", b"other-ns", obj)
.await
.unwrap();
let resp = purge(&cache, "/*").await;
assert_eq!(StatusCode::OK, resp.status);
assert_eq!(
"purged: 1, fail: 0",
std::str::from_utf8(&resp.body).unwrap()
);
let purged = cache
.http_cache
.cache
.get("purge-key", b"purge-ns")
.await
.unwrap();
assert_eq!(true, purged.is_none());
let kept = cache
.http_cache
.cache
.get("other-key", b"other-ns")
.await
.unwrap();
assert_eq!(true, kept.is_some());
let resp = purge(&cache, "/vicanso/pingap").await;
assert_eq!(StatusCode::NO_CONTENT, resp.status);
}
#[tokio::test]
async fn test_purge_forbidden_ip() {
let cache = Cache::try_from(
&toml::from_str::<PluginConf>(
r###"
namespace = "purge-deny"
purge_ip_list = ["192.168.1.1"]
"###,
)
.unwrap(),
)
.unwrap();
let resp = purge(&cache, "/*").await;
assert_eq!(StatusCode::FORBIDDEN, resp.status);
}
}