use axum::body::Body;
use axum::http::{Request, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
const MAX_ETAG_BODY_BYTES: usize = 10 * 1024 * 1024;
fn strong_etag(body: &[u8]) -> String {
use sha2::{Digest, Sha256};
use std::fmt::Write as _;
let digest = Sha256::digest(body);
let mut hex = String::with_capacity(64);
for b in digest.iter() {
let _ = write!(hex, "{b:02x}");
}
format!("\"{hex}\"")
}
fn if_none_match_matches(header_value: &str, etag: &str) -> bool {
header_value
.split(',')
.map(|t| t.trim())
.map(|candidate| candidate.strip_prefix("W/").unwrap_or(candidate))
.any(|candidate| candidate == "*" || candidate == etag)
}
pub async fn etag_middleware(req: Request<Body>, next: Next) -> Response {
let is_get = req.method() == axum::http::Method::GET;
let if_none_match = req
.headers()
.get(axum::http::header::IF_NONE_MATCH)
.and_then(|v| v.to_str().ok())
.map(str::to_string);
let response = next.run(req).await;
if !is_get
|| !response.status().is_success()
|| response.headers().contains_key(axum::http::header::ETAG)
{
return response;
}
let (mut parts, body) = response.into_parts();
let oversized = parts
.headers
.get(axum::http::header::CONTENT_LENGTH)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<u64>().ok())
.is_some_and(|len| len > MAX_ETAG_BODY_BYTES as u64);
if oversized {
return Response::from_parts(parts, body);
}
let bytes = match axum::body::to_bytes(body, MAX_ETAG_BODY_BYTES).await {
Ok(b) => b,
Err(_) => {
return StatusCode::BAD_GATEWAY.into_response();
}
};
let etag = strong_etag(&bytes);
let matches = if_none_match
.as_deref()
.map(|inm| if_none_match_matches(inm, &etag))
.unwrap_or(false);
if let Ok(etag_value) = axum::http::HeaderValue::from_str(&etag) {
parts.headers.insert(axum::http::header::ETAG, etag_value);
}
if matches {
parts.status = StatusCode::NOT_MODIFIED;
parts.headers.remove(axum::http::header::CONTENT_LENGTH);
return Response::from_parts(parts, Body::empty());
}
Response::from_parts(parts, Body::from(bytes))
}
#[cfg(test)]
mod tests {
use super::*;
use tower::ServiceExt;
fn app() -> axum::Router {
axum::Router::new()
.route("/resource", axum::routing::get(|| async { "payload-v1" }))
.route("/post-only", axum::routing::post(|| async { "created" }))
.layer(axum::middleware::from_fn(etag_middleware))
}
async fn send(method: &str, uri: &str, if_none_match: Option<&str>) -> Response {
let mut builder = Request::builder().method(method).uri(uri);
if let Some(inm) = if_none_match {
builder = builder.header(axum::http::header::IF_NONE_MATCH, inm);
}
app()
.oneshot(builder.body(Body::empty()).unwrap())
.await
.unwrap()
}
#[tokio::test]
async fn get_response_carries_strong_etag() {
let res = send("GET", "/resource", None).await;
assert_eq!(res.status(), 200);
let etag = res
.headers()
.get(axum::http::header::ETAG)
.unwrap()
.to_str()
.unwrap();
assert!(
etag.starts_with('"') && etag.ends_with('"'),
"quoted: {etag}"
);
assert_eq!(etag.len(), 66, "quoted sha256 hex = 64 + 2 quotes");
}
#[tokio::test]
async fn matching_if_none_match_returns_304() {
let res = send("GET", "/resource", None).await;
let etag = res
.headers()
.get(axum::http::header::ETAG)
.unwrap()
.to_str()
.unwrap()
.to_string();
let res = send("GET", "/resource", Some(&etag)).await;
assert_eq!(res.status(), 304);
let body = axum::body::to_bytes(res.into_body(), usize::MAX)
.await
.unwrap();
assert!(body.is_empty(), "304 must elide the body");
}
#[tokio::test]
async fn star_if_none_match_returns_304() {
let res = send("GET", "/resource", Some("*")).await;
assert_eq!(res.status(), 304);
}
#[tokio::test]
async fn weak_if_none_match_candidate_matches_strong_etag() {
let res = send("GET", "/resource", None).await;
let etag = res
.headers()
.get(axum::http::header::ETAG)
.unwrap()
.to_str()
.unwrap()
.to_string();
let weak = format!("W/{etag}");
let res = send("GET", "/resource", Some(&weak)).await;
assert_eq!(res.status(), 304, "weak candidate must use weak comparison");
}
#[tokio::test]
async fn mismatched_if_none_match_returns_200_with_body() {
let res = send("GET", "/resource", Some("\"deadbeef\"")).await;
assert_eq!(res.status(), 200);
let body = axum::body::to_bytes(res.into_body(), usize::MAX)
.await
.unwrap();
assert_eq!(body, "payload-v1");
}
#[tokio::test]
async fn non_get_methods_are_not_fingerprinted() {
let res = send("POST", "/post-only", None).await;
assert_eq!(res.status(), 200);
assert!(
res.headers().get(axum::http::header::ETAG).is_none(),
"POST responses must not gain an ETag"
);
}
}