feat: add multi-tenancy support via tenant ID propagation

README.md

@@ -18,6 +18,7 @@ The same docker image can run on AWS Lambda, Amazon EC2, AWS Fargate, and local
 - Enables graceful shutdown
 - Supports response payload compression
 - Supports response streaming
+- Supports multi-tenancy via tenant ID propagation
 - Supports non-http event triggers
 
 ## Usage
@@ -109,16 +110,16 @@ The readiness check port/path and traffic port can be configured using environme
 | AWS_LWA_ERROR_STATUS_CODES                                  | comma-separated list of HTTP status codes that will cause Lambda invocations to fail (e.g. "500,502-504,422") | None  |
 | AWS_LWA_LAMBDA_RUNTIME_API_PROXY                              | overwrites `AWS_LAMBDA_RUNTIME_API` to allow proxying request (not affecting registration)                    | None       |
 
-**AWS_LWA_PORT** - Lambda Web Adapter will send traffic to this port. This is the port your web application listening on. Inside Lambda execution environment,
-the web application runs as a non-root user, and not allowed to listen on ports lower than 1024. Please also avoid port 9001 and 3000.
-Lambda Runtime API is on port 9001. CloudWatch Lambda Insight extension uses port 3000.
-
 > **Deprecation Notice:** The following non-namespaced environment variables are deprecated and will be removed in version 2.0:
 > `HOST`, `READINESS_CHECK_PORT`, `READINESS_CHECK_PATH`, `READINESS_CHECK_PROTOCOL`, `REMOVE_BASE_PATH`, `ASYNC_INIT`.
 > Please migrate to the `AWS_LWA_` prefixed versions. Note: `PORT` is not deprecated and remains a supported fallback for `AWS_LWA_PORT`.
 >
 > Additionally, `AWS_LWA_READINESS_CHECK_MIN_UNHEALTHY_STATUS` is deprecated. Use `AWS_LWA_READINESS_CHECK_HEALTHY_STATUS` instead.  
 
+**AWS_LWA_PORT** - Lambda Web Adapter will send traffic to this port. This is the port your web application listening on. Inside Lambda execution environment,
+the web application runs as a non-root user, and not allowed to listen on ports lower than 1024. Please also avoid port 9001 and 3000.
+Lambda Runtime API is on port 9001. CloudWatch Lambda Insight extension uses port 3000.
+
 **AWS_LWA_ASYNC_INIT** - Lambda managed runtimes offer up to 10 seconds for function initialization. During this period of time, Lambda functions have burst of CPU to accelerate initialization.
 If a lambda function couldn't complete the initialization within 10 seconds, Lambda will restart the function, and bill for the initialization.
 To help functions to use this 10 seconds free initialization time and avoid the restart, Lambda Web Adapter supports asynchronous initialization.
@@ -172,6 +173,30 @@ Lambda Web Adapter forwards this information to the web application in a Http He
 
 Lambda Web Adapter forwards this information to the web application in a Http Header named "x-amzn-lambda-context". In the web application, you can retrieve the value of this http header and deserialize it into a JSON object. Check out [Express.js in Zip](examples/expressjs-zip) on how to use it.
 
