diff --git a/.env b/.env deleted file mode 100644 index 1873f15d..00000000 --- a/.env +++ /dev/null @@ -1,16 +0,0 @@ -# Test configuration for RIE and dockerized tests -# Customize these values as needed for testing both local and on github - -# Handlers to build -HANDLERS_TO_BUILD="basic-lambda basic-sqs http-basic-lambda basic-lambda-concurrent" - -HANDLER=basic-lambda - -# Output directory for built binaries -OUTPUT_DIR=test/dockerized/tasks - -# Max concurrent Lambda invocations for LMI mode -RIE_MAX_CONCURRENCY=4 - -# Branch of containerized-test-runner-for-aws-lambda to clone -TEST_RUNNER_BRANCH=main diff --git a/.github/workflows/dockerized-test.yml b/.github/workflows/dockerized-test.yml index cb6cf495..be86e147 100644 --- a/.github/workflows/dockerized-test.yml +++ b/.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' diff --git a/Dockerfile.rie b/Dockerfile.rie index 1a46b577..b55545b2 100644 --- a/Dockerfile.rie +++ b/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 diff --git a/Dockerfile.test b/Dockerfile.test index b36b1f28..ece55bf6 100644 --- a/Dockerfile.test +++ b/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 diff --git a/Makefile b/Makefile index 2b652e12..ccc415a7 100644 --- a/Makefile +++ b/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 \ diff --git a/examples/invocation-id-concurrent/Cargo.toml b/examples/invocation-id-concurrent/Cargo.toml new file mode 100644 index 00000000..8d0841ce --- /dev/null +++ b/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"] } diff --git a/examples/invocation-id-concurrent/src/main.rs b/examples/invocation-id-concurrent/src/main.rs new file mode 100644 index 00000000..fd4af07f --- /dev/null +++ b/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 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) -> Result { + 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() }); + } +} 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..e7976c15 100644 --- a/lambda-runtime/src/layers/api_response.rs +++ b/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::(&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 +149,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 +170,7 @@ pub enum RuntimeApiResponseFuture, PhantomData<( (), Response, @@ -183,11 +200,47 @@ 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"), }) } } + +#[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| 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")), + ); + } +} 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..a4b63702 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,12 @@ 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).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 +159,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 +212,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 +240,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 +270,92 @@ 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_streaming_event_completion_request_with_invocation_id() { + use crate::StreamResponse; + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + runtime.block_on(async { + let stream = tokio_stream::iter(vec![Ok::(Bytes::from_static(b"chunk"))]); + let stream_response: StreamResponse<_> = stream.into(); + let response = FunctionResponse::StreamingResponse(stream_response); + + let req: EventCompletionRequest<'_, _, (), _, _, _> = + EventCompletionRequest::new("id", Some("invocation_id"), response); + + let http_req = req.into_req().expect("into_req should succeed"); + let expected = Uri::from_static("/2018-06-01/runtime/invocation/id/response"); + assert_eq!(http_req.method(), Method::POST); + assert_eq!(http_req.uri(), &expected); + + assert_eq!( + http_req.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).unwrap(), + "invocation_id" + ); + }); + } + + #[test] + fn test_streaming_event_completion_request_invocation_id_not_added_when_none() { + use crate::StreamResponse; + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + runtime.block_on(async { + let stream = tokio_stream::iter(vec![Ok::(Bytes::from_static(b"chunk"))]); + let stream_response: StreamResponse<_> = stream.into(); + let response = FunctionResponse::StreamingResponse(stream_response); + + let req: EventCompletionRequest<'_, _, (), _, _, _> = EventCompletionRequest::new("id", None, response); + + let http_req = req.into_req().expect("into_req should succeed"); + + assert!(http_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 +369,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] @@ -343,7 +462,7 @@ mod tests { let stream_response: StreamResponse<_> = stream.into(); let response = FunctionResponse::StreamingResponse(stream_response); - let req: EventCompletionRequest<'_, _, (), _, _, _> = EventCompletionRequest::new("id", response); + let req: EventCompletionRequest<'_, _, (), _, _, _> = EventCompletionRequest::new("id", None, response); let http_req = req.into_req().expect("into_req should succeed"); 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..edd93240 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_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}; @@ -106,7 +113,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 +124,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,24 +138,24 @@ 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, }; @@ -165,7 +172,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 +298,7 @@ where #[cfg(test)] mod test { + use super::*; use crate::Config; use std::sync::Arc; diff --git a/scripts/build-examples.sh b/scripts/build-examples.sh index b479059d..2e577c3a 100755 --- a/scripts/build-examples.sh +++ b/scripts/build-examples.sh @@ -11,13 +11,25 @@ echo "Building handlers: ${HANDLERS_TO_BUILD}" for handler in ${HANDLERS_TO_BUILD}; do dir="examples/$handler" - [ ! -f "$dir/Cargo.toml" ] && echo "✗ $handler not found" && continue - + if [ ! -f "$dir/Cargo.toml" ]; then + echo "✗ $handler not found" + continue + fi + echo "Building $handler..." - (cd "$dir" && cargo build --release) || continue - - [ -f "$dir/target/release/$handler" ] && cp "$dir/target/release/$handler" "$OUTPUT_DIR/" && echo "✓ $handler" + if ! (cd "$dir" && cargo build --release); then + continue + fi + + if [ ! -f "$dir/target/release/$handler" ]; then + echo "✗ $handler artifact not found" + continue + fi + + cp "$dir/target/release/$handler" "$OUTPUT_DIR/" + echo "✓ $handler" done echo "" ls -lh "$OUTPUT_DIR/" 2>/dev/null || echo "No binaries built" +exit 0 diff --git a/scripts/download-rie.sh b/scripts/download-rie.sh new file mode 100644 index 00000000..0d22509a --- /dev/null +++ b/scripts/download-rie.sh @@ -0,0 +1,46 @@ +#!/bin/sh + +set -eu + +if [ "$#" -ne 5 ]; then + echo "Usage: $0 TARGETARCH RIE_VERSION RIE_SHA256_AMD64 RIE_SHA256_ARM64 RIE_PATH" >&2 + exit 1 +fi + +TARGETARCH=$1 +RIE_VERSION=$2 +RIE_SHA256_AMD64=$3 +RIE_SHA256_ARM64=$4 +RIE_PATH=$5 + +case "${TARGETARCH}" in + amd64) + RIE_ASSET=aws-lambda-rie + RIE_SHA256=${RIE_SHA256_AMD64} + ;; + arm64) + RIE_ASSET=aws-lambda-rie-arm64 + RIE_SHA256=${RIE_SHA256_ARM64} + ;; + *) + echo "Unsupported target architecture: ${TARGETARCH}" >&2 + exit 1 + ;; +esac + +: "${RIE_PATH:?RIE_PATH must be set}" +RIE_TMP=$(mktemp) +trap 'rm -f "${RIE_TMP}"' EXIT + +curl \ + --fail \ + --location \ + --silent \ + --show-error \ + --retry 3 \ + --retry-all-errors \ + --output "${RIE_TMP}" \ + "https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/download/v${RIE_VERSION}/${RIE_ASSET}" + +echo "${RIE_SHA256} ${RIE_TMP}" | sha256sum --check --status +install -m 0755 "${RIE_TMP}" "${RIE_PATH}" diff --git a/scripts/test-rie.sh b/scripts/test-rie.sh index c5949fe8..9de6fedd 100755 --- a/scripts/test-rie.sh +++ b/scripts/test-rie.sh @@ -18,7 +18,7 @@ fi CONTAINER_PID=$! echo "Container started. Test with:" -if [ "$EXAMPLE" = "basic-lambda" ] || [ "$EXAMPLE" = "basic-lambda-concurrent" ]; then +if [ "$EXAMPLE" = "basic-lambda" ] || [ "$EXAMPLE" = "basic-lambda-concurrent" ] || [ "$EXAMPLE" = "invocation-id-concurrent" ]; then echo "curl -XPOST 'http://localhost:9000/2015-03-31/functions/function/invocations' -d '{\"command\": \"test from RIE\"}' -H 'Content-Type: application/json'" else echo "For example '$EXAMPLE', check examples/$EXAMPLE/src/main.rs for the expected payload format." diff --git a/test/dockerized/scenarios/concurrent_scenarios.py b/test/dockerized/scenarios/concurrent_scenarios.py index 8c2ce842..c7a1127f 100644 --- a/test/dockerized/scenarios/concurrent_scenarios.py +++ b/test/dockerized/scenarios/concurrent_scenarios.py @@ -9,8 +9,11 @@ from containerized_test_runner.models import Request, ConcurrentTest HANDLER = "basic-lambda-concurrent" +INVOCATION_ID_HANDLER = "invocation-id-concurrent" IMAGE = os.environ.get("TEST_IMAGE", "local/test-base") +SAME_REQUEST_ID = "shared-request-id" DEFAULT_CONCURRENCY = 10 +TIMEOUT = 5 def _make_env(concurrency: int = DEFAULT_CONCURRENCY) -> dict: @@ -21,6 +24,12 @@ def _make_env(concurrency: int = DEFAULT_CONCURRENCY) -> dict: } +def _invocation_id_env(concurrency: int = DEFAULT_CONCURRENCY, timeout: int = TIMEOUT) -> dict: + return _make_env(concurrency) | { + "AWS_LAMBDA_FUNCTION_TIMEOUT": str(timeout), + } + + def get_concurrent_scenarios(): scenarios = [] @@ -62,3 +71,28 @@ def get_concurrent_scenarios(): )) return scenarios + + +def get_invocation_id_scenarios(): + batches = [ + [Request.create( + payload={"command": "invoke-A", "sleep": TIMEOUT + 2}, + assertions=[{"transform": ".errorType", "error": "Sandbox.Timedout"}], + headers={"X-Amzn-RequestId": SAME_REQUEST_ID}, + )], + [Request.create( + payload={"command": "invoke-B", "sleep": TIMEOUT - 1}, + assertions=[{"response": {"from": "invoke-B"}}], + headers={"X-Amzn-RequestId": SAME_REQUEST_ID}, + )], + ] + + + return [ConcurrentTest( + name="invocation_id", + handler=INVOCATION_ID_HANDLER, + environment_variables=_invocation_id_env(timeout=TIMEOUT), + request_batches=batches, + image=IMAGE, + )] +