RustProxy/src/proxy/forward/success_prepare.rs

{
let origin_validation_headers = if revalidation_sent
    && response.status() == StatusCode::NOT_MODIFIED
{
    response.headers().clone()
} else { HeaderMap::new() };
if let Some(exchange) = &exchange {
    exchange
        .upstream_status
        .store(response.status().as_u16(), Ordering::Relaxed);
}
if let Some(lease) = &active_lease {
    lease.response(
        response.status().as_u16(),
        final_attempt_elapsed,
        request_body_empty,
    );
    if retries > 0 && response.status().as_u16() < 500 {
        lease.retry_succeeded();
    }
}
if cache_settings.allows_stale_error_status(response.status()) {
    if let Some(stale) = serve_stale(
        "cache_stale_if_error",
        upstream_started.elapsed(),
        &active_backend.id,
        Some(response.status()),
    ) {
        return stale;
    }
}
let h2_direct_tunnel = websocket_upstream_h2 && response.status().is_success();
if websocket_upstream_h2 && response.status() == StatusCode::SWITCHING_PROTOCOLS {
    if let Some(session) = h2_direct_session.take() {
        session.cancellation().cancel();
    }
    if let Some(exchange) = &exchange {
        exchange.websocket_error("upstream_rejected", "unexpected_h2_101");
    }
    return error(
        StatusCode::BAD_GATEWAY,
        "RustProxy: HTTP/2 WebSocket upstream returned HTTP/1.1 switching protocols\n",
    );
}
if h2_direct_tunnel {
    let Some(expected) = &handshake else {
        unreachable!("direct H2 WebSocket requires a handshake")
    };
    if websocket::validate_h2_response(response.headers(), expected).is_err() {
        if let Some(session) = h2_direct_session.take() {
            session.cancellation().cancel();
        }
        if let Some(exchange) = &exchange {
            exchange.websocket_error("upstream_rejected", "invalid_h2_handshake");
        }
        return error(
            StatusCode::BAD_GATEWAY,
            "RustProxy: invalid upstream HTTP/2 WebSocket handshake\n",
        );
    }
    if let Some(session) = h2_direct_session.as_ref() {
        session.establish();
    }
    if let Some(exchange) = &exchange {
        exchange.websocket_phase("upstream_2xx_received");
    }
} else if websocket_upstream_h2 {
    if let Some(session) = h2_direct_session.take() {
        session.cancellation().cancel();
    }
}
let mut upgraded = if response.status() == StatusCode::SWITCHING_PROTOCOLS {
    let Some(expected) = &handshake else {
        return error(
            StatusCode::BAD_GATEWAY,
            "RustProxy: unsolicited upstream upgrade\n",
        );
    };
    if websocket::validate_response(
        response.headers(),
        expected,
        websocket_legacy_support,
    )
    .is_err()
    {
        if let Some(exchange) = &exchange {
            exchange.websocket_error("upstream_rejected", "invalid_101_handshake");
        }
        return error(
            StatusCode::BAD_GATEWAY,
            "RustProxy: invalid upstream WebSocket handshake\n",
        );
    }
    if let Some(exchange) = &exchange {
        exchange.websocket_phase("upstream_101_received");
    }
    Some(hyper::upgrade::on(&mut response))
} else {
    None
};
if handshake.as_ref().is_some_and(|handshake| {
    handshake.h3() || (handshake.h2() && !websocket_upstream_h2)
})
    && response.status().is_success()
    && response.status() != StatusCode::SWITCHING_PROTOCOLS
{
    if let Some(exchange) = &exchange {
        exchange.websocket_error("upstream_rejected", "upstream_2xx_without_101");
    }
    return error(
        StatusCode::BAD_GATEWAY,
        "RustProxy: HTTP/1.1 WebSocket upstream returned 2xx without switching protocols\n",
    );
}
let origin_plan = if cache_ticket.is_some()
    && !(ignore_static_cookie && crate::cache::varies_on(response.headers(), "cookie"))
{
    crate::cache::plan(
        response.status(),
        response.headers(),
        cache_settings,
        upstream_started.elapsed(),
    ).map(|plan| {
        let plan = if force_vary_cookie { plan.with_vary("cookie") } else { plan };
        if force_vary_authorization { plan.with_vary("authorization") } else { plan }
    })
} else {
    None
};
let origin_reason = (cache_ticket.is_some() && origin_plan.is_none()).then(|| {
    if ignore_static_cookie && crate::cache::varies_on(response.headers(), "cookie") {
        "response_vary_cookie"
    } else {
        crate::cache::response_reason(response.status(), response.headers(), cache_settings)
    }
});
let origin_no_transform = crate::transform::no_transform(response.headers());
let skip_compression = crate::compression::no_transform(response.headers());
let origin_validators = crate::transform::OriginValidators {
    etag: response.headers().get("etag").cloned(),
    last_modified: response.headers().get("last-modified").cloned(),
};
let (mut parts, body) = response.into_parts();
parts.extensions.insert(origin_validators);
if skip_compression || range_requested {
    parts.extensions.insert(crate::compression::Skip);
}
headers::strip_hop_by_hop(&mut parts.headers);
if websocket_legacy_support && handshake.is_some() {
    parts.headers.remove("content-length");
    parts.headers.remove("transfer-encoding");
}
apply_response_rules(
    &mut parts.headers,
    &snapshot,
    route.frontend.as_deref(),
    Some(route),
    admin_control_plane,
);
if let Some(pool) = &route.pool {
    pool.strip_affinity_response(&mut parts.headers);
}
if let Some(cookie) = active_affinity_cookie {
    parts.headers.append("set-cookie", cookie);
}
let h3_tunnel = handshake.as_ref().is_some_and(|handshake| handshake.h3())
    && upgraded.is_some();
let h2_bridge_tunnel = handshake.as_ref().is_some_and(|handshake| handshake.h2())
    && !websocket_upstream_h2
    && upgraded.is_some();
let mut response = if h3_tunnel {
    let Some(upstream) = upgraded.take() else {
        unreachable!("checked above")
    };
    let Some(exchange) = exchange.clone() else {
        return error(
            StatusCode::INTERNAL_SERVER_ERROR,
            "RustProxy: WebSocket exchange state unavailable\n",
        );
    };
    drop(body);
    let stop_token = stop
        .as_ref()
        .map(|value| value.0.clone())
        .unwrap_or_else(|| state.shutdown.clone());
    let upstream = tokio::select! {
        _ = stop_token.cancelled() => {
            exchange.failed.store(true, Ordering::Relaxed);
            exchange.websocket_error("upgrade_cancelled", "shutdown");
            state.metrics.websocket_errors.fetch_add(1, Ordering::Relaxed);
            return error(
                StatusCode::SERVICE_UNAVAILABLE,
                "RustProxy: WebSocket upgrade cancelled\n",
            );
        }
        result = tokio::time::timeout(Duration::from_secs(10), upstream) => match result {
            Ok(Ok(upstream)) => upstream,
            _ => {
                exchange.failed.store(true, Ordering::Relaxed);
                exchange.websocket_error("upgrade_failed", "backend_on_upgrade");
                state.metrics.websocket_errors.fetch_add(1, Ordering::Relaxed);
                return error(
                    StatusCode::BAD_GATEWAY,
                    "RustProxy: upstream WebSocket upgrade failed\n",
                );
            }
        }
    };
    exchange.websocket_phase("tunnel_ready");
    parts.status = StatusCode::OK;
    parts.headers.remove("sec-websocket-accept");
    parts.headers.remove("content-length");
    parts.headers.remove("transfer-encoding");
    parts.extensions.insert(crate::compression::Skip);
    parts
        .extensions
        .insert(websocket::H3Tunnel::new(
            upstream,
            exchange,
            active_backend.origin.clone(),
            lease,
            stop,
        ));
    Response::from_parts(parts, Body::empty())
} else if h2_bridge_tunnel {
    let Some(upstream) = upgraded.take() else {
        unreachable!("checked above")
    };
    let Some(exchange) = exchange.clone() else {
        return error(
            StatusCode::INTERNAL_SERVER_ERROR,
            "RustProxy: WebSocket exchange state unavailable\n",
        );
    };
    let Some(downstream_body) = h2_downstream.take() else {
        return error(
            StatusCode::INTERNAL_SERVER_ERROR,
            "RustProxy: HTTP/2 WebSocket request stream unavailable\n",
        );
    };
    drop(body);
    let stop_token = stop
        .as_ref()
        .map(|value| value.0.clone())
        .unwrap_or_else(|| state.shutdown.clone());
    let upstream = tokio::select! {
        _ = stop_token.cancelled() => {
            exchange.failed.store(true, Ordering::Relaxed);
            exchange.websocket_error("upgrade_cancelled", "shutdown");
            state.metrics.websocket_errors.fetch_add(1, Ordering::Relaxed);
            return error(
                StatusCode::SERVICE_UNAVAILABLE,
                "RustProxy: WebSocket upgrade cancelled\n",
            );
        }
        result = tokio::time::timeout(Duration::from_secs(10), upstream) => match result {
            Ok(Ok(upstream)) => upstream,
            _ => {
                exchange.failed.store(true, Ordering::Relaxed);
                exchange.websocket_error("upgrade_failed", "backend_on_upgrade");
                state.metrics.websocket_errors.fetch_add(1, Ordering::Relaxed);
                return error(
                    StatusCode::BAD_GATEWAY,
                    "RustProxy: upstream WebSocket upgrade failed\n",
                );
            }
        }
    };
    exchange.websocket_phase("tunnel_ready");
    parts.status = StatusCode::OK;
    for name in [
        "connection",
        "upgrade",
        "sec-websocket-accept",
        "content-length",
        "transfer-encoding",
    ] {
        parts.headers.remove(name);
    }
    parts.extensions.insert(crate::compression::Skip);
    let body = websocket::bridge_h2_to_h1(
        state.clone(),
        downstream_body,
        upstream,
        exchange,
        active_backend.origin.clone(),
        lease,
        stop,
    );
    Response::from_parts(parts, body)
} else if h2_direct_tunnel {
    let Some(exchange) = exchange.clone() else {
        return error(
            StatusCode::INTERNAL_SERVER_ERROR,
            "RustProxy: WebSocket exchange state unavailable\n",
        );
    };
    let Some(session) = h2_direct_session.take() else {
        return error(
            StatusCode::INTERNAL_SERVER_ERROR,
            "RustProxy: HTTP/2 WebSocket tunnel state unavailable\n",
        );
    };
    parts.headers.remove("content-length");
    parts.headers.remove("transfer-encoding");
    parts.extensions.insert(crate::compression::Skip);
    Response::from_parts(
        parts,
        websocket::wrap_h2_response_body(body, session, exchange),
    )
} else {
    if let (Some(upstream), Some(downstream), Some(exchange)) =
        (upgraded.take(), downstream, exchange.clone())
    {
        parts
            .headers
            .insert("connection", HeaderValue::from_static("Upgrade"));
        parts
            .headers
            .insert("upgrade", HeaderValue::from_static("websocket"));
        websocket::spawn(
            state.clone(),
            downstream,
            upstream,
            exchange,
            active_backend.origin.clone(),
            lease,
            stop,
        );
    }
    Response::from_parts(parts, Body::new(body))
};
let intercept_origin_error = !route.internal && route.frontend.as_ref().is_some_and(|frontend| {
    frontend.pages.intercept_statuses.contains(&response.status().as_u16())
});
response = intercept_route_origin_error(response, route, &snapshot);
if intercept_origin_error {
    apply_response_rules(
        response.headers_mut(),
        &snapshot,
        route.frontend.as_deref(),
        Some(route),
        admin_control_plane,
    );
    response.headers_mut().insert(
        "cache-control",
        HeaderValue::from_static("no-store"),
    );
}
response.extensions_mut().insert(ResponseRulesApplied);
response
    .extensions_mut()
    .insert(OriginTiming(upstream_started.elapsed()));
response
    .extensions_mut()
    .insert(OriginBackend(active_backend.id.clone()));
if !range_requested && handshake.is_none() {
    if let Some(settings) = transform_settings.filter(|settings| settings.enabled) {
    let transform_started = Instant::now();
    response = match crate::transform::apply(
        response,
        &public_suffix,
        settings,
        origin_no_transform,
    )
    .await
    {
        Ok(response) => {
            state.transform_metrics.record_response(
                scope,
                &response,
                transform_started.elapsed(),
            );
            response
        }
        Err(cause) => {
            state
                .transform_metrics
                .record_error(scope, transform_started.elapsed());
            state.cache.record_reason(scope, &host, "transform_error");
            if let Some(exchange) = &exchange {
                *exchange
                    .cache_reason
                    .lock()
                    .unwrap_or_else(|e| e.into_inner()) = "transform_error";
            }
            tracing::debug!(error = %cause, "static response transformation failed");
            return error(
                StatusCode::BAD_GATEWAY,
                "RustProxy: upstream response body could not be transformed\n",
            );
        }
    };
    }
}
    (origin_validation_headers, origin_plan, origin_reason, response)
}