feat: add lambda-runtime-invocation-id header

.github/workflows/dockerized-test.yml

@@ -19,20 +19,10 @@ jobs:
 
     - uses: dtolnay/rust-toolchain@stable
     - uses: Swatinem/rust-cache@v2
-    - name: Load environment variables
-      run: |
-        if [ -f .env ]; then
-          set -a
-          source .env
-          set +a
-          echo "HANDLERS_TO_BUILD=${HANDLERS_TO_BUILD}" >> $GITHUB_ENV
-          echo "OUTPUT_DIR=${OUTPUT_DIR}" >> $GITHUB_ENV
-        fi
-
     - name: Build Lambda artifacts for testing
       run: |
         mkdir -p test/dockerized/tasks
-        OUTPUT_DIR="$(pwd)/test/dockerized/tasks" make build-examples
+        HANDLERS_TO_BUILD="basic-lambda basic-sqs http-basic-lambda basic-lambda-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" ./scripts/build-examples.sh
         ls -la test/dockerized/tasks/
 
     - name: Build base test image with RIE and custom entrypoint
@@ -57,7 +47,7 @@ jobs:
     - name: Build Lambda artifacts for testing
       run: |
         mkdir -p test/dockerized/tasks
-        HANDLERS_TO_BUILD="basic-lambda-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" make build-examples
+        HANDLERS_TO_BUILD="basic-lambda-concurrent invocation-id-concurrent" OUTPUT_DIR="$(pwd)/test/dockerized/tasks" ./scripts/build-examples.sh
         ls -la test/dockerized/tasks/
 
     - name: Build base test image with RIE and custom entrypoint
@@ -68,6 +58,8 @@ jobs:
 
     - name: Run concurrent scenarios
       uses: aws/containerized-test-runner-for-aws-lambda@main
+      env:
+        CONTAINER_READY_DELAY_SECS: 5
       with:
         suiteFileArray: '[]'
         dockerImageName: 'local/test-base'

Dockerfile.rie

@@ -4,8 +4,20 @@ RUN dnf install -y gcc
 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
 ENV PATH="/root/.cargo/bin:${PATH}"
 
-ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie /usr/local/bin/aws-lambda-rie
-RUN chmod +x /usr/local/bin/aws-lambda-rie
+ARG TARGETARCH
+ENV RIE_VERSION=1.36 \
+    RIE_SHA256_AMD64=ba57f2683260127135ad5ba9bafea141f90492143cbaeb9312cde6dae8d1c08e \
+    RIE_SHA256_ARM64=7826415f278663274e279085ff96d7c9da210a30213fa72279e56e59f028ce76 \
+    RIE_PATH=/usr/local/bin/aws-lambda-rie
+
+COPY scripts/download-rie.sh /tmp/download-rie.sh
+RUN sh /tmp/download-rie.sh \
+    "${TARGETARCH}" \
+    "${RIE_VERSION}" \
+    "${RIE_SHA256_AMD64}" \
+    "${RIE_SHA256_ARM64}" \
+    "${RIE_PATH}" \
+    && rm /tmp/download-rie.sh
 
 ARG EXAMPLE=basic-lambda
 

Dockerfile.test

@@ -1,7 +1,19 @@
 FROM public.ecr.aws/lambda/provided:al2023
 
-ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie /usr/local/bin/aws-lambda-rie
-RUN chmod +x /usr/local/bin/aws-lambda-rie
+ARG TARGETARCH
+ENV RIE_VERSION=1.36 \
+    RIE_SHA256_AMD64=ba57f2683260127135ad5ba9bafea141f90492143cbaeb9312cde6dae8d1c08e \
+    RIE_SHA256_ARM64=7826415f278663274e279085ff96d7c9da210a30213fa72279e56e59f028ce76 \
+    RIE_PATH=/usr/local/bin/aws-lambda-rie
+
+COPY scripts/download-rie.sh /tmp/download-rie.sh
+RUN sh /tmp/download-rie.sh \
+    "${TARGETARCH}" \
+    "${RIE_VERSION}" \
+    "${RIE_SHA256_AMD64}" \
+    "${RIE_SHA256_ARM64}" \
+    "${RIE_PATH}" \
+    && rm /tmp/download-rie.sh
 
 COPY scripts/custom-lambda-entrypoint.sh /usr/local/bin/lambda-entrypoint
 RUN chmod +x /usr/local/bin/lambda-entrypoint

