use axum::http::{
HeaderMap, HeaderValue, Method, Request, Response, StatusCode,
header::{CACHE_CONTROL, SET_COOKIE},
};
use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use tower_layer::Layer;
use tower_service::Service;
const DEFAULT_BROWSER_SECS: u32 = 60;
pub fn cdn_cache_for(cdn_secs: u32) -> CdnCacheFor {
CdnCacheFor {
cdn_secs,
browser_secs: DEFAULT_BROWSER_SECS.min(cdn_secs),
}
}
#[derive(Clone, Copy, Debug)]
pub struct CdnCacheFor {
cdn_secs: u32,
browser_secs: u32,
}
impl CdnCacheFor {
pub fn browser(mut self, secs: u32) -> Self {
self.browser_secs = secs.min(self.cdn_secs);
self
}
fn header(&self) -> HeaderValue {
let value = format!(
"public, max-age={}, s-maxage={}",
self.browser_secs, self.cdn_secs
);
HeaderValue::from_str(&value).expect("digits and ASCII are a valid header value")
}
}
impl<S> Layer<S> for CdnCacheFor {
type Service = CdnCacheForService<S>;
fn layer(&self, inner: S) -> Self::Service {
CdnCacheForService {
inner,
policy: *self,
}
}
}
#[derive(Clone, Debug)]
pub struct CdnCacheForService<S> {
inner: S,
policy: CdnCacheFor,
}
type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
impl<S, B, ResBody> Service<Request<B>> for CdnCacheForService<S>
where
S: Service<Request<B>, Response = Response<ResBody>>,
S::Future: Send + 'static,
{
type Response = Response<ResBody>;
type Error = S::Error;
type Future = BoxFuture<Result<Self::Response, Self::Error>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, request: Request<B>) -> Self::Future {
let cacheable = matches!(*request.method(), Method::GET | Method::HEAD);
let policy = self.policy;
let response = self.inner.call(request);
Box::pin(async move {
let mut response = response.await?;
mark_shared(policy, cacheable, response.status(), response.headers_mut());
Ok(response)
})
}
}
fn mark_shared(policy: CdnCacheFor, cacheable: bool, status: StatusCode, headers: &mut HeaderMap) {
if cacheable && status.is_success() {
headers.insert(CACHE_CONTROL, policy.header());
}
}
pub fn cdn_cache_guard(api_prefix: &'static str) -> CdnCacheGuard {
CdnCacheGuard { api_prefix }
}
#[derive(Clone, Copy, Debug)]
pub struct CdnCacheGuard {
api_prefix: &'static str,
}
impl<S> Layer<S> for CdnCacheGuard {
type Service = CdnCacheGuardService<S>;
fn layer(&self, inner: S) -> Self::Service {
CdnCacheGuardService {
inner,
api_prefix: self.api_prefix,
}
}
}
#[derive(Clone, Debug)]
pub struct CdnCacheGuardService<S> {
inner: S,
api_prefix: &'static str,
}
impl<S, B, ResBody> Service<Request<B>> for CdnCacheGuardService<S>
where
S: Service<Request<B>, Response = Response<ResBody>>,
S::Future: Send + 'static,
{
type Response = Response<ResBody>;
type Error = S::Error;
type Future = BoxFuture<Result<Self::Response, Self::Error>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, request: Request<B>) -> Self::Future {
let api = under_prefix(request.uri().path(), self.api_prefix);
let response = self.inner.call(request);
Box::pin(async move {
let mut response = response.await?;
apply_guard(api, response.headers_mut());
Ok(response)
})
}
}
fn under_prefix(path: &str, prefix: &str) -> bool {
path.strip_prefix(prefix.trim_end_matches('/'))
.is_some_and(|rest| rest.is_empty() || rest.starts_with('/'))
}
fn is_shared(headers: &HeaderMap) -> bool {
headers
.get_all(CACHE_CONTROL)
.iter()
.filter_map(|value| value.to_str().ok())
.flat_map(|value| value.split(','))
.any(|part| part.trim().eq_ignore_ascii_case("public"))
}
fn apply_guard(api: bool, headers: &mut HeaderMap) {
if is_shared(headers) {
headers.remove(SET_COOKIE);
} else if api && !headers.contains_key(CACHE_CONTROL) {
headers.insert(CACHE_CONTROL, HeaderValue::from_static("private, no-cache"));
}
}
#[cfg(test)]
mod tests {
use super::*;
fn with_cookie(cache_control: Option<&'static str>) -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(SET_COOKIE, HeaderValue::from_static("session=abc"));
if let Some(value) = cache_control {
headers.insert(CACHE_CONTROL, HeaderValue::from_static(value));
}
headers
}
#[test]
fn shared_responses_never_carry_a_cookie() {
let mut headers = with_cookie(Some("public, max-age=60, s-maxage=300"));
apply_guard(true, &mut headers);
assert!(headers.get(SET_COOKIE).is_none());
}
#[test]
fn other_api_responses_are_private_and_keep_their_cookie() {
let mut headers = with_cookie(None);
apply_guard(true, &mut headers);
assert!(headers.get(SET_COOKIE).is_some());
assert_eq!(headers[CACHE_CONTROL], "private, no-cache");
let mut explicit = with_cookie(Some("no-store"));
apply_guard(true, &mut explicit);
assert_eq!(explicit[CACHE_CONTROL], "no-store");
}
#[test]
fn pages_and_assets_keep_their_headers() {
let mut headers = with_cookie(None);
apply_guard(false, &mut headers);
assert!(headers.get(CACHE_CONTROL).is_none());
}
#[test]
fn prefixes_match_whole_segments() {
assert!(under_prefix("/api/v1/trending", "/api"));
assert!(under_prefix("/api", "/api/"));
assert!(!under_prefix("/apiary", "/api"));
assert!(!under_prefix("/assets/app.js", "/api"));
}
#[test]
fn only_successful_gets_are_marked() {
let policy = cdn_cache_for(300);
let mut ok = HeaderMap::new();
mark_shared(policy, true, StatusCode::OK, &mut ok);
assert_eq!(ok[CACHE_CONTROL], "public, max-age=60, s-maxage=300");
let mut error = HeaderMap::new();
mark_shared(policy, true, StatusCode::INTERNAL_SERVER_ERROR, &mut error);
assert!(error.get(CACHE_CONTROL).is_none());
let mut post = HeaderMap::new();
mark_shared(policy, false, StatusCode::OK, &mut post);
assert!(post.get(CACHE_CONTROL).is_none());
}
#[test]
fn the_browser_never_keeps_a_copy_longer_than_the_cdn() {
assert_eq!(
cdn_cache_for(30).header(),
"public, max-age=30, s-maxage=30"
);
assert_eq!(
cdn_cache_for(300).browser(600).header(),
"public, max-age=300, s-maxage=300"
);
}
}