From b60f865d411a00feb4cc56b10d881ff7ecc95184 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Wed, 29 Jul 2026 13:04:13 +0100 Subject: [PATCH] feat: add lambda-runtime-invocation-id header add lambda-runtime-invocation-id in order to support cross wiring protection. Fixes #1155 --- lambda-runtime/src/constants.rs | 9 ++ lambda-runtime/src/layers/api_response.rs | 20 ++-- lambda-runtime/src/lib.rs | 2 + lambda-runtime/src/requests.rs | 108 ++++++++++++++++++---- lambda-runtime/src/runtime.rs | 7 +- lambda-runtime/src/types.rs | 54 +++++++++-- 6 files changed, 166 insertions(+), 34 deletions(-) create mode 100644 lambda-runtime/src/constants.rs diff --git a/lambda-runtime/src/constants.rs b/lambda-runtime/src/constants.rs new file mode 100644 index 00000000..98c789d0 --- /dev/null +++ b/lambda-runtime/src/constants.rs @@ -0,0 +1,9 @@ +/// Header names used in the Lambda Runtime API. +pub(crate) const LAMBDA_RUNTIME_REQUEST_ID: &str = "lambda-runtime-aws-request-id"; +pub(crate) const LAMBDA_RUNTIME_DEADLINE_MS: &str = "lambda-runtime-deadline-ms"; +pub(crate) const LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN: &str = "lambda-runtime-invoked-function-arn"; +pub(crate) const LAMBDA_RUNTIME_TRACE_ID: &str = "lambda-runtime-trace-id"; +pub(crate) const LAMBDA_RUNTIME_CLIENT_CONTEXT: &str = "lambda-runtime-client-context"; +pub(crate) const LAMBDA_RUNTIME_COGNITO_IDENTITY: &str = "lambda-runtime-cognito-identity"; +pub(crate) const LAMBDA_RUNTIME_TENANT_ID: &str = "lambda-runtime-aws-tenant-id"; +pub(crate) const LAMBDA_RUNTIME_INVOCATION_ID: &str = "lambda-runtime-invocation-id"; diff --git a/lambda-runtime/src/layers/api_response.rs b/lambda-runtime/src/layers/api_response.rs index 5bb3c96f..378199e4 100644 --- a/lambda-runtime/src/layers/api_response.rs +++ b/lambda-runtime/src/layers/api_response.rs @@ -123,9 +123,10 @@ where }; let request_id = req.context.request_id.clone(); + let invocation_id = req.context.invocation_id.clone(); let lambda_event = match deserializer::deserialize::(&req.body, req.context) { Ok(lambda_event) => lambda_event, - Err(err) => match build_event_error_request(&request_id, err) { + Err(err) => match build_event_error_request(request_id, invocation_id, err) { Ok(request) => return RuntimeApiResponseFuture::Ready(Box::new(Some(Ok(request)))), Err(err) => { error!(error = ?err, "failed to build error response for Lambda Runtime API"); @@ -137,16 +138,20 @@ where // Once the handler input has been generated successfully, pass it through to inner services // allowing processing both before reaching the handler function and after the handler completes. let fut = self.inner.call(lambda_event); - RuntimeApiResponseFuture::Future(fut, request_id, PhantomData) + RuntimeApiResponseFuture::Future(fut, request_id, invocation_id, PhantomData) } } -fn build_event_error_request(request_id: &str, err: T) -> Result, BoxError> +fn build_event_error_request( + request_id: String, + invocation_id: Option, + err: T, +) -> Result, BoxError> where T: Into + Debug, { error!(error = ?err, "Request payload deserialization into LambdaEvent failed. The handler will not be called. Log at TRACE level to see the payload."); - EventErrorRequest::new(request_id, err).into_req() + EventErrorRequest::new(&request_id, invocation_id.as_deref(), err).into_req() } #[pin_project(project = RuntimeApiResponseFutureProj)] @@ -154,6 +159,7 @@ pub enum RuntimeApiResponseFuture, PhantomData<( (), Response, @@ -183,9 +189,9 @@ where fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll { task::Poll::Ready(match self.as_mut().project() { - RuntimeApiResponseFutureProj::Future(fut, request_id, _) => match ready!(fut.poll(cx)) { - Ok(ok) => EventCompletionRequest::new(request_id, ok).into_req(), - Err(err) => EventErrorRequest::new(request_id, err).into_req(), + RuntimeApiResponseFutureProj::Future(fut, request_id, invocation_id, _) => match ready!(fut.poll(cx)) { + Ok(ok) => EventCompletionRequest::new(request_id, invocation_id.as_deref(), ok).into_req(), + Err(err) => EventErrorRequest::new(request_id, invocation_id.as_deref(), err).into_req(), }, RuntimeApiResponseFutureProj::Ready(ready) => ready.take().expect("future polled after completion"), }) diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index 69c04ebc..59344194 100644 --- a/lambda-runtime/src/lib.rs +++ b/lambda-runtime/src/lib.rs @@ -23,6 +23,8 @@ pub use tower::{self, service_fn, Service}; #[macro_use] mod macros; +mod constants; + /// Diagnostic utilities to convert Rust types into Lambda Error types. pub mod diagnostic; pub use diagnostic::Diagnostic; diff --git a/lambda-runtime/src/requests.rs b/lambda-runtime/src/requests.rs index b03f14c7..16105e5b 100644 --- a/lambda-runtime/src/requests.rs +++ b/lambda-runtime/src/requests.rs @@ -1,4 +1,7 @@ -use crate::{types::ToStreamErrorTrailer, Diagnostic, Error, FunctionResponse, IntoFunctionResponse}; +use crate::{ + constants::LAMBDA_RUNTIME_INVOCATION_ID, types::ToStreamErrorTrailer, Diagnostic, Error, FunctionResponse, + IntoFunctionResponse, +}; use bytes::Bytes; use http::{header::CONTENT_TYPE, Method, Request, Uri}; use lambda_runtime_api_client::{body::Body, build_request}; @@ -88,6 +91,7 @@ where E: Into + Send + Debug, { pub(crate) request_id: &'a str, + pub(crate) invocation_id: Option<&'a str>, pub(crate) body: R, pub(crate) _unused_b: PhantomData, pub(crate) _unused_s: PhantomData, @@ -102,9 +106,14 @@ where E: Into + Send + Debug, { /// Initialize a new EventCompletionRequest - pub(crate) fn new(request_id: &'a str, body: R) -> EventCompletionRequest<'a, R, B, S, D, E> { + pub(crate) fn new( + request_id: &'a str, + invocation_id: Option<&'a str>, + body: R, + ) -> EventCompletionRequest<'a, R, B, S, D, E> { EventCompletionRequest { request_id, + invocation_id, body, _unused_b: PhantomData::, _unused_s: PhantomData::, @@ -129,7 +138,16 @@ where let body = serde_json::to_vec(&body)?; let body = Body::from(body); - let req = build_request().method(Method::POST).uri(uri).body(body)?; + let mut req = build_request() + .method(Method::POST) + .uri(uri) + .header(LAMBDA_RUNTIME_INVOCATION_ID, self.request_id) + .body(body)?; + + if let Some(id) = self.invocation_id { + req.headers_mut().insert(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); + } + Ok(req) } FunctionResponse::StreamingResponse(mut response) => { @@ -145,6 +163,11 @@ where // See the details in Lambda Developer Doc: https://docs.aws.amazon.com/lambda/latest/dg/runtimes-custom.html#runtimes-custom-response-streaming req_headers.append("Trailer", "Lambda-Runtime-Function-Error-Type".parse()?); req_headers.append("Trailer", "Lambda-Runtime-Function-Error-Body".parse()?); + + if let Some(id) = self.invocation_id { + req_headers.append(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); + } + req_headers.insert( "Content-Type", "application/vnd.awslambda.http-integration-response".parse()?, @@ -193,29 +216,22 @@ where } } -#[test] -fn test_event_completion_request() { - let req = EventCompletionRequest::new("id", "hello, world!"); - let req = req.into_req().unwrap(); - let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response"); - assert_eq!(req.method(), Method::POST); - assert_eq!(req.uri(), &expected); - assert!(match req.headers().get("User-Agent") { - Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"), - None => false, - }); -} - // /runtime/invocation/{AwsRequestId}/error pub(crate) struct EventErrorRequest<'a> { pub(crate) request_id: &'a str, + pub(crate) invocation_id: Option<&'a str>, pub(crate) diagnostic: Diagnostic, } impl<'a> EventErrorRequest<'a> { - pub(crate) fn new(request_id: &'a str, diagnostic: impl Into) -> EventErrorRequest<'a> { + pub(crate) fn new( + request_id: &'a str, + invocation_id: Option<&'a str>, + diagnostic: impl Into, + ) -> EventErrorRequest<'a> { EventErrorRequest { request_id, + invocation_id, diagnostic: diagnostic.into(), } } @@ -228,11 +244,16 @@ impl IntoRequest for EventErrorRequest<'_> { let body = serde_json::to_vec(&self.diagnostic)?; let body = Body::from(body); - let req = build_request() + let mut req = build_request() .method(Method::POST) .uri(uri) .header("lambda-runtime-function-error-type", "unhandled") .body(body)?; + + if let Some(id) = self.invocation_id { + req.headers_mut().insert(LAMBDA_RUNTIME_INVOCATION_ID, id.parse()?); + } + Ok(req) } } @@ -253,10 +274,41 @@ mod tests { }); } + #[test] + fn test_event_completion_request() { + let req = EventCompletionRequest::new("id", Option::Some("invocation_id"), "hello, world!"); + let req = req.into_req().unwrap(); + let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response"); + assert_eq!(req.method(), Method::POST); + assert_eq!(req.uri(), &expected); + + assert!(req + .headers() + .get("User-Agent") + .unwrap() + .to_str() + .unwrap() + .starts_with("aws-lambda-rust/")); + + assert_eq!( + req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).unwrap(), + "invocation_id" + ); + } + + #[test] + fn test_event_completion_request_invocation_id_not_added_when_none() { + let req = EventCompletionRequest::new("id", Option::None, "hello, world!"); + let req = req.into_req().unwrap(); + + assert!(req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none()); + } + #[test] fn test_event_error_request() { let req = EventErrorRequest { request_id: "id", + invocation_id: Option::Some("invocation_id"), diagnostic: Diagnostic { error_type: "InvalidEventDataError".into(), error_message: "Error parsing event data".into(), @@ -270,6 +322,26 @@ mod tests { Some(header) => header.to_str().unwrap().starts_with("aws-lambda-rust/"), None => false, }); + + assert!(match req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID) { + Some(header) => header.to_str().unwrap() == "invocation_id", + None => false, + }); + } + + #[test] + fn test_event_error_request_invocation_id_not_added_when_none() { + let req = EventErrorRequest { + request_id: "id", + invocation_id: None, + diagnostic: Diagnostic { + error_type: "InvalidEventDataError".into(), + error_message: "Error parsing event data".into(), + }, + }; + let req = req.into_req().unwrap(); + + assert!(req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none()); } #[test] diff --git a/lambda-runtime/src/runtime.rs b/lambda-runtime/src/runtime.rs index ae00bc20..b7335576 100644 --- a/lambda-runtime/src/runtime.rs +++ b/lambda-runtime/src/runtime.rs @@ -790,7 +790,11 @@ mod endpoint_tests { let base = server.base_url().parse().expect("Invalid mock server Uri"); let client = Client::builder().with_endpoint(base).build(); - let req = EventCompletionRequest::new("156cb537-e2d4-11e8-9b34-d36013741fb9", "{}"); + let req = EventCompletionRequest::new( + "156cb537-e2d4-11e8-9b34-d36013741fb9", + Option::Some("invocation_id"), + "{}", + ); let req = req.into_req()?; let rsp = client.call(req).await?; @@ -822,6 +826,7 @@ mod endpoint_tests { let req = EventErrorRequest { request_id: "156cb537-e2d4-11e8-9b34-d36013741fb9", + invocation_id: Option::Some("invocation_id"), diagnostic, }; let req = req.into_req()?; diff --git a/lambda-runtime/src/types.rs b/lambda-runtime/src/types.rs index 2f8b3698..20481fd9 100644 --- a/lambda-runtime/src/types.rs +++ b/lambda-runtime/src/types.rs @@ -1,4 +1,11 @@ -use crate::{Error, RefConfig}; +use crate::{ + constants::{ + LAMBDA_RUNTIME_CLIENT_CONTEXT, LAMBDA_RUNTIME_COGNITO_IDENTITY, LAMBDA_RUNTIME_DEADLINE_MS, + LAMBDA_RUNTIME_INVOCATION_ID, LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN, LAMBDA_RUNTIME_REQUEST_ID, + LAMBDA_RUNTIME_TENANT_ID, LAMBDA_RUNTIME_TRACE_ID, + }, + Error, RefConfig, +}; use base64::prelude::*; use bytes::Bytes; use http::{header::ToStrError, HeaderMap, HeaderValue, StatusCode}; @@ -85,6 +92,11 @@ pub struct Context { /// Includes information such as the function name, memory allocation, /// version, and log streams. pub env_config: RefConfig, + /// The invocation ID assigned by the Lambda runtime for cross-wiring protection. + /// Echoed back on `/response` and `/error` to allow RAPID to reject stale responses + /// from timed-out invocations. `None` when running against older RAPID versions + /// that don't send this header. + pub invocation_id: Option, } impl Default for Context { @@ -98,6 +110,7 @@ impl Default for Context { identity: None, tenant_id: None, env_config: std::sync::Arc::new(crate::Config::default()), + invocation_id: None, } } } @@ -106,7 +119,7 @@ impl Context { /// Create a new [Context] struct based on the function configuration /// and the incoming request data. pub fn new(request_id: &str, env_config: RefConfig, headers: &HeaderMap) -> Result { - let client_context: Option = if let Some(value) = headers.get("lambda-runtime-client-context") { + let client_context: Option = if let Some(value) = headers.get(LAMBDA_RUNTIME_CLIENT_CONTEXT) { let raw = value.to_str()?; if raw.is_empty() { None @@ -117,7 +130,7 @@ impl Context { None }; - let identity: Option = if let Some(value) = headers.get("lambda-runtime-cognito-identity") { + let identity: Option = if let Some(value) = headers.get(LAMBDA_RUNTIME_COGNITO_IDENTITY) { let raw = value.to_str()?; if raw.is_empty() { None @@ -131,26 +144,29 @@ impl Context { let ctx = Context { request_id: request_id.to_owned(), deadline: headers - .get("lambda-runtime-deadline-ms") + .get(LAMBDA_RUNTIME_DEADLINE_MS) .expect("missing lambda-runtime-deadline-ms header") .to_str()? .parse::()?, invoked_function_arn: headers - .get("lambda-runtime-invoked-function-arn") + .get(LAMBDA_RUNTIME_INVOKED_FUNCTION_ARN) .unwrap_or(&HeaderValue::from_static( "No header lambda-runtime-invoked-function-arn found.", )) .to_str()? .to_owned(), xray_trace_id: headers - .get("lambda-runtime-trace-id") + .get(LAMBDA_RUNTIME_TRACE_ID) .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), client_context, identity, tenant_id: headers - .get("lambda-runtime-aws-tenant-id") + .get(LAMBDA_RUNTIME_TENANT_ID) .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), env_config, + invocation_id: headers + .get(LAMBDA_RUNTIME_INVOCATION_ID) + .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()), }; Ok(ctx) @@ -165,7 +181,7 @@ impl Context { /// Extract the invocation request id from the incoming request. pub(crate) fn invoke_request_id(headers: &HeaderMap) -> Result<&str, ToStrError> { headers - .get("lambda-runtime-aws-request-id") + .get(LAMBDA_RUNTIME_REQUEST_ID) .expect("missing lambda-runtime-aws-request-id header") .to_str() } @@ -291,6 +307,8 @@ where #[cfg(test)] mod test { + use http::HeaderName; + use super::*; use crate::Config; use std::sync::Arc; @@ -535,4 +553,24 @@ mod test { let context = Context::new("id", config, &headers).unwrap(); assert_eq!(context.tenant_id, None); } + + #[test] + fn context_with_invocation_id_resolves() { + let config = Arc::new(Config::default()); + let mut headers = HeaderMap::new(); + + let context = Context::new("id", config, &headers).unwrap(); + + assert_eq!(context.invocation_id, None); + + let config = Arc::new(Config::default()); + headers.insert( + "lambda-runtime-invocation-id", + HeaderValue::from_static("invocation-123"), + ); + + let context = Context::new("id", config, &headers).unwrap(); + + assert_eq!(context.invocation_id, Some("invocation-123".to_string())); + } }