Makefile

@@ -7,13 +7,10 @@ INTEG_EXTENSIONS := extension-fn extension-trait logs-trait
 INTEG_ARCH := x86_64-unknown-linux-musl
 RIE_MAX_CONCURRENCY ?= 4
 TEST_RUNNER_BRANCH ?= main
+CONTAINER_READY_DELAY_SECS ?= 5
 OUTPUT_DIR ?= test/dockerized/tasks
-HANDLERS_TO_BUILD ?=
-HANDLER ?=
-
-# Load environment variables from .env file if it exists
--include .env
-export
+HANDLERS_TO_BUILD ?= basic-lambda basic-sqs http-basic-lambda basic-lambda-concurrent
+HANDLER ?= basic-lambda
 
 .PHONY: help pr-check integration-tests check-event-features fmt build-examples build-test-runner test-rie test-rie-lmi nuke test-dockerized test-dockerized-concurrent
 
@@ -125,7 +122,7 @@ fmt:
 	cargo +nightly fmt --all
 
 build-examples:
-	HANDLERS_TO_BUILD=${HANDLERS_TO_BUILD} OUTPUT_DIR=${OUTPUT_DIR} ./scripts/build-examples.sh
+	HANDLERS_TO_BUILD="$(subst ",,$(HANDLERS_TO_BUILD))" OUTPUT_DIR="$(OUTPUT_DIR)" ./scripts/build-examples.sh
 
 nuke:
 	docker kill $$(docker ps -q)
@@ -145,6 +142,7 @@ build-test-runner: build-examples
 	@echo "Building test runner Docker image..."
 	@docker build -t test-runner:local -f .test-runner/Dockerfile .test-runner
 
+test-dockerized-concurrent: HANDLERS_TO_BUILD := basic-lambda-concurrent invocation-id-concurrent
 test-dockerized-concurrent: build-test-runner
 	@echo "Running concurrent scenarios in Docker..."
 	@docker network rm concurrent-test-net 2>/dev/null || true
@@ -156,6 +154,7 @@ test-dockerized-concurrent: build-test-runner
 		-e TASK_FOLDER=./test/dockerized/tasks \
 		-e GITHUB_WORKSPACE=/workspace \
 		-e DOCKER_SHARED_NETWORK=concurrent-test-net \
+		-e CONTAINER_READY_DELAY_SECS=$(CONTAINER_READY_DELAY_SECS) \
 		-v /var/run/docker.sock:/var/run/docker.sock \
 		-v "$(CURDIR):/workspace" \
 		-w /workspace \

examples/invocation-id-concurrent/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "invocation-id-concurrent"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+lambda_runtime = { path = "../../lambda-runtime", features = ["concurrency-tokio"] }
+serde = "1.0.219"
+tokio = { version = "1", features = ["macros", "rt", "time"] }

examples/invocation-id-concurrent/src/main.rs