+## Multi-Tenancy Support
+
+Lambda Web Adapter supports multi-tenancy by automatically propagating the tenant ID from the Lambda runtime context to your web application.
+
+When the Lambda runtime includes a `tenant_id` in the invocation context, the adapter forwards it as an `X-Amz-Tenant-Id` HTTP header to your web application. This allows your application to identify the tenant for each request without any additional configuration.
+
+In your web application, you can read the tenant ID from the request header:
+
+```python
+# FastAPI example
+@app.get("/")
+def handler(request: Request):
+    tenant_id = request.headers.get("x-amz-tenant-id")
+```
+
+```javascript
+// Express.js example
+app.get('/', (req, res) => {
+    const tenantId = req.headers['x-amz-tenant-id'];
+});
+```
+
+If no tenant ID is present in the Lambda context, the header is simply omitted.
+
 ## Lambda Managed Instances
 
 Lambda Web Adapter supports [Lambda Managed Instances](https://docs.aws.amazon.com/lambda/latest/dg/lambda-managed-instances.html), which allows a single Lambda execution environment to handle multiple concurrent requests. This can improve throughput and reduce costs for I/O-bound workloads.

src/lib.rs

@@ -932,6 +932,16 @@ impl Adapter<HttpConnector, Body> {
             HeaderValue::from_bytes(serde_json::to_string(&lambda_context)?.as_bytes())?,
         );
 
+        // Multi-tenancy support: propagate tenant_id from Lambda context
+        if let Some(ref tenant_id) = lambda_context.tenant_id {
+            if let Ok(value) = HeaderValue::from_str(tenant_id) {
+                req_headers.insert(HeaderName::from_static("x-amz-tenant-id"), value);
+                tracing::debug!(tenant_id = %tenant_id, "propagating tenant_id header");
+            } else {
+                tracing::warn!(tenant_id = %tenant_id, "tenant_id contains invalid header characters, skipping");
+            }
+        }
+
         if let Some(authorization_source) = self.authorization_source.as_deref() {
             if let Some(original) = req_headers.remove(authorization_source) {
                 req_headers.insert("authorization", original);
@@ -1292,4 +1302,89 @@ mod tests {
             "Compression should remain enabled when invoke mode is Buffered"
         );
     }
+
+    /// Helper to create a Lambda Context with an optional tenant_id.
+    fn make_lambda_context(tenant_id: Option<&str>) -> lambda_http::Context {
+        use lambda_http::lambda_runtime::Config;
+        let mut headers = http::HeaderMap::new();
+        headers.insert("lambda-runtime-aws-request-id", "test-id".parse().unwrap());
+        headers.insert("lambda-runtime-deadline-ms", "123".parse().unwrap());
+        headers.insert("lambda-runtime-client-context", "{}".parse().unwrap());
+        if let Some(tid) = tenant_id {
+            headers.insert("lambda-runtime-aws-tenant-id", tid.parse().unwrap());
+        }
+        let conf = Config {
+            function_name: "test_function".into(),
+            memory: 128,
+            version: "latest".into(),
+            log_stream: "/aws/lambda/test_function".into(),
+            log_group: "2023/09/15/[$LATEST]ab831cef03e94457a94b6efcbe22406a".into(),
+        };
+        lambda_http::Context::new("test-id", Arc::new(conf), &headers).unwrap()
+    }
+
+    #[tokio::test]
+    async fn test_tenant_id_header_propagated() {
+        let app_server = MockServer::start();
+        app_server.mock(|when, then| {
+            when.method(GET).path("/hello").header("x-amz-tenant-id", "tenant-abc");
+            then.status(200).body("OK");
+        });
+
+        let options = AdapterOptions {
+            host: app_server.host(),
+            port: app_server.port().to_string(),
+            readiness_check_port: app_server.port().to_string(),
+            readiness_check_path: "/".to_string(),
+            ..Default::default()
+        };
+
+        let adapter = Adapter::new(&options).expect("Failed to create adapter");
+
+        // Build a minimal ALB request
+        let alb_req = lambda_http::request::LambdaRequest::Alb({
+            let mut req = lambda_http::aws_lambda_events::alb::AlbTargetGroupRequest::default();
+            req.http_method = Method::GET;
+            req.path = Some("/hello".into());
+            req
+        });
+        let mut request = Request::from(alb_req);
+        request.extensions_mut().insert(make_lambda_context(Some("tenant-abc")));
+
+        let response = adapter.fetch_response(request).await.expect("Request failed");
+        assert_eq!(200, response.status().as_u16());
+    }
+
+    #[tokio::test]
+    async fn test_tenant_id_header_absent_when_no_tenant() {
+        let app_server = MockServer::start();
+        app_server.mock(|when, then| {
+            when.method(GET)
+                .path("/hello")
+                .is_true(|req| !req.headers().iter().any(|(k, _)| k == "x-amz-tenant-id"));
+            then.status(200).body("OK");
+        });
+
+        let options = AdapterOptions {
+            host: app_server.host(),
+            port: app_server.port().to_string(),
+            readiness_check_port: app_server.port().to_string(),
+            readiness_check_path: "/".to_string(),
+            ..Default::default()
+        };
+
+        let adapter = Adapter::new(&options).expect("Failed to create adapter");
+
+        let alb_req = lambda_http::request::LambdaRequest::Alb({
+            let mut req = lambda_http::aws_lambda_events::alb::AlbTargetGroupRequest::default();
+            req.http_method = Method::GET;
+            req.path = Some("/hello".into());
+            req
+        });
+        let mut request = Request::from(alb_req);
+        request.extensions_mut().insert(make_lambda_context(None));
+
+        let response = adapter.fetch_response(request).await.expect("Request failed");
+        assert_eq!(200, response.status().as_u16());
+    }
 }

tests/integ_tests/main.rs

@@ -1296,3 +1296,80 @@ async fn test_concurrent_post_body_isolation() {
     json_a.assert_calls(2);
     json_b.assert_calls(2);
 }
+
+fn add_lambda_context_with_tenant(request: &mut Request<Body>, tenant_id: &str) {
+    let mut headers = HeaderMap::new();
+    headers.insert("lambda-runtime-aws-request-id", "my_id".parse().unwrap());
+    headers.insert("lambda-runtime-deadline-ms", "123".parse().unwrap());
+    headers.insert("lambda-runtime-client-context", "{}".parse().unwrap());
+    headers.insert("lambda-runtime-aws-tenant-id", tenant_id.parse().unwrap());
+
+    let conf = Config {
+        function_name: "test_function".into(),
+        memory: 128,
+        version: "latest".into(),
+        log_stream: "/aws/lambda/test_function".into(),
+        log_group: "2023/09/15/[$LATEST]ab831cef03e94457a94b6efcbe22406a".into(),
+    };
+
+    let context = Context::new("my_id", Arc::new(conf), &headers).expect("Couldn't convert HeaderMap to Context");
+    request.extensions_mut().insert(context);
+}
+
+#[tokio::test]
+async fn test_tenant_id_propagated_from_client_context() {
+    let app_server = MockServer::start();
+    let endpoint = app_server.mock(|when, then| {
+        when.method(GET)
+            .path("/hello")
+            .header("x-amz-tenant-id", "tenant-abc-123");
+        then.status(200).body("OK");
+    });
+
+    let mut adapter = Adapter::new(&AdapterOptions {
+        host: app_server.host(),
+        port: app_server.port().to_string(),
+        readiness_check_port: app_server.port().to_string(),
+        readiness_check_path: "/healthcheck".to_string(),
+        ..Default::default()
+    })
+    .expect("Failed to create adapter");
+
+    let req = LambdaEventBuilder::new().with_path("/hello").build();
+    let mut request = Request::from(req);
+    add_lambda_context_with_tenant(&mut request, "tenant-abc-123");
+
+    let response = adapter.call(request).await.expect("Request failed");
+
+    endpoint.assert();
+    assert_eq!(200, response.status());
+}
+
+#[tokio::test]
+async fn test_no_tenant_id_header_when_absent() {
+    let app_server = MockServer::start();
+    let endpoint = app_server.mock(|when, then| {
+        when.method(GET)
+            .path("/hello")
+            .is_true(|req| !req.headers().iter().any(|(k, _)| k == "x-amz-tenant-id"));
+        then.status(200).body("OK");
+    });
+
+    let mut adapter = Adapter::new(&AdapterOptions {
+        host: app_server.host(),
+        port: app_server.port().to_string(),
+        readiness_check_port: app_server.port().to_string(),
+        readiness_check_path: "/healthcheck".to_string(),
+        ..Default::default()
+    })
+    .expect("Failed to create adapter");
+
+    let req = LambdaEventBuilder::new().with_path("/hello").build();
+    let mut request = Request::from(req);
+    add_lambda_context_to_request(&mut request);
+
+    let response = adapter.call(request).await.expect("Request failed");
+
+    endpoint.assert();
+    assert_eq!(200, response.status());
+}