use crate::session::Session;
use axum::body::{to_bytes, Body};
use axum::extract::{Request, State};
use axum::http::{header, Method, StatusCode};
use axum::middleware::{from_fn_with_state, FromFnLayer, Next};
use axum::response::{IntoResponse, Response};
use larust_core::AppError;
use serde::{Deserialize, Serialize};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
const MAX_CACHEABLE_BODY_BYTES: usize = 2 * 1024 * 1024;
pub struct ResponseCacheState {
ttl: Duration,
}
#[derive(Serialize, Deserialize)]
struct CachedResponse {
content_type: Option<String>,
body: Vec<u8>,
}
impl CachedResponse {
fn into_response(self) -> Response {
let mut builder = Response::builder().status(StatusCode::OK);
if let Some(content_type) = &self.content_type {
builder = builder.header(header::CONTENT_TYPE, content_type);
}
builder
.body(Body::from(self.body))
.expect("a cached response with a plain header/body pair always builds")
}
}
fn cache_key_for_url(url: &str) -> String {
format!("__larust_responsecache:{url}")
}
pub async fn forget(url: &str) -> Result<(), AppError> {
larust_cache::forget(&cache_key_for_url(url)).await
}
type Extractors = (State<Arc<ResponseCacheState>>, Request);
type PerSessionExtractors = (State<Arc<ResponseCacheState>>, Session, Request);
type MiddlewareFn = fn(State<Arc<ResponseCacheState>>, Request, Next) -> MiddlewareFuture;
type PerSessionMiddlewareFn =
fn(State<Arc<ResponseCacheState>>, Session, Request, Next) -> MiddlewareFuture;
type MiddlewareFuture = Pin<Box<dyn Future<Output = Response> + Send>>;
pub fn for_minutes(minutes: u64) -> FromFnLayer<MiddlewareFn, Arc<ResponseCacheState>, Extractors> {
for_duration(Duration::from_secs(minutes * 60))
}
pub fn for_duration(
ttl: Duration,
) -> FromFnLayer<MiddlewareFn, Arc<ResponseCacheState>, Extractors> {
from_fn_with_state(
Arc::new(ResponseCacheState { ttl }),
middleware as MiddlewareFn,
)
}
pub fn for_minutes_per_session(
minutes: u64,
) -> FromFnLayer<PerSessionMiddlewareFn, Arc<ResponseCacheState>, PerSessionExtractors> {
for_duration_per_session(Duration::from_secs(minutes * 60))
}
pub fn for_duration_per_session(
ttl: Duration,
) -> FromFnLayer<PerSessionMiddlewareFn, Arc<ResponseCacheState>, PerSessionExtractors> {
from_fn_with_state(
Arc::new(ResponseCacheState { ttl }),
middleware_per_session as PerSessionMiddlewareFn,
)
}
fn middleware(
State(state): State<Arc<ResponseCacheState>>,
request: Request,
next: Next,
) -> MiddlewareFuture {
Box::pin(async move {
if request.method() != Method::GET {
return next.run(request).await;
}
let key = cache_key_for_url(&request.uri().to_string());
if let Some(hit) = lookup(&key).await {
return hit;
}
let response = next.run(request).await;
store_and_respond(&state, &key, response).await
})
}
fn middleware_per_session(
State(state): State<Arc<ResponseCacheState>>,
session: Session,
request: Request,
next: Next,
) -> MiddlewareFuture {
Box::pin(async move {
if request.method() != Method::GET {
return next.run(request).await;
}
let Some(session_id) = session.id() else {
return next.run(request).await;
};
let key = format!(
"{}:{session_id}",
cache_key_for_url(&request.uri().to_string())
);
if let Some(hit) = lookup(&key).await {
return hit;
}
let response = next.run(request).await;
store_and_respond(&state, &key, response).await
})
}
async fn lookup(key: &str) -> Option<Response> {
match larust_cache::get::<CachedResponse>(key).await {
Ok(Some(cached)) => Some(cached.into_response()),
Ok(None) => None,
Err(error) => {
tracing::warn!(%error, "responsecache lookup failed, serving live");
None
}
}
}
async fn store_and_respond(state: &ResponseCacheState, key: &str, response: Response) -> Response {
if response.status() != StatusCode::OK || response_is_private(&response) {
return response;
}
let content_type = response
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::to_string);
let (parts, body) = response.into_parts();
let bytes = match to_bytes(body, usize::MAX).await {
Ok(bytes) => bytes,
Err(error) => {
tracing::error!(%error, "failed to read response body for caching");
return (
StatusCode::INTERNAL_SERVER_ERROR,
"response body read failed",
)
.into_response();
}
};
if bytes.len() <= MAX_CACHEABLE_BODY_BYTES {
let cached = CachedResponse {
content_type,
body: bytes.to_vec(),
};
if let Err(error) = larust_cache::put(key, &cached, state.ttl).await {
tracing::warn!(%error, "responsecache store failed");
}
}
Response::from_parts(parts, Body::from(bytes))
}
fn response_is_private(response: &Response) -> bool {
if response.headers().contains_key(header::SET_COOKIE) {
return true;
}
response
.headers()
.get_all(header::CACHE_CONTROL)
.iter()
.any(|value| {
value.to_str().map_or(true, |value| {
value.split(',').any(|directive| {
matches!(
directive.trim().to_ascii_lowercase().as_str(),
"private" | "no-cache" | "no-store"
)
})
})
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn shared_cache_rejects_cookie_setting_and_private_responses() {
let cookie_response = Response::builder()
.header(header::SET_COOKIE, "session=secret; HttpOnly")
.body(Body::empty())
.unwrap();
assert!(response_is_private(&cookie_response));
for directive in ["private", "no-cache", "no-store", "public, private"] {
let response = Response::builder()
.header(header::CACHE_CONTROL, directive)
.body(Body::empty())
.unwrap();
assert!(
response_is_private(&response),
"{directive:?} must not enter the shared cache"
);
}
}
#[test]
fn shared_cache_allows_an_explicitly_public_response() {
let response = Response::builder()
.header(header::CACHE_CONTROL, "public, max-age=60")
.body(Body::empty())
.unwrap();
assert!(!response_is_private(&response));
}
}