Simple conditional Middleware for Axum with support for If-None-Match/ETag and If-Modified-Since/Last-Modified. Either checks If-None-Match against ETag of the response of the next layer (computes from body if ETag is not present), or if If-None-Match is not present in the request checks Last-Modified header of the next layer response against th…

simple_conditional_mw.rs

use etag::EntityTag;
// Based on alextes's implementation (https://gist.github.com/alextes/4095b1fca7d58bd3100825e10e882e57)
//   which already provided support for If-None-Match/ETag
//   but with support for precomputed ETags and If-Modified-Since/Last-Modified
async fn conditional_mw<B>(request: Request<B>, next: Next<B>) -> Result<Response, StatusCode> {
    let if_none_match_header = request.headers().get(header::IF_NONE_MATCH).cloned();
    let if_modified_since_header = request.headers().get(header::IF_MODIFIED_SINCE).cloned();
    // let path = request.uri().path().to_owned();
    
    let response = next.run(request).await;
    let (mut original_parts, mut original_body) = response.into_parts();
    let (etag, new_parts, new_body): (EntityTag, Parts, BoxBody) = {
        if let Some(et) = original_parts.headers.get(header::ETAG)
                                        .and_then(|e| e.to_str().unwrap().parse::<EntityTag>().ok()) {
            // response already has an ETag, assume it is better
            (et, original_parts, original_body)
        } else {
            let bytes = {
                let mut body_bytes = vec![];
                while let Some(inner) = original_body.data().await {
                    body_bytes.extend(inner.unwrap());
                }
                body_bytes
            };
            if bytes.len() == 0 {
                // no body => no ETag to compute, just return
                return Ok(original_parts.into_response());
            } else {
                // body is present, compute ETag
                let etag = EntityTag::from_data(&bytes);
                original_parts.headers.insert(
                    header::ETAG,
                    HeaderValue::from_str(&etag.to_string()).unwrap(),
                );
                (etag, original_parts, BoxBody::new(Body::from(bytes).map_err(axum::Error::new)))
            }
        }
    };
    
    match if_none_match_header {
        Some(if_none_match) => {
            // request has If-None-Match header, return 304 if valid and matching
            match if_none_match.to_str().unwrap().parse::<EntityTag>().ok().map(|e| etag.weak_eq(&e)) {
                Some(true) => Ok((StatusCode::NOT_MODIFIED, new_parts).into_response()),
                _ => Ok((new_parts, new_body).into_response())
            }
        },
        None => {
            // request has no If-None-Match, return 304 if If-Modified-Since is present, valid and outdated
            let parse = |h: &HeaderValue| DateTime::parse_from_rfc2822(h.to_str().unwrap()).ok();
            match if_modified_since_header.and_then(|e| parse(&e)).and_then(|if_modified_since|
                  new_parts.headers.get(header::LAST_MODIFIED).and_then(|e| parse(e)).map(|last_modified| if_modified_since >= last_modified)) {
                Some(true) => Ok((StatusCode::NOT_MODIFIED, new_parts).into_response()),
                _ => Ok((new_parts, new_body).into_response())
            }
        }
    }
}