RustProxy/src/proxy/cache_helpers.rs

fn schedule_proactive_refresh(
    context: CacheRefreshContext,
    ticket: crate::cache::Ticket,
    remaining_ttl: u64,
) {
    let Some(schedule_guard) = context
        .state
        .cache
        .begin_proactive_schedule(&ticket, &context.cache_settings)
    else {
        return;
    };
    let threshold = context
        .cache_settings
        .proactive_refresh_before_expiry_seconds;
    tokio::spawn(async move {
        let _schedule_guard = schedule_guard;
        let mut delay = remaining_ttl.saturating_sub(threshold).max(1);
        let mut awaiting_initial_publish = true;
        let mut initial_publish_retries = 0u8;
        loop {
            if delay > 0 {
                tokio::time::sleep(Duration::from_secs(delay)).await;
            }
            if !context.state.cache.ticket_current(&ticket) {
                return;
            }
            let Some(cached) = context.state.cache.revalidation_tiered(&ticket).await else {
                if awaiting_initial_publish && initial_publish_retries < 60 {
                    initial_publish_retries = initial_publish_retries.saturating_add(1);
                    delay = 1;
                    continue;
                }
                return;
            };
            awaiting_initial_publish = false;
            let Some(inflight_guard) = context
                .state
                .cache
                .try_begin_proactive_refresh(&context.cache_settings)
            else {
                delay = 5;
                continue;
            };
            let Some(revalidation_guard) =
                context.state.cache.begin_background_revalidation(&cached)
            else {
                drop(inflight_guard);
                delay = 1;
                continue;
            };
            let refreshed = refresh_cache_once(&context, ticket.clone(), cached).await;
            drop(revalidation_guard);
            drop(inflight_guard);
            if !refreshed {
                return;
            }
            context.state.cache.record(
                &context.scope,
                &context.host,
                "HIT",
                "cache_proactive_refreshed",
            );
            let Some(next) = context.state.cache.revalidation_tiered(&ticket).await else {
                return;
            };
            let remaining = next.remaining_ttl();
            if remaining == 0 {
                return;
            }
            delay = remaining.saturating_sub(threshold).max(1);
        }
    });
}

fn enforce_local_page_no_cache(response: &mut Response) {
    for name in [
        "age",
        "expires",
        "pragma",
        "etag",
        "last-modified",
        "cdn-cache-control",
        "surrogate-control",
        "x-cache",
        "x-cache-hits",
        "x-cache-status",
        "x-cache-resident-age",
        "x-cache-reason",
        "cache-status",
    ] {
        response.headers_mut().remove(name);
    }
    response
        .headers_mut()
        .insert("cache-control", HeaderValue::from_static("no-store"));
}

fn error(status: StatusCode, detail: &'static str) -> Response {
    tracing::debug!(
        status = status.as_u16(),
        detail = detail.trim_end(),
        "proxy generated error"
    );
    let mut response = if let Some(html) = crate::pages::builtin_gateway_error(status) {
        let mut response = (status, html).into_response();
        response.headers_mut().insert(
            "content-type",
            HeaderValue::from_static("text/html; charset=utf-8"),
        );
        response.headers_mut().insert(
            "x-content-type-options",
            HeaderValue::from_static("nosniff"),
        );
        response
    } else {
        let body = format!(
            "{} {}\n",
            status.as_u16(),
            status.canonical_reason().unwrap_or("Error")
        );
        (status, body).into_response()
    };
    response
        .headers_mut()
        .insert("cache-control", HeaderValue::from_static("no-store"));
    response.headers_mut().insert(
        "x-content-type-options",
        HeaderValue::from_static("nosniff"),
    );
    response.extensions_mut().insert(LocalGeneratedResponse);
    response
}

fn method_policy_denial(policy: &crate::vhost::MethodPolicy, method: &Method) -> Option<Response> {
    if policy.allows(method) {
        return None;
    }
    let mut response = error(
        StatusCode::METHOD_NOT_ALLOWED,
        "RustProxy: HTTP method denied by access policy\n",
    );
    if let Some(value) = policy.allow_header() {
        response.headers_mut().insert(http::header::ALLOW, value);
    }
    Some(response)
}

