RustProxy/src/proxy/web_cache_revalidation/mod.rs

use crate::cache::{Cache, RangeSelection, Revalidation, Settings, Ticket};
use axum::{body::Body, response::Response};
use http::{HeaderMap, HeaderValue, Method, StatusCode};
use std::time::Duration;

pub(super) struct WebRevalidated {
    pub response: Response,
    pub transformed: bool,
    pub reason: &'static str,
}

pub(super) fn web_origin_request_headers(
    client: &HeaderMap,
    cached: Option<&Revalidation>,
) -> (HeaderMap, bool) {
    let mut headers = client.clone();
    let Some(cached) = cached else {
        return (headers, false);
    };
    headers.remove("if-none-match");
    headers.remove("if-modified-since");
    let sent = cached.apply_validators(&mut headers);
    (headers, sent)
}

fn web_merged_validation_headers(cached: &Revalidation, validation: &HeaderMap) -> HeaderMap {
    let mut merged = cached.merge_validation_headers(validation);
    if cached.transformed() {
        for name in [
            "content-encoding",
            "vary",
            "accept-ranges",
            "content-md5",
            "digest",
            "content-digest",
        ] {
            merged.remove(name);
            for value in cached.headers().get_all(name) {
                merged.append(name, value.clone());
            }
        }
    }
    merged
}

pub(super) fn finish_web_revalidation(
    cache: &Cache,
    ticket: &Ticket,
    cached: &Revalidation,
    validation: &HeaderMap,
    settings: &Settings,
    elapsed: Duration,
    client: &HeaderMap,
    method: &Method,
    ignore_cookie: bool,
    force_cache: bool,
    apply_rules: impl FnOnce(&mut HeaderMap),
) -> Option<WebRevalidated> {
    let origin_merged = web_merged_validation_headers(cached, validation);
    let origin_safe =
        crate::cache::plan(StatusCode::OK, &origin_merged, settings, elapsed).is_some();
    let mut merged = origin_merged;
    apply_rules(&mut merged);
    if force_cache {
        if client.contains_key("cookie") {
            merged.append("vary", HeaderValue::from_static("Cookie"));
        }
        if client.contains_key("authorization") {
            merged.append("vary", HeaderValue::from_static("Authorization"));
        }
    }
    let can_store = origin_safe && !(ignore_cookie && crate::cache::varies_on(&merged, "cookie"));
    let client_not_modified = crate::cache::conditional_not_modified(client, &merged);
    let range = crate::cache::request_range(method, client);
    let if_range = client.get("if-range");
    let published = if can_store {
        cache.publish_revalidated(
            ticket.clone(),
            cached,
            merged.clone(),
            validation,
            settings,
            elapsed,
        )
    } else {
        None
    };
    let (mut response, hits, ttl, resident_age, transformed, reason) =
        if let Some(mut hit) = published {
            if !client_not_modified {
                if let RangeSelection::Single(selection) = range {
                    hit = hit.apply_range(selection, if_range);
                }
            }
            (
                hit.response,
                hit.hits,
                hit.remaining_ttl,
                hit.resident_age_seconds,
                hit.transformed,
                "cache_revalidated",
            )
        } else {
            if !cache.ticket_current(ticket) {
                return None;
            }
            cache.discard(ticket);
            let response = if !client_not_modified {
                if let RangeSelection::Single(selection) = range {
                    cached.response_with_range(merged, selection, if_range)
                } else {
                    cached.response(merged)
                }
            } else {
                cached.response(merged)
            };
            (
                response,
                0,
                0,
                0,
                cached.transformed(),
                "cache_revalidated_uncached",
            )
        };
    if client_not_modified {
        response = crate::cache::make_not_modified(response);
    }
    if *method == Method::HEAD {
        *response.body_mut() = Body::empty();
    }
    let request_forces_validation =
        crate::cache::request_policy(method, client, true, "", false, settings)
            .revalidation_reason()
            .is_some()
            || cached
                .request_freshness_revalidation_reason(&crate::cache::request_freshness_headers(
                    client, settings,
                ))
                .is_some();
    settings.annotate_reason_transformed_forwarded(
        response.headers_mut(),
        "HIT",
        hits,
        ttl,
        reason,
        transformed,
        cached.cache_status_forward_reason(request_forces_validation),
        Some(StatusCode::NOT_MODIFIED),
    );
    settings.annotate_resident_age(response.headers_mut(), resident_age);
    Some(WebRevalidated {
        response,
        transformed,
        reason,
    })
}

#[cfg(test)]
mod tests;