feat: add lambda-runtime-invocation-id header

darklight3it · darklight3it · commit b60f865d411a · 2026-07-29T13:36:10.000+01:00
add lambda-runtime-invocation-id in order to support cross wiring protection. Fixes #1155
diff --git a/lambda-runtime/src/constants.rs 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
@@ -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::<EventPayload>(&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,23 +138,28 @@ 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<T>(request_id: &str, err: T) -> Result<http::Request<Body>, BoxError>
+fn build_event_error_request<T>(
+    request_id: String,
+    invocation_id: Option<String>,
+    err: T,
+) -> Result<http::Request<Body>, BoxError>
 where
     T: Into<Diagnostic> + Debug,
 {
     error!(error = ?err, "Request payload deserialization into LambdaEvent<T> 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)]
 pub enum RuntimeApiResponseFuture<F, Response, BufferedResponse, StreamingResponse, StreamItem, StreamError> {
     Future(
         #[pin] F,
         String,
+        Option<String>,
         PhantomData<(
             (),
             Response,
@@ -183,9 +189,9 @@ where
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
         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
@@ -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
@@ -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<Error> + Send + Debug,
 {
     pub(crate) request_id: &'a str,
+    pub(crate) invocation_id: Option<&'a str>,
     pub(crate) body: R,
     pub(crate) _unused_b: PhantomData<B>,
     pub(crate) _unused_s: PhantomData<S>,
@@ -102,9 +106,14 @@ where
     E: Into<Error> + 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::<B>,
             _unused_s: PhantomData::<S>,
@@ -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<Diagnostic>) -> EventErrorRequest<'a> {
+    pub(crate) fn new(
+        request_id: &'a str,
+        invocation_id: Option<&'a str>,
+        diagnostic: impl Into<Diagnostic>,
+    ) -> 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
@@ -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
