RustProxy/src/proxy/error_fallback.rs

#[derive(Clone)]
struct InternalFallbackRoute {
    vhost_id: String,
    target: String,
    depth: u8,
}

#[derive(Clone)]
struct FinalRouteContext(Arc<crate::config::CompiledRoute>);

#[derive(Clone, Copy)]
struct ErrorFallbackResponse;

#[derive(Clone)]
struct ErrorFallbackRequest {
    method: Method,
    uri: Uri,
    version: Version,
    headers: HeaderMap,
    transport: Option<Transport>,
    snapshot: Arc<crate::config::Snapshot>,
    exchange: Option<Arc<Exchange>>,
    connection_lease: Option<Arc<ConnectionLease>>,
    stop: Option<ServerShutdown>,
    depth: u8,
}

impl ErrorFallbackRequest {
    fn capture(request: &Request, snapshot: &Arc<crate::config::Snapshot>) -> Self {
        Self {
            method: request.method().clone(),
            uri: request.uri().clone(),
            version: request.version(),
            headers: request.headers().clone(),
            transport: request.extensions().get::<Transport>().copied(),
            snapshot: snapshot.clone(),
            exchange: request.extensions().get::<Arc<Exchange>>().cloned(),
            connection_lease: request.extensions().get::<Arc<ConnectionLease>>().cloned(),
            stop: request.extensions().get::<ServerShutdown>().cloned(),
            depth: request
                .extensions()
                .get::<InternalFallbackRoute>()
                .map(|dispatch| dispatch.depth)
                .unwrap_or(0),
        }
    }

    fn rebuild(
        &self,
        body: Bytes,
        peer: PeerAddr,
        local: LocalAddr,
        dispatch: InternalFallbackRoute,
    ) -> Request {
        let mut request = Request::new(Body::from(body));
        *request.method_mut() = self.method.clone();
        *request.uri_mut() = self.uri.clone();
        *request.version_mut() = self.version;
        *request.headers_mut() = self.headers.clone();
        request.extensions_mut().insert(peer);
        request.extensions_mut().insert(local);
        if let Some(transport) = self.transport {
            request.extensions_mut().insert(transport);
        }
        request.extensions_mut().insert(self.snapshot.clone());
        if let Some(exchange) = &self.exchange {
            request.extensions_mut().insert(exchange.clone());
        }
        if let Some(lease) = &self.connection_lease {
            request.extensions_mut().insert(lease.clone());
        }
        if let Some(stop) = &self.stop {
            request.extensions_mut().insert(stop.clone());
        }
        request.extensions_mut().insert(dispatch);
        request
    }
}

enum ErrorFallbackBody {
    Disabled,
    Replayable(Bytes),
    Unavailable,
}

const MAX_ERROR_FALLBACK_DEPTH: u8 = 8;
const MAX_ERROR_FALLBACK_REPLAY_BODY_BYTES: u64 = 8 * 1024 * 1024;

#[allow(clippy::result_large_err)]
async fn prepare_error_fallback_body(
    request: &mut Request,
    enabled: bool,
) -> Result<ErrorFallbackBody, Response> {
    if !enabled {
        return Ok(ErrorFallbackBody::Disabled);
    }
    if request.body().is_end_stream() {
        return Ok(ErrorFallbackBody::Replayable(Bytes::new()));
    }
    let Some(length) = request.body().size_hint().exact() else {
        if request.method() == Method::GET || request.method() == Method::HEAD {
            return match empty_for_cache(request).await {
                Ok(true) => Ok(ErrorFallbackBody::Replayable(Bytes::new())),
                Ok(false) => Ok(ErrorFallbackBody::Unavailable),
                Err(_) => Err(error(
                    StatusCode::BAD_REQUEST,
                    "RustProxy: invalid or incomplete request body\n",
                )),
            };
        }
        return Ok(ErrorFallbackBody::Unavailable);
    };
    if length > MAX_ERROR_FALLBACK_REPLAY_BODY_BYTES {
        return Ok(ErrorFallbackBody::Unavailable);
    }
    let body = std::mem::replace(request.body_mut(), Body::empty());
    let collected = body.collect().await.map_err(|_| {
        error(
            StatusCode::BAD_REQUEST,
            "RustProxy: invalid or incomplete request body\n",
        )
    })?;
    let bytes = collected.to_bytes();
    *request.body_mut() = Body::from(bytes.clone());
    Ok(ErrorFallbackBody::Replayable(bytes))
}

#[allow(clippy::too_many_arguments)]
async fn execute_error_fallback(
    state: Arc<App>,
    peer: PeerAddr,
    local: LocalAddr,
    snapshot: &Arc<crate::config::Snapshot>,
    host: &str,
    tls: bool,
    route: &crate::config::CompiledRoute,
    status: StatusCode,
    template: &Option<ErrorFallbackRequest>,
    body: &ErrorFallbackBody,
) -> Option<Response> {
    let fallback = route
        .error_fallback
        .as_ref()
        .filter(|fallback| fallback.matches(status))?;
    let template = template.as_ref()?;
    let frontend = route.frontend.as_ref()?;
    if template.depth >= MAX_ERROR_FALLBACK_DEPTH {
        tracing::warn!(vhost = %frontend.id, route = %route.id, target = %fallback.target, "error fallback depth limit reached");
        return None;
    }
    let Some(target) = snapshot
        .select_internal_on_arc(host, &frontend.id, &fallback.target, local.0, tls)
        .cloned()
    else {
        tracing::warn!(
            vhost = %frontend.id,
            route = %route.id,
            target = %fallback.target,
            status = status.as_u16(),
            "error fallback target could not be resolved on this frontend"
        );
        return None;
    };
    let ErrorFallbackBody::Replayable(body) = body else {
        tracing::warn!(vhost = %frontend.id, route = %route.id, target = %fallback.target, "error fallback request body is not replayable");
        return None;
    };
    let request = template.rebuild(
        body.clone(),
        peer,
        local,
        InternalFallbackRoute {
            vhost_id: frontend.id.clone(),
            target: fallback.target.clone(),
            depth: template.depth.saturating_add(1),
        },
    );
    let mut response = Box::pin(forward_request(state, peer, local, request)).await;
    response.extensions_mut().insert(ErrorFallbackResponse);
    if response.extensions().get::<FinalRouteContext>().is_none() {
        response.extensions_mut().insert(FinalRouteContext(target));
    }
    Some(response)
}

fn intercept_route_origin_error(
    response: Response,
    route: &crate::config::CompiledRoute,
    snapshot: &crate::config::Snapshot,
) -> Response {
    if route.internal {
        return response;
    }
    let mut response = if let Some(frontend) = route.frontend.as_ref() {
        crate::pages::intercept_application_error(
            response,
            &frontend.pages,
            &snapshot.envelope.config.pages,
        )
    } else {
        response
    };
    response
        .extensions_mut()
        .insert(crate::pages::OriginProvidedResponse);
    response
}