@@ -0,0 +1,96 @@
+// This example requires the following input to succeed:
+// { "command": "do something" }
+
+use lambda_runtime::{service_fn, tracing, Diagnostic, Error, LambdaEvent};
+use serde::{Deserialize, Serialize};
+
+#[derive(Deserialize)]
+struct Request {
+    #[serde(rename = "command")]
+    _command: String,
+    sleep: u32,
+}
+
+#[derive(Serialize, Debug, PartialEq)]
+struct Response {
+    from: String,
+}
+
+#[derive(Debug)]
+struct HandlerError(String);
+
+impl std::fmt::Display for HandlerError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl From<HandlerError> for Diagnostic {
+    fn from(e: HandlerError) -> Diagnostic {
+        Diagnostic {
+            error_type: "HandlerError".into(),
+            error_message: e.0,
+        }
+    }
+}
+
+
+/**
+ * Cross-wiring protection: duplicate request-id after timeout.
+
+    Timeline:
+        t=0: Invoke A starts, handler sleeps 7s
+        t=5: A times out (timeout=5s). Batch 1 completes with timeout error.
+        t=5: Invoke B starts (same request-id), handler sleeps 4s
+        t=7: A's handler wakes up, posts stale /response/{same-id}
+        t=9: B's handler wakes up, posts correct /response/{same-id}
+
+        With invocation-id: A's stale post at t=7 gets 410 Gone. B responds at t=9 correctly.
+        Without: A's stale response at t=7 is accepted for B (cross-wired).
+ */
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    // required to enable CloudWatch error logging by the runtime
+    tracing::init_default_subscriber();
+    let max_concurrency = std::env::var("AWS_LAMBDA_MAX_CONCURRENCY").unwrap_or_else(|_| "not set".to_string());
+    tracing::info!(AWS_LAMBDA_MAX_CONCURRENCY = %max_concurrency, "starting concurrent handler");
+
+    let func = service_fn(my_handler);
+    if let Err(err) = lambda_runtime::run_concurrent(func).await {
+        tracing::error!(error = %err, "run error");
+        return Err(err);
+    }
+    Ok(())
+}
+
+pub(crate) async fn my_handler(event: LambdaEvent<Request>) -> Result<Response, HandlerError> {
+    if event.payload.sleep > 0 {
+        tokio::time::sleep(tokio::time::Duration::from_secs(event.payload.sleep.into())).await;
+    }
+
+    Ok(Response {
+        from: event.payload._command,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use lambda_runtime::{Context, LambdaEvent};
+
+    #[tokio::test]
+    async fn handler_echoes_marker() {
+        let event = LambdaEvent {
+            payload: Request {
+                _command: "invoke-B".into(),
+                sleep: 0,
+            },
+            context: Context::default(),
+        };
+
+        let result = my_handler(event).await.unwrap();
+
+        assert_eq!(result, Response { from: "invoke-B".into() });
+    }
+}

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";

lambda-runtime/src/layers/api_response.rs

@@ -1,4 +1,5 @@
 use crate::{
+    constants::LAMBDA_RUNTIME_INVOCATION_ID,
     deserializer,
     requests::{EventCompletionRequest, IntoRequest},
     runtime::LambdaInvocation,
@@ -123,9 +124,20 @@ where
         };
 
         let request_id = req.context.request_id.clone();
+
+        // 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
+        let invocation_id = req
+            .parts
+            .headers
+            .get(LAMBDA_RUNTIME_INVOCATION_ID)
+            .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string());
+
         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 +149,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,11 +200,47 @@ 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"),
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{constants::LAMBDA_RUNTIME_INVOCATION_ID, runtime::LambdaInvocation, Context};
+    use http::{HeaderValue, Response};
+    use serde_json::json;
+    use tower::{service_fn, Service};
+
+    #[tokio::test]
+    async fn forwards_invocation_id_from_next_response_headers() {
+        let mut response = Response::new(());
+        response
+            .headers_mut()
+            .insert(LAMBDA_RUNTIME_INVOCATION_ID, HeaderValue::from_static("invocation-123"));
+        let (parts, _) = response.into_parts();
+
+        let mut service = RuntimeApiResponseService::new(service_fn(|_event: LambdaEvent<serde_json::Value>| async {
+            Ok::<_, Diagnostic>(json!({"ok": true}))
+        }));
+
+        let request = service
+            .call(LambdaInvocation {
+                parts,
+                body: bytes::Bytes::from_static(b"{}"),
+                context: Context::default(),
+            })
+            .await
+            .expect("response request should be created");
+
+        assert_eq!(
+            request.headers().get(LAMBDA_RUNTIME_INVOCATION_ID),
+            Some(&HeaderValue::from_static("invocation-123")),
+        );
+    }
+}

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;