fn local_location_response(spec: &crate::vhost::LocationResponse, method: &Method) -> Response {
    let status = StatusCode::from_u16(spec.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
    let body_allowed = !matches!(status.as_u16(), 100..=199 | 204 | 304) && *method != Method::HEAD;
    let mut response = Response::builder()
        .status(status)
        .body(if body_allowed {
            Body::from(spec.body.clone())
        } else {
            Body::empty()
        })
        .unwrap_or_else(|_| (StatusCode::INTERNAL_SERVER_ERROR, Body::empty()).into_response());
    if let Ok(value) = HeaderValue::from_str(&spec.content_type) {
        response.headers_mut().insert("content-type", value);
    }
    response
}

fn encoded_origin_path(path: &str) -> Option<String> {
    let mut url = url::Url::parse("http://rustproxy.invalid/").ok()?;
    url.set_path(path);
    Some(url.path().to_owned())
}
fn response_context(snapshot: &crate::config::Snapshot, request: &Request) -> SelectedRouteContext {
    let tls = request
        .extensions()
        .get::<Transport>()
        .copied()
        .unwrap_or_default()
        .tls;
    let selected = (|| {
        let local = request.extensions().get::<LocalAddr>()?.0;
        let original_host = effective_host(request, tls).ok()?;
        let host = request_host(original_host.to_str().ok()?).ok()?;
        let frontend = snapshot.frontend(&host, local, tls).cloned();
        let path = crate::acl::routing_path_view(request.uri().path()).ok();
        let route = path.as_ref().and_then(|path| {
            snapshot
                .select_on_arc(&host, path.as_ref(), local, tls)
                .cloned()
        });
        let decoded_path = path.and_then(|path| match path {
            std::borrow::Cow::Borrowed(_) => None,
            std::borrow::Cow::Owned(path) => Some(path),
        });
        Some(SelectedRouteContext {
            frontend,
            route,
            validated: Some(ValidatedRoutingIdentity {
                original_host,
                host,
                decoded_path,
            }),
        })
    })();
    selected.unwrap_or(SelectedRouteContext {
        frontend: None,
        route: None,
        validated: None,
    })
}

fn admin_session_passthrough(
    snapshot: &crate::config::Snapshot,
    host: &str,
    tls: bool,
    backend: &crate::config::CompiledBackend,
) -> bool {
    tls && backend.accept_local_listener_risk
        && backend.targets_admin_listener
        && snapshot
            .envelope
            .config
            .administration
            .explicitly_allowed_panel_host(host)
}

fn apply_response_rules(
    headers: &mut http::HeaderMap,
    snapshot: &crate::config::Snapshot,
    frontend: Option<&crate::config::CompiledFrontend>,
    route: Option<&crate::config::CompiledRoute>,
    preserve_admin_session: bool,
) {
    crate::headers::apply(headers, &snapshot.headers.response);
    if let Some(frontend) = frontend {
        crate::headers::apply(headers, &frontend.headers.response);
    }
    if let Some(route) = route {
        crate::headers::apply(headers, &route.headers.response);
    }
    if !preserve_admin_session {
        headers::strip_admin_set_cookie(headers);
    }
}
fn apply_response_age_rules(
    headers: &mut http::HeaderMap,
    snapshot: &crate::config::Snapshot,
    frontend: Option<&crate::config::CompiledFrontend>,
    route: Option<&crate::config::CompiledRoute>,
) {
    crate::headers::apply_wire_age(headers, &snapshot.headers.response);
    if let Some(frontend) = frontend {
        crate::headers::apply_wire_age(headers, &frontend.headers.response);
    }
    if let Some(route) = route {
        crate::headers::apply_wire_age(headers, &route.headers.response);
    }
}
fn finish_headers(
    response: &mut Response,
    snapshot: &crate::config::Snapshot,
    frontend: Option<&crate::config::CompiledFrontend>,
    local: Option<std::net::SocketAddr>,
    tls: bool,
    request_id: &str,
    active: &[crate::vhost::Listener],
) {
    if let Some(frontend) = frontend {
        frontend.security.apply(response.headers_mut(), tls);
    }
    frontend
        .map(|f| &f.server_header)
        .unwrap_or(&snapshot.envelope.config.server_header)
        .apply(response.headers_mut());
    frontend
        .map(|f| &f.diagnostic_headers)
        .unwrap_or(&snapshot.envelope.config.diagnostic_headers)
        .apply(response.headers_mut(), request_id);
    if frontend.is_some_and(|f| f.tinyauth.is_some() || f.basic_auth.is_some()) {
        response.headers_mut().insert(
            "cache-control",
            HeaderValue::from_static("private, no-store"),
        );
    }
    crate::http3::advertise(
        response.headers_mut(),
        local,
        tls,
        active,
        &snapshot.listeners,
        !frontend.is_some_and(|frontend| !frontend.http3_enabled),
    );
}
fn apply_timing_headers(
    response: &mut Response,
    frontend: Option<&crate::config::CompiledFrontend>,
    exchange: &Exchange,
) {
    if !frontend.is_some_and(|frontend| frontend.timing_headers) {
        return;
    }
    let total_ms = exchange.started.elapsed().as_secs_f64() * 1000.0;
    let mut values = vec![format!("proxy;dur={total_ms:.2}")];
    if let Some(origin) = response.extensions().get::<OriginTiming>() {
        values.push(format!("origin;dur={:.2}", origin.0.as_secs_f64() * 1000.0));
    }
    let cache = match exchange.cache_status.load(Ordering::Relaxed) {
        1 => Some("HIT"),
        2 => Some("MISS"),
        3 => Some("BYPASS"),
        _ => None,
    };
    if let Some(cache) = cache {
        values.push(format!("cache;desc=\"{cache}\""));
    }
    response.headers_mut().remove("server-timing");
    if let Ok(value) = HeaderValue::from_str(&values.join(", ")) {
        response.headers_mut().insert("server-timing", value);
    }
}

fn apply_backend_diagnostic_values(
    response: &mut Response,
    backend_header: bool,
    backend_timing_header: bool,
) {
    response.headers_mut().remove("x-rustproxy-backend");
    response.headers_mut().remove("x-rust-proxy-backend");
    response.headers_mut().remove("x-rust-proxy-backend-time");
    response.headers_mut().remove("x-rustproxy-backend-time");
    if backend_header {
        if let Some(origin) = response.extensions().get::<OriginBackend>() {
            if let Ok(value) = HeaderValue::from_str(&origin.0) {
                response.headers_mut().insert("x-rust-proxy-backend", value);
            }
        }
    }
    if backend_timing_header {
        if let Some(origin) = response.extensions().get::<OriginTiming>() {
            let value = format!("{:.2}ms", origin.0.as_secs_f64() * 1000.0);
            if let Ok(value) = HeaderValue::from_str(&value) {
                response
                    .headers_mut()
                    .insert("x-rustproxy-backend-time", value);
            }
        }
    }
}
fn apply_backend_diagnostic_headers(
    response: &mut Response,
    frontend: Option<&crate::config::CompiledFrontend>,
) {
    apply_backend_diagnostic_values(
        response,
        frontend.is_some_and(|value| value.backend_header),
        frontend.is_some_and(|value| value.backend_timing_header),
    );
}
async fn empty_for_cache(request: &mut Request) -> Result<bool, axum::Error> {
    if request.body().is_end_stream() {
        return Ok(true);
    }
    if request.body().size_hint().lower() > 0 {
        return Ok(false);
    }
    let bounded_h3_navigation = request.version() == Version::HTTP_3
        && (request.method() == Method::GET || request.method() == Method::HEAD);
    for _ in 0..8 {
        let next = if bounded_h3_navigation {
            match tokio::time::timeout(
                std::time::Duration::from_millis(35),
                request.body_mut().frame(),
            )
            .await
            {
                Ok(frame) => frame,
                Err(_) => return Ok(false),
            }
        } else {
            request.body_mut().frame().await
        };
        match next {
            None => {
                *request.body_mut() = Body::empty();
                return Ok(true);
            }
            Some(Err(error)) => return Err(error),
            Some(Ok(frame)) => {
                if frame.data_ref().is_some_and(|data| data.is_empty()) {
                    if request.body().is_end_stream() {
                        *request.body_mut() = Body::empty();
                        return Ok(true);
                    }
                    continue;
                }
                let tail = std::mem::replace(request.body_mut(), Body::empty());
                *request.body_mut() = Body::new(body::ReplayBody::new(frame, tail));
                return Ok(false);
            }
        }
    }
    Ok(false)
}

async fn normalize_h3_h1_unknown_body(
    request: &mut Request,
    upstream_version: Version,
) -> Result<(), axum::Error> {
    let method = request.method();
    if request.version() != Version::HTTP_3
        || upstream_version != Version::HTTP_11
        || method == Method::CONNECT
        || request.headers().contains_key(http::header::CONTENT_LENGTH)
        || request.body().size_hint().exact().is_some()
    {
        return Ok(());
    }

    let navigation_method = method == Method::GET || method == Method::HEAD;
    for _ in 0..8 {
        let next = if navigation_method {
            match tokio::time::timeout(
                std::time::Duration::from_millis(35),
                request.body_mut().frame(),
            )
            .await
            {
                Ok(frame) => frame,
                Err(_) => return Ok(()),
            }
        } else {
            request.body_mut().frame().await
        };
        match next {
            None => {
                *request.body_mut() = Body::empty();
                return Ok(());
            }
            Some(Err(error)) => return Err(error),
            Some(Ok(frame)) => {
                if frame.data_ref().is_some_and(|data| data.is_empty()) {
                    if request.body().is_end_stream() {
                        *request.body_mut() = Body::empty();
                        return Ok(());
                    }
                    continue;
                }
                let tail = std::mem::replace(request.body_mut(), Body::empty());
                *request.body_mut() = Body::new(body::ReplayBody::new(frame, tail));
                return Ok(());
            }
        }
    }
    Ok(())
}