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..4d61e682 --- /dev/null +++ b/examples/invocation-id-concurrent/Cargo.toml @@ -0,0 +1,12 @@ +[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"] } + +[dev-dependencies] +serde_json = "=1.0.151" diff --git a/examples/invocation-id-concurrent/src/main.rs b/examples/invocation-id-concurrent/src/main.rs new file mode 100644 index 00000000..acca09a0 --- /dev/null +++ b/examples/invocation-id-concurrent/src/main.rs @@ -0,0 +1,103 @@ +// 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(default, 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() }); + } + + #[test] + fn request_defaults_missing_command() { + let request: Request = serde_json::from_str(r#"{"sleep": 0}"#).unwrap(); + + assert_eq!(request._command, ""); + } +} 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_client.rs b/lambda-runtime/src/layers/api_client.rs index bc2bcd0f..d91ea78f 100644 --- a/lambda-runtime/src/layers/api_client.rs +++ b/lambda-runtime/src/layers/api_client.rs @@ -114,8 +114,12 @@ where // Adding more information on top of 410 Gone, to make it more clear since we cannot access the body of the message if status == 410 { log_or_print!( - tracing: tracing::error!("Lambda function timeout!"), - fallback: eprintln!("Lambda function timeout!") + tracing: tracing::error!( + "Lambda Runtime API rejected the response: invocation timed out or response was stale" + ), + fallback: eprintln!( + "Lambda Runtime API rejected the response: invocation timed out or response was stale" + ) ); } @@ -219,7 +223,9 @@ mod tests { // Verify the error was logged assert!(logs_contain("Lambda Runtime API returned non-200 response")); - assert!(logs_contain("Lambda function timeout!")); + assert!(logs_contain( + "Lambda Runtime API rejected the response: invocation timed out or response was stale" + )); } #[tokio::test] diff --git a/lambda-runtime/src/layers/api_response.rs b/lambda-runtime/src/layers/api_response.rs index 5bb3c96f..11d9c5b3 100644 --- a/lambda-runtime/src/layers/api_response.rs +++ b/lambda-runtime/src/layers/api_response.rs @@ -1,5 +1,7 @@ use crate::{ + constants::LAMBDA_RUNTIME_INVOCATION_ID, deserializer, + rate_limiter::RateLimiter, requests::{EventCompletionRequest, IntoRequest}, runtime::LambdaInvocation, Diagnostic, EventErrorRequest, IntoFunctionResponse, LambdaEvent, @@ -8,9 +10,11 @@ use futures::{ready, Stream}; use lambda_runtime_api_client::{body::Body, BoxError}; use pin_project::pin_project; use serde::{Deserialize, Serialize}; -use std::{fmt::Debug, future::Future, marker::PhantomData, pin::Pin, task}; +use std::{fmt::Debug, future::Future, marker::PhantomData, pin::Pin, task, time::Duration}; use tower::Service; -use tracing::{error, trace}; +use tracing::{error, trace, warn}; + +static MALFORMED_INVOCATION_ID_LIMITER: RateLimiter = RateLimiter::new(Duration::from_secs(60)); /// Tower service that turns the result or an error of a handler function into a Lambda Runtime API /// response. @@ -123,9 +127,31 @@ 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 = match req.parts.headers.get(LAMBDA_RUNTIME_INVOCATION_ID) { + Some(value) => match value.to_str() { + Ok(value) => Some(value.to_owned()), + Err(error) => { + if MALFORMED_INVOCATION_ID_LIMITER.allow() { + warn!( + error = ?error, + rate_limit_interval_ms = MALFORMED_INVOCATION_ID_LIMITER.interval().as_millis(), + "Ignoring malformed Lambda runtime invocation ID header; this warning is rate limited" + ); + } + None + } + }, + None => None, + }; + 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 +163,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 +184,7 @@ pub enum RuntimeApiResponseFuture, PhantomData<( (), Response, @@ -183,11 +214,76 @@ 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")), + ); + } + + #[tokio::test] + async fn malformed_invocation_id_does_not_block_deserialization_error_response() { + let mut response = Response::new(()); + response.headers_mut().insert( + LAMBDA_RUNTIME_INVOCATION_ID, + HeaderValue::from_bytes(&[0xff]).expect("header value should accept opaque bytes"), + ); + 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 { + request_id: "request-123".to_owned(), + ..Context::default() + }, + }) + .await + .expect("deserialization errors should produce an error request"); + + assert_eq!(request.uri().path(), "/2018-06-01/runtime/invocation/request-123/error",); + assert!(request.headers().get(LAMBDA_RUNTIME_INVOCATION_ID).is_none()); + } +} diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index 69c04ebc..ce6688ba 100644 --- a/lambda-runtime/src/lib.rs +++ b/lambda-runtime/src/lib.rs @@ -23,10 +23,14 @@ 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; +mod rate_limiter; + mod deserializer; /// Tower middleware to be applied to runtime invocations. pub mod layers; diff --git a/lambda-runtime/src/rate_limiter.rs b/lambda-runtime/src/rate_limiter.rs new file mode 100644 index 00000000..f2fa6514 --- /dev/null +++ b/lambda-runtime/src/rate_limiter.rs @@ -0,0 +1,111 @@ +//! Thread-safe, process-local rate limiting for infrequent runtime events. + +use std::{ + sync::Mutex, + time::{Duration, Instant}, +}; + +/// Allows an operation at most once during each configured interval. +/// +/// A `RateLimiter` is intended to be shared by concurrent runtime tasks. When +/// stored in a `static`, it is initialized once per Lambda execution environment +/// and retains its state across warm invocations. A new cold-started environment +/// receives a new limiter. +pub(crate) struct RateLimiter { + /// Minimum duration between allowed operations. + interval: Duration, + /// Timestamp of the most recent allowed operation. + last_allowed: Mutex>, +} + +impl RateLimiter { + /// Creates a rate limiter with the specified minimum interval. + pub(crate) const fn new(interval: Duration) -> RateLimiter { + RateLimiter { + interval, + last_allowed: Mutex::new(None), + } + } + + /// Returns the minimum duration between allowed operations. + pub(crate) const fn interval(&self) -> Duration { + self.interval + } + + /// + /// The first call is allowed. Subsequent calls are rejected until the + /// configured interval has elapsed since the previous allowed call. + /// Concurrent callers are serialized while checking and updating the last + /// allowed timestamp so only one caller crosses the interval boundary. + pub(crate) fn allow(&self) -> bool { + let mut last_allowed = match self.last_allowed.lock() { + Ok(guard) => guard, + Err(poisoned) => { + // The limiter state is disposable, so reset it and recover instead of + // allowing a poisoned mutex to crash the runtime or suppress future logs. + let mut guard = poisoned.into_inner(); + *guard = None; + self.last_allowed.clear_poison(); + guard + } + }; + + if last_allowed + .as_ref() + .is_some_and(|value| value.elapsed() < self.interval) + { + return false; + } + + *last_allowed = Some(Instant::now()); + + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::{ + panic::{catch_unwind, AssertUnwindSafe}, + thread, + }; + + #[test] + fn allows_first_call() { + let limiter = RateLimiter::new(Duration::from_secs(60)); + + assert!(limiter.allow()); + } + + #[test] + fn rejects_calls_inside_interval() { + let limiter = RateLimiter::new(Duration::from_secs(60)); + + assert!(limiter.allow()); + assert!(!limiter.allow()); + } + + #[test] + fn allows_call_after_interval() { + let limiter = RateLimiter::new(Duration::from_millis(10)); + + assert!(limiter.allow()); + thread::sleep(Duration::from_millis(15)); + + assert!(limiter.allow()); + } + + #[test] + fn recovers_from_poisoned_mutex() { + let limiter = RateLimiter::new(Duration::from_secs(60)); + + let _ = catch_unwind(AssertUnwindSafe(|| { + let _guard = limiter.last_allowed.lock().unwrap(); + panic!("poison the limiter mutex"); + })); + + assert!(limiter.allow()); + assert!(!limiter.allow()); + } +} diff --git a/lambda-runtime/src/requests.rs b/lambda-runtime/src/requests.rs index b03f14c7..933e3100 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.insert(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..8370289b 100644 --- a/test/dockerized/scenarios/concurrent_scenarios.py +++ b/test/dockerized/scenarios/concurrent_scenarios.py @@ -1,26 +1,34 @@ -""" -Multi-concurrency test scenarios for basic-lambda-concurrent. - -The handler expects: { "command": "" } -and responds with: { "req_id": "", "msg": "Command executed." } -""" +"""Multi-concurrency test scenarios.""" import os -from containerized_test_runner.models import Request, ConcurrentTest +from containerized_test_runner.models import ConcurrentTest, Request 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: +def _make_env(handler: str = HANDLER, concurrency: int = DEFAULT_CONCURRENCY) -> dict: return { - "_HANDLER": HANDLER, + "_HANDLER": handler, "AWS_LAMBDA_MAX_CONCURRENCY": str(concurrency), "AWS_LAMBDA_LOG_FORMAT": "JSON", } +def _invocation_id_env( + handler: str = INVOCATION_ID_HANDLER, + concurrency: int = DEFAULT_CONCURRENCY, + timeout: int = TIMEOUT, +) -> dict: + return _make_env(handler, concurrency) | { + "AWS_LAMBDA_FUNCTION_TIMEOUT": str(timeout), + } + + def get_concurrent_scenarios(): scenarios = [] @@ -62,3 +70,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(handler=INVOCATION_ID_HANDLER, timeout=TIMEOUT), + request_batches=batches, + image=IMAGE, + )] +