RustProxy/src/proxy/web_cache_revalidation.rs

//! Conditional validation for a local Web Server origin. The reverse-proxy
//! source has an equivalent 304 path; both must preserve the stored body and
//! its hit counter when the origin reports that the representation is unchanged.

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,
}

/// The client's validators describe the *public* representation (possibly
/// transformed). Only validators belonging to our stored origin representation
/// may be sent to the Web Server. Never allow an unrelated client ETag to turn
/// into a 304 that would incorrectly validate our cached body.
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)
}

/// A transformed body uses public representation-specific headers. A local
/// origin's 304 can describe its own encoding/ranges, which must not be copied
/// onto the cached, already transformed body.
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
}

/// Handle a real 304 from the local origin, *after* matching it to a stored
/// object using our own validator. A successful validation increases HITs and
/// refreshes freshness metadata; an uncacheable 304 is served once and purged.
/// `apply_rules` keeps vhost response rules identical to the ordinary Web
/// Server 200/fill path. Invalidated tickets are never resurrected.
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);
    // Vhost response rules must never turn a newly private/no-store origin
    // into a cacheable shared representation.
    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();
    }
    settings.annotate_reason(response.headers_mut(), "HIT", hits, ttl, reason);
    settings.annotate_transformed(response.headers_mut(), transformed);
    settings.annotate_resident_age(response.headers_mut(), resident_age);
    Some(WebRevalidated { response, transformed, reason })
}

#[cfg(test)]
mod tests;