diff --git a/src/invocation-plane-services/http-invocation/crates/server/src/app.rs b/src/invocation-plane-services/http-invocation/crates/server/src/app.rs index 034bb02ebd..7883fdfe07 100644 --- a/src/invocation-plane-services/http-invocation/crates/server/src/app.rs +++ b/src/invocation-plane-services/http-invocation/crates/server/src/app.rs @@ -166,6 +166,8 @@ pub async fn app( }, )) .layer(trace_layer.clone()) + // outermost so it also covers responses from the timeout and panic layers + .layer(axum_mw::from_fn(routes::error_source_middleware)) .with_state(app_state.clone()); // using a full router instead of a handler function so the metrics layer can be attached @@ -192,6 +194,7 @@ pub async fn app( }, )) .layer(trace_layer) + .layer(axum_mw::from_fn(routes::error_source_middleware)) .with_state(app_state); let host_router = Router::new() diff --git a/src/invocation-plane-services/http-invocation/crates/server/src/routes/error_source.rs b/src/invocation-plane-services/http-invocation/crates/server/src/routes/error_source.rs new file mode 100644 index 0000000000..c8f7154383 --- /dev/null +++ b/src/invocation-plane-services/http-invocation/crates/server/src/routes/error_source.rs @@ -0,0 +1,131 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use axum::extract::Request; +use axum::middleware::Next; +use axum::response::Response; +use http::{HeaderName, HeaderValue}; + +/// Names the component that produced an error (4xx or 5xx) response. Clients and the +/// vanity gateway read it to tell service errors from errors relayed from a +/// worker. The service owns this header: any value on a response before this +/// layer runs, including one supplied by a worker, is discarded. +pub(crate) const NVCF_ERROR_SOURCE: HeaderName = HeaderName::from_static("nvcf-error-source"); + +const INVOCATION_SERVICE: HeaderValue = HeaderValue::from_static("invocation_service"); +const WORKER: HeaderValue = HeaderValue::from_static("worker"); + +/// Response extension set by handlers that relay a response produced by a worker. +/// Every other response is treated as generated by the invocation service. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct WorkerResponse; + +/// Replaces `NVCF-Error-Source` on every response and sets it on 4xx and 5xx +/// responses to `worker` or `invocation_service`. +pub(crate) async fn error_source_middleware(request: Request, next: Next) -> Response { + let mut response = next.run(request).await; + // `remove` drops every value stored under the name, so repeated + // worker-supplied copies are discarded too. + response.headers_mut().remove(NVCF_ERROR_SOURCE); + if response.status().is_client_error() || response.status().is_server_error() { + let source = if response.extensions().get::().is_some() { + WORKER + } else { + INVOCATION_SERVICE + }; + response.headers_mut().insert(NVCF_ERROR_SOURCE, source); + } + response +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::Body; + use axum::middleware::from_fn; + use axum::routing::get; + use axum::Router; + use http::{Request as HttpRequest, StatusCode}; + use tower::ServiceExt; + + fn spoofed_headers(status: StatusCode, worker: bool) -> Response { + let mut response = Response::new(Body::empty()); + *response.status_mut() = status; + response + .headers_mut() + .append(NVCF_ERROR_SOURCE, HeaderValue::from_static("spoofed")); + response + .headers_mut() + .append(NVCF_ERROR_SOURCE, HeaderValue::from_static("worker")); + if worker { + response.extensions_mut().insert(WorkerResponse); + } + response + } + + async fn serve(status: StatusCode, worker: bool) -> Response { + let app = Router::new() + .route( + "/", + get(move || async move { spoofed_headers(status, worker) }), + ) + .layer(from_fn(error_source_middleware)); + app.oneshot(HttpRequest::get("/").body(Body::empty()).unwrap()) + .await + .unwrap() + } + + fn values(response: &Response) -> Vec<&str> { + response + .headers() + .get_all(NVCF_ERROR_SOURCE) + .iter() + .map(|value| value.to_str().unwrap()) + .collect() + } + + #[tokio::test] + async fn worker_error_is_stamped_worker_and_spoofed_copies_are_removed() { + for status in [StatusCode::BAD_REQUEST, StatusCode::INTERNAL_SERVER_ERROR] { + let response = serve(status, true).await; + assert_eq!(values(&response), vec!["worker"], "{status}"); + } + } + + #[tokio::test] + async fn service_error_is_stamped_invocation_service_even_with_spoofed_copies() { + for status in [ + StatusCode::NOT_FOUND, + StatusCode::TOO_MANY_REQUESTS, + StatusCode::GATEWAY_TIMEOUT, + ] { + let response = serve(status, false).await; + assert_eq!(values(&response), vec!["invocation_service"], "{status}"); + } + } + + #[tokio::test] + async fn non_error_responses_carry_no_error_source() { + for (status, worker) in [ + (StatusCode::OK, true), + (StatusCode::ACCEPTED, true), + (StatusCode::FOUND, true), + (StatusCode::OK, false), + ] { + let response = serve(status, worker).await; + assert!(values(&response).is_empty(), "{status} worker={worker}"); + } + } +} diff --git a/src/invocation-plane-services/http-invocation/crates/server/src/routes/get_exec.rs b/src/invocation-plane-services/http-invocation/crates/server/src/routes/get_exec.rs index b4794c0de6..78b9404029 100644 --- a/src/invocation-plane-services/http-invocation/crates/server/src/routes/get_exec.rs +++ b/src/invocation-plane-services/http-invocation/crates/server/src/routes/get_exec.rs @@ -18,6 +18,7 @@ use crate::nvcf_api::NVCFService; use crate::routes::app_error::AppError; use crate::routes::get_pexec::pexec_status; use crate::routes::post_exec::InvokeFunctionResponse; +use crate::routes::WorkerResponse; use crate::settings::AppConfig; use crate::worker_streams::WorkerStreamService; use axum::extract::{Extension, Path, State}; @@ -35,7 +36,14 @@ pub async fn exec_status( State(worker_stream_service): State>, State(app_config): State>, Extension(Authorization(bearer)): Extension>, -) -> Result<(StatusCode, Json), AppError> { +) -> Result< + ( + StatusCode, + Option>, + Json, + ), + AppError, +> { let response = pexec_status( request_id, nvcf_service, diff --git a/src/invocation-plane-services/http-invocation/crates/server/src/routes/mod.rs b/src/invocation-plane-services/http-invocation/crates/server/src/routes/mod.rs index 4062f80f73..cb741713cd 100644 --- a/src/invocation-plane-services/http-invocation/crates/server/src/routes/mod.rs +++ b/src/invocation-plane-services/http-invocation/crates/server/src/routes/mod.rs @@ -16,6 +16,7 @@ mod app_error; mod attach; mod body_stream; +mod error_source; mod get_exec; mod get_pexec; mod health; @@ -29,6 +30,8 @@ mod tlb; pub use attach::get_attach::request_attach; pub use attach::post_attach::response_attach; +pub(crate) use error_source::error_source_middleware; +pub use error_source::WorkerResponse; pub use get_exec::exec_status; pub use get_pexec::pexec_status_route; pub use health::get_health; diff --git a/src/invocation-plane-services/http-invocation/crates/server/src/routes/post_exec.rs b/src/invocation-plane-services/http-invocation/crates/server/src/routes/post_exec.rs index 26810a03cb..7ec3b5376d 100644 --- a/src/invocation-plane-services/http-invocation/crates/server/src/routes/post_exec.rs +++ b/src/invocation-plane-services/http-invocation/crates/server/src/routes/post_exec.rs @@ -21,6 +21,7 @@ use crate::routes::app_error::AppError; use crate::routes::input_asset_header::InputAssetHeader; use crate::routes::pexec; use crate::routes::post_pexec::{FunctionRouting, NVCF_POLL_SECONDS}; +use crate::routes::WorkerResponse; use crate::s3::S3Service; use crate::settings::AppConfig; use crate::worker_streams::WorkerStreamService; @@ -94,7 +95,14 @@ pub async fn exec( State(app_config): State>, Extension(Authorization(bearer)): Extension>, request: Request, -) -> Result<(StatusCode, Json), AppError> { +) -> Result< + ( + StatusCode, + Option>, + Json, + ), + AppError, +> { let function_id = function_routing.function_id; let function_version_id = function_routing.function_version_id; @@ -160,7 +168,14 @@ pub async fn exec( #[tracing::instrument(level = Level::TRACE, fields(request_id))] pub async fn map_passthrough_response_to_wrapped_response( response: Response, -) -> Result<(StatusCode, Json), AppError> { +) -> Result< + ( + StatusCode, + Option>, + Json, + ), + AppError, +> { let (parts, body) = response.into_parts(); let req_id = parts .headers @@ -170,11 +185,18 @@ pub async fn map_passthrough_response_to_wrapped_response( .unwrap_or_default(); let span = tracing::Span::current(); span.record("request_id", format!("{}", req_id)); + // The wrapper reports the worker's status, so it keeps the worker origin. + let worker_response = parts + .extensions + .get::() + .copied() + .map(Extension); Ok(( match parts.status { StatusCode::FOUND => StatusCode::OK, _ => parts.status, }, + worker_response, Json(InvokeFunctionResponse { req_id, status: match parts.status { diff --git a/src/invocation-plane-services/http-invocation/crates/server/src/routes/post_pexec.rs b/src/invocation-plane-services/http-invocation/crates/server/src/routes/post_pexec.rs index 63b2226259..8590dd48df 100644 --- a/src/invocation-plane-services/http-invocation/crates/server/src/routes/post_pexec.rs +++ b/src/invocation-plane-services/http-invocation/crates/server/src/routes/post_pexec.rs @@ -37,6 +37,7 @@ use crate::routes::{ input_asset_header::InputAssetHeader, nvcf_status_header::NVCFStatusHeader, tlb::FunctionId, + WorkerResponse, }; use crate::s3::{Error, S3Service}; use crate::settings::AppConfig; @@ -506,7 +507,10 @@ pub(crate) async fn handle_streaming_response( let mut response = select! { biased; response = attach_response_rx => { - response? + let mut response = response?; + // marks the status as worker-produced for the error source layer + response.extensions_mut().insert(WorkerResponse); + response } _ = tokio::time::sleep(poll_duration + Duration::from_secs(2)) => { gateway_timeout_response(function_id, request_id) diff --git a/src/invocation-plane-services/http-invocation/crates/server/tests/mocks/nvcf_worker_mock.rs b/src/invocation-plane-services/http-invocation/crates/server/tests/mocks/nvcf_worker_mock.rs index 0d3ded50c0..eb3ddb11d7 100644 --- a/src/invocation-plane-services/http-invocation/crates/server/tests/mocks/nvcf_worker_mock.rs +++ b/src/invocation-plane-services/http-invocation/crates/server/tests/mocks/nvcf_worker_mock.rs @@ -386,6 +386,7 @@ impl WorkHandler for DefaultWorkHandler { pub struct CustomHeadersWorkHandler { headers: http::HeaderMap, + status: StatusCode, } #[async_trait] @@ -396,7 +397,7 @@ impl WorkHandler for CustomHeadersWorkHandler { request: &WorkerInvokeFunctionRequest, ) -> anyhow::Result<()> { let mut builder = http::Response::builder() - .status(StatusCode::OK) + .status(self.status) .header(CONTENT_TYPE, "text/event-stream"); for header in &self.headers { builder = builder.header(header.0.clone(), header.1.clone()); @@ -410,7 +411,15 @@ impl WorkHandler for CustomHeadersWorkHandler { impl CustomHeadersWorkHandler { pub fn new(headers: http::HeaderMap) -> Self { - Self { headers } + Self { + headers, + status: StatusCode::OK, + } + } + + pub fn with_status(mut self, status: StatusCode) -> Self { + self.status = status; + self } } diff --git a/src/invocation-plane-services/http-invocation/crates/server/tests/test_response_headers.rs b/src/invocation-plane-services/http-invocation/crates/server/tests/test_response_headers.rs index 30d9c9a2c4..586ca76ff0 100644 --- a/src/invocation-plane-services/http-invocation/crates/server/tests/test_response_headers.rs +++ b/src/invocation-plane-services/http-invocation/crates/server/tests/test_response_headers.rs @@ -29,7 +29,10 @@ use tower::{Service, ServiceExt}; // Typically you'd place this in the same module or a dedicated macros module: macro_rules! hydrate_fixtures { - ($response_headers:expr) => {{ + ($response_headers:expr) => { + hydrate_fixtures!($response_headers, StatusCode::OK) + }; + ($response_headers:expr, $status:expr) => {{ // The curly braces inside the macro body let us return an expression (a block). // This block can contain multiple statements and a final expression. let (_localstack, nats, mock_nvcf_api, config) = fixtures().await; @@ -43,7 +46,7 @@ macro_rules! hydrate_fixtures { function_version_id: VERSION_ID_1, instance_id: INSTANCE_ID.into(), }, - Box::new(CustomHeadersWorkHandler::new($response_headers)), + Box::new(CustomHeadersWorkHandler::new($response_headers).with_status($status)), PublishMode::Attach(app.clone()), ) .await? @@ -182,3 +185,77 @@ async fn test_reqid_returned_for_empty_response_headers() -> anyhow::Result<()> Ok(()) } + +fn spoofed_error_source_headers() -> http::HeaderMap { + http::HeaderMap::from_iter(vec![ + ( + http::header::HeaderName::from_static("nvcf-error-source"), + HeaderValue::from_static("invocation_service"), + ), + ( + http::header::HeaderName::from_static("nvcf-error-source"), + HeaderValue::from_static("worker"), + ), + ]) +} + +fn pexec_request(api_key: &str) -> anyhow::Result> { + Ok(axum::http::Request::builder() + .method(Method::POST) + .uri(format!( + "/v2/nvcf/pexec/functions/{FUNCTION_ID}/versions/{VERSION_ID_1}" + )) + .header(http::header::AUTHORIZATION, format!("Bearer {api_key}")) + .body(Body::from("a body"))?) +} + +fn error_source_values(response: &http::Response) -> Vec { + response + .headers() + .get_all("nvcf-error-source") + .iter() + .map(|v| v.to_str().unwrap().to_string()) + .collect() +} + +#[tokio::test] +async fn test_error_source_not_set_on_success() -> anyhow::Result<()> { + // Worker-supplied copies must not leak through on a successful response. + let (mut app, _fixtures_guard) = hydrate_fixtures!(spoofed_error_source_headers()); + + let response = app.call(pexec_request(API_KEY)?).await?; + + assert_eq!(response.status(), StatusCode::OK); + assert!(error_source_values(&response).is_empty()); + + Ok(()) +} + +#[tokio::test] +async fn test_error_source_worker_for_relayed_worker_error() -> anyhow::Result<()> { + // The worker claims invocation_service and worker; only the authoritative + // value stamped by the invocation service may remain. + let (mut app, _fixtures_guard) = hydrate_fixtures!( + spoofed_error_source_headers(), + StatusCode::INTERNAL_SERVER_ERROR + ); + + let response = app.call(pexec_request(API_KEY)?).await?; + + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(error_source_values(&response), vec!["worker"]); + + Ok(()) +} + +#[tokio::test] +async fn test_error_source_invocation_service_for_service_error() -> anyhow::Result<()> { + let (mut app, _fixtures_guard) = hydrate_fixtures!(spoofed_error_source_headers()); + + let response = app.call(pexec_request("not-a-valid-key")?).await?; + + assert!(!response.status().is_success()); + assert_eq!(error_source_values(&response), vec!["invocation_service"]); + + Ok(()) +}