#![warn(missing_docs)]
#![deny(unsafe_code)]
use bytes::Bytes;
use http::{header::HeaderValue, Request, Response};
use http_body::{Body as HttpBody, Frame};
use http_body_util::BodyExt;
use http_cache::{CacheManager, HttpResponse, HttpVersion};
use http_cache_semantics::{BeforeRequest, CachePolicy};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::error::Error as StdError;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::{Duration, SystemTime};
use tower::{Layer, Service};
type BoxError = Box<dyn StdError + Send + Sync>;
#[derive(Debug, Default)]
pub struct CacheMetrics {
pub hits: AtomicU64,
pub misses: AtomicU64,
pub stores: AtomicU64,
pub skipped: AtomicU64,
}
impl CacheMetrics {
pub fn new() -> Self {
Self::default()
}
pub fn hit_rate(&self) -> f64 {
let hits = self.hits.load(Ordering::Relaxed);
let total = hits + self.misses.load(Ordering::Relaxed);
if total == 0 {
0.0
} else {
hits as f64 / total as f64
}
}
pub fn reset(&self) {
self.hits.store(0, Ordering::Relaxed);
self.misses.store(0, Ordering::Relaxed);
self.stores.store(0, Ordering::Relaxed);
self.skipped.store(0, Ordering::Relaxed);
}
}
pub trait Keyer: Clone + Send + Sync + 'static {
fn cache_key<B>(&self, req: &Request<B>) -> String;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultKeyer;
impl Keyer for DefaultKeyer {
fn cache_key<B>(&self, req: &Request<B>) -> String {
format!("{} {}", req.method(), req.uri().path())
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct QueryKeyer;
impl Keyer for QueryKeyer {
fn cache_key<B>(&self, req: &Request<B>) -> String {
format!("{} {}", req.method(), req.uri())
}
}
#[derive(Clone)]
pub struct CustomKeyer<F> {
func: F,
}
impl<F> CustomKeyer<F> {
pub fn new(func: F) -> Self {
Self { func }
}
}
impl<F> Keyer for CustomKeyer<F>
where
F: Fn(&Request<()>) -> String + Clone + Send + Sync + 'static,
{
fn cache_key<B>(&self, req: &Request<B>) -> String {
let mut temp_req = Request::builder()
.method(req.method())
.uri(req.uri())
.version(req.version())
.body(())
.expect("request built from valid parts");
*temp_req.headers_mut() = req.headers().clone();
(self.func)(&temp_req)
}
}
#[derive(Debug, Clone)]
pub struct ServerCacheOptions {
pub default_ttl: Option<Duration>,
pub max_ttl: Option<Duration>,
pub min_ttl: Option<Duration>,
pub cache_status_headers: bool,
pub max_body_size: usize,
pub cache_by_default: bool,
pub respect_vary: bool,
pub respect_authorization: bool,
}
impl Default for ServerCacheOptions {
fn default() -> Self {
Self {
default_ttl: Some(Duration::from_secs(60)),
max_ttl: Some(Duration::from_secs(3600)),
min_ttl: None,
cache_status_headers: true,
max_body_size: 128 * 1024 * 1024,
cache_by_default: false,
respect_vary: true,
respect_authorization: true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CachedResponse {
pub status: u16,
pub headers: HashMap<String, Vec<String>>,
pub body: Vec<u8>,
pub cached_at: SystemTime,
pub ttl: Duration,
pub vary: Option<Vec<String>>,
}
impl CachedResponse {
pub fn is_stale(&self) -> bool {
SystemTime::now()
.duration_since(self.cached_at)
.unwrap_or(Duration::MAX)
> self.ttl
}
pub fn into_response(
self,
) -> Result<Response<Bytes>, Box<dyn std::error::Error + Send + Sync>> {
let mut builder = Response::builder().status(self.status);
for (key, values) in self.headers {
for value in values {
if let Ok(header_value) = HeaderValue::from_str(&value) {
builder = builder.header(&key, header_value);
}
}
}
Ok(builder.body(Bytes::from(self.body))?)
}
}
#[derive(Debug)]
pub enum ResponseBody {
Cached(Bytes),
Fresh(Bytes),
Uncacheable(Bytes),
}
impl HttpBody for ResponseBody {
type Data = Bytes;
type Error = BoxError;
fn poll_frame(
mut self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Option<std::result::Result<Frame<Self::Data>, Self::Error>>> {
let bytes = match &mut *self {
ResponseBody::Cached(b)
| ResponseBody::Fresh(b)
| ResponseBody::Uncacheable(b) => {
std::mem::replace(b, Bytes::new())
}
};
if bytes.is_empty() {
Poll::Ready(None)
} else {
Poll::Ready(Some(Ok(Frame::data(bytes))))
}
}
fn is_end_stream(&self) -> bool {
match self {
ResponseBody::Cached(b)
| ResponseBody::Fresh(b)
| ResponseBody::Uncacheable(b) => b.is_empty(),
}
}
}
#[derive(Clone)]
pub struct ServerCacheLayer<M, K = DefaultKeyer>
where
M: CacheManager,
K: Keyer,
{
manager: M,
keyer: K,
options: ServerCacheOptions,
metrics: Arc<CacheMetrics>,
}
impl<M> ServerCacheLayer<M, DefaultKeyer>
where
M: CacheManager,
{
pub fn new(manager: M) -> Self {
Self {
manager,
keyer: DefaultKeyer,
options: ServerCacheOptions::default(),
metrics: Arc::new(CacheMetrics::new()),
}
}
}
impl<M, K> ServerCacheLayer<M, K>
where
M: CacheManager,
K: Keyer,
{
pub fn with_keyer(manager: M, keyer: K) -> Self {
Self {
manager,
keyer,
options: ServerCacheOptions::default(),
metrics: Arc::new(CacheMetrics::new()),
}
}
pub fn with_options(mut self, options: ServerCacheOptions) -> Self {
self.options = options;
self
}
pub fn metrics(&self) -> &Arc<CacheMetrics> {
&self.metrics
}
pub async fn invalidate(&self, cache_key: &str) -> Result<(), BoxError> {
self.manager.delete(cache_key).await
}
pub async fn invalidate_request<B>(
&self,
req: &Request<B>,
) -> Result<(), BoxError> {
let cache_key = self.keyer.cache_key(req);
self.invalidate(&cache_key).await
}
}
impl<S, M, K> Layer<S> for ServerCacheLayer<M, K>
where
M: CacheManager + Clone,
K: Keyer,
{
type Service = ServerCacheService<S, M, K>;
fn layer(&self, inner: S) -> Self::Service {
ServerCacheService {
inner,
manager: self.manager.clone(),
keyer: self.keyer.clone(),
options: self.options.clone(),
metrics: self.metrics.clone(),
}
}
}
#[derive(Clone)]
pub struct ServerCacheService<S, M, K>
where
M: CacheManager,
K: Keyer,
{
inner: S,
manager: M,
keyer: K,
options: ServerCacheOptions,
metrics: Arc<CacheMetrics>,
}
impl<S, ReqBody, ResBody, M, K> Service<Request<ReqBody>>
for ServerCacheService<S, M, K>
where
S: Service<Request<ReqBody>, Response = Response<ResBody>>
+ Clone
+ Send
+ 'static,
S::Error: Into<BoxError>,
S::Future: Send + 'static,
M: CacheManager + Clone,
K: Keyer,
ReqBody: Send + 'static,
ResBody: HttpBody + Send + 'static,
ResBody::Data: Send,
ResBody::Error: Into<BoxError>,
{
type Response = Response<ResponseBody>;
type Error = BoxError;
type Future = Pin<
Box<
dyn std::future::Future<
Output = std::result::Result<Self::Response, Self::Error>,
> + Send,
>,
>;
fn poll_ready(
&mut self,
cx: &mut Context<'_>,
) -> Poll<std::result::Result<(), Self::Error>> {
self.inner.poll_ready(cx).map_err(Into::into)
}
fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
let manager = self.manager.clone();
let keyer = self.keyer.clone();
let options = self.options.clone();
let metrics = self.metrics.clone();
let mut inner = self.inner.clone();
Box::pin(async move {
let (req_parts, req_body) = req.into_parts();
let temp_req = Request::from_parts(req_parts.clone(), ());
let cache_key = keyer.cache_key(&temp_req);
if let Ok(Some((cached_resp, policy))) =
manager.get(&cache_key).await
{
if let Ok(cached) =
serde_json::from_slice::<CachedResponse>(&cached_resp.body)
{
let before_req =
policy.before_request(&req_parts, SystemTime::now());
let has_explicit_ttl = cached
.headers
.get("cache-control")
.is_some_and(|cc_values| {
cc_values.iter().any(|cc| {
cc.contains("max-age")
|| cc.contains("s-maxage")
})
});
let is_fresh = match before_req {
BeforeRequest::Fresh(_) => {
true
}
BeforeRequest::Stale { .. } => {
if has_explicit_ttl {
false
} else {
!cached.is_stale()
}
}
};
if is_fresh {
metrics.hits.fetch_add(1, Ordering::Relaxed);
let mut response = cached.into_response()?;
if options.cache_status_headers {
response.headers_mut().insert(
"x-cache",
HeaderValue::from_static("HIT"),
);
}
return Ok(response.map(ResponseBody::Cached));
}
}
}
let req = Request::from_parts(req_parts.clone(), req_body);
metrics.misses.fetch_add(1, Ordering::Relaxed);
let response = inner.call(req).await.map_err(Into::into)?;
let (res_parts, body) = response.into_parts();
if let Some(ttl) = should_cache(&req_parts, &res_parts, &options) {
let body_bytes = match collect_body(body).await {
Ok(bytes) => bytes,
Err(e) => {
return Err(e);
}
};
if body_bytes.len() <= options.max_body_size {
metrics.stores.fetch_add(1, Ordering::Relaxed);
let cached = CachedResponse {
status: res_parts.status.as_u16(),
headers: {
let mut map: HashMap<String, Vec<String>> =
HashMap::new();
for (k, v) in res_parts.headers.iter() {
if let Ok(s) = v.to_str() {
map.entry(k.to_string())
.or_default()
.push(s.to_string());
}
}
map
},
body: body_bytes.to_vec(),
cached_at: SystemTime::now(),
ttl,
vary: extract_vary_headers(&res_parts),
};
let cached_json = serde_json::to_vec(&cached)
.map_err(|e| Box::new(e) as BoxError)?;
let http_response = HttpResponse {
body: cached_json,
headers: Default::default(),
status: 200,
url: cache_key.clone().parse().unwrap_or_else(|_| {
"http://localhost/".parse().unwrap()
}),
version: HttpVersion::Http11,
metadata: None,
};
let policy_req = Request::from_parts(req_parts.clone(), ());
let policy_res =
Response::from_parts(res_parts.clone(), ());
let policy = CachePolicy::new(&policy_req, &policy_res);
let manager_clone = manager.clone();
tokio::spawn(async move {
let _ = manager_clone
.put(cache_key, http_response, policy)
.await;
});
} else {
metrics.skipped.fetch_add(1, Ordering::Relaxed);
}
let mut response = Response::from_parts(res_parts, body_bytes);
if options.cache_status_headers {
response
.headers_mut()
.insert("x-cache", HeaderValue::from_static("MISS"));
}
return Ok(response.map(ResponseBody::Fresh));
}
metrics.skipped.fetch_add(1, Ordering::Relaxed);
let body_bytes = collect_body(body).await?;
Ok(Response::from_parts(res_parts, body_bytes)
.map(ResponseBody::Uncacheable))
})
}
}
async fn collect_body<B>(body: B) -> std::result::Result<Bytes, BoxError>
where
B: HttpBody,
B::Error: Into<BoxError>,
{
body.collect()
.await
.map(|collected| collected.to_bytes())
.map_err(Into::into)
}
fn extract_vary_headers(parts: &http::response::Parts) -> Option<Vec<String>> {
parts
.headers
.get(http::header::VARY)
.and_then(|v| v.to_str().ok())
.map(|s| s.split(',').map(|h| h.trim().to_string()).collect())
}
fn has_directive(cache_control: &str, directive: &str) -> bool {
cache_control
.split(',')
.map(|d| d.trim())
.any(|d| d == directive || d.starts_with(&format!("{}=", directive)))
}
fn response_permits_authorized_caching(cc_str: &str) -> bool {
has_directive(cc_str, "public")
|| has_directive(cc_str, "s-maxage")
|| has_directive(cc_str, "must-revalidate")
}
fn should_cache(
req_parts: &http::request::Parts,
res_parts: &http::response::Parts,
options: &ServerCacheOptions,
) -> Option<Duration> {
if !res_parts.status.is_success() {
return None;
}
let has_authorization =
req_parts.headers.contains_key(http::header::AUTHORIZATION);
if let Some(cc) = res_parts.headers.get(http::header::CACHE_CONTROL) {
let cc_str = cc.to_str().ok()?;
if has_authorization
&& options.respect_authorization
&& !response_permits_authorized_caching(cc_str)
{
return None;
}
if has_directive(cc_str, "no-store") {
return None;
}
if has_directive(cc_str, "no-cache") {
return None;
}
if has_directive(cc_str, "private") {
return None;
}
if let Some(s_maxage) = parse_s_maxage(cc_str) {
let ttl = Duration::from_secs(s_maxage);
let ttl = apply_ttl_constraints(ttl, options);
return Some(ttl);
}
if let Some(max_age) = parse_max_age(cc_str) {
let ttl = Duration::from_secs(max_age);
let ttl = apply_ttl_constraints(ttl, options);
return Some(ttl);
}
if has_directive(cc_str, "public") {
return options.default_ttl;
}
} else {
if has_authorization && options.respect_authorization {
return None;
}
}
if let Some(expires) = res_parts.headers.get(http::header::EXPIRES) {
if let Ok(expires_str) = expires.to_str() {
if let Some(ttl) = parse_expires(expires_str) {
let ttl = apply_ttl_constraints(ttl, options);
return Some(ttl);
}
}
}
if options.cache_by_default {
options.default_ttl
} else {
None
}
}
fn apply_ttl_constraints(
ttl: Duration,
options: &ServerCacheOptions,
) -> Duration {
let mut result = ttl;
if let Some(max) = options.max_ttl {
result = result.min(max);
}
if let Some(min) = options.min_ttl {
result = result.max(min);
}
result
}
fn parse_max_age(cache_control: &str) -> Option<u64> {
for directive in cache_control.split(',') {
let directive = directive.trim();
if let Some(value) = directive.strip_prefix("max-age=") {
return value.parse().ok();
}
}
None
}
fn parse_s_maxage(cache_control: &str) -> Option<u64> {
for directive in cache_control.split(',') {
let directive = directive.trim();
if let Some(value) = directive.strip_prefix("s-maxage=") {
return value.parse().ok();
}
}
None
}
fn parse_expires(expires: &str) -> Option<Duration> {
let expires_time = httpdate::parse_http_date(expires).ok()?;
let now = SystemTime::now();
expires_time.duration_since(now).ok()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_keyer() {
let keyer = DefaultKeyer;
let req = Request::get("/users/123").body(()).unwrap();
let key = keyer.cache_key(&req);
assert_eq!(key, "GET /users/123");
}
#[test]
fn test_query_keyer() {
let keyer = QueryKeyer;
let req = Request::get("/users?page=1").body(()).unwrap();
let key = keyer.cache_key(&req);
assert_eq!(key, "GET /users?page=1");
}
#[test]
fn test_parse_max_age() {
assert_eq!(parse_max_age("max-age=3600"), Some(3600));
assert_eq!(parse_max_age("public, max-age=3600"), Some(3600));
assert_eq!(parse_max_age("max-age=3600, public"), Some(3600));
assert_eq!(parse_max_age("public"), None);
}
#[test]
fn test_parse_s_maxage() {
assert_eq!(parse_s_maxage("s-maxage=7200"), Some(7200));
assert_eq!(parse_s_maxage("public, s-maxage=7200"), Some(7200));
assert_eq!(parse_s_maxage("s-maxage=7200, max-age=3600"), Some(7200));
assert_eq!(parse_s_maxage("public"), None);
}
#[test]
fn test_apply_ttl_constraints() {
let options = ServerCacheOptions {
min_ttl: Some(Duration::from_secs(10)),
max_ttl: Some(Duration::from_secs(100)),
..Default::default()
};
assert_eq!(
apply_ttl_constraints(Duration::from_secs(5), &options),
Duration::from_secs(10)
);
assert_eq!(
apply_ttl_constraints(Duration::from_secs(50), &options),
Duration::from_secs(50)
);
assert_eq!(
apply_ttl_constraints(Duration::from_secs(200), &options),
Duration::from_secs(100)
);
}
}