v0.29.2

src/proxy/error_fallback.rs

@@ -0,0 +1,155 @@
+#[derive(Clone)]
+struct InternalFallbackRoute {
+    vhost_id: String,
+    target: String,
+    depth: u8,
+}
+
+#[derive(Clone)]
+struct FinalRouteContext(crate::config::CompiledRoute);
+
+#[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;
+
+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 {
+        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: &ErrorFallbackRequest,
+    body: &ErrorFallbackBody,
+) -> Option<Response> {
+    let fallback = route
+        .error_fallback
+        .as_ref()
+        .filter(|fallback| fallback.matches(status))?;
+    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(host, &frontend.id, &fallback.target, local.0, tls)
+        .cloned()
+    else {
+        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;
+    if response.extensions().get::<FinalRouteContext>().is_none() {
+        response.extensions_mut().insert(FinalRouteContext(target));
+    }
+    Some(response)
+}