From eccd045d80f7333f39a4a5dde6da822aeaed8aea Mon Sep 17 00:00:00 2001 From: Manoj Mahapatra Date: Wed, 18 Feb 2026 06:15:04 -0800 Subject: [PATCH] fix: Local Server should pass HTTP headers down to the Lambda Runtime (#643) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Issue \# https://github.com/awslabs/swift-aws-lambda-runtime/issues/607 ## Description of changes The local HTTP server was not forwarding user‑provided headers to the runtime’s response. It passes all headers through to the runtime. This it makes local behavior match the Lambda runtime API contract and allows developers to opt into metadata by sending the appropriate runtime headers. ## New/existing dependencies impact assessment, if applicable N/A ## Conventional Commits By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license. --------- Co-authored-by: Sébastien Stormacq --- Examples/MultiTenant/Package.swift | 11 +- Examples/MultiTenant/README.md | 23 +++- .../Sources/{ => MultiTenant}/main.swift | 0 .../Sources/MultiTenantLocal/main.swift | 117 ++++++++++++++++++ .../HTTPServer/Lambda+LocalServer.swift | 15 ++- .../LambdaLocalServerTests.swift | 53 +++++++- 6 files changed, 212 insertions(+), 7 deletions(-) rename Examples/MultiTenant/Sources/{ => MultiTenant}/main.swift (100%) create mode 100644 Examples/MultiTenant/Sources/MultiTenantLocal/main.swift diff --git a/Examples/MultiTenant/Package.swift b/Examples/MultiTenant/Package.swift index 82b9c277..43c15c54 100644 --- a/Examples/MultiTenant/Package.swift +++ b/Examples/MultiTenant/Package.swift @@ -6,7 +6,8 @@ let package = Package( name: "swift-aws-lambda-runtime-example", platforms: [.macOS(.v15)], products: [ - .executable(name: "MultiTenant", targets: ["MultiTenant"]) + .executable(name: "MultiTenant", targets: ["MultiTenant"]), + .executable(name: "MultiTenantLocal", targets: ["MultiTenantLocal"]), ], dependencies: [ // For local development (default) @@ -24,6 +25,12 @@ let package = Package( .product(name: "AWSLambdaRuntime", package: "swift-aws-lambda-runtime"), .product(name: "AWSLambdaEvents", package: "swift-aws-lambda-events"), ] - ) + ), + .executableTarget( + name: "MultiTenantLocal", + dependencies: [ + .product(name: "AWSLambdaRuntime", package: "swift-aws-lambda-runtime") + ] + ), ] ) diff --git a/Examples/MultiTenant/README.md b/Examples/MultiTenant/README.md index 0e8bba07..d4b0e671 100644 --- a/Examples/MultiTenant/README.md +++ b/Examples/MultiTenant/README.md @@ -6,7 +6,8 @@ This example demonstrates how to build a multi-tenant Lambda function using Swif This example implements a request tracking system that maintains separate counters and request histories for each tenant. The Lambda function: -- Accepts requests from multiple tenants via API Gateway +- Accepts requests from multiple tenants via API Gateway (MultiTenant) +- Accepts direct JSON payloads without API Gateway (MultiTenantLocal) - Maintains isolated execution environments per tenant - Tracks request counts and timestamps for each tenant - Returns tenant-specific data in JSON format @@ -41,6 +42,8 @@ The example consists of: 3. **Lambda Handler** - Processes API Gateway requests and manages tenant data +4. **MultiTenantLocal Handler** - Processes plain JSON requests without API Gateway + ## Code Structure ```swift @@ -175,6 +178,24 @@ requestParameters: ## Testing +### Local Testing (No API Gateway) + +Run the simple handler locally and invoke it directly using the local server. This example does not depend on API Gateway or its request/response types. + +1. Start the local server: + ```bash + swift run MultiTenantLocal + ``` + +2. Invoke with a tenant header and a JSON body: + ```bash + # Test multi-tenant function locally + curl -X POST http://127.0.0.1:7000/invoke \ + -H "Content-Type: application/json" \ + -H "Lambda-Runtime-Aws-Tenant-Id: tenant-123" \ + -d '{"message" : "hello"}' + ``` + ### Using API Gateway The tenant ID is passed as a query parameter. API Gateway automatically maps it to the `X-Amz-Tenant-Id` header: diff --git a/Examples/MultiTenant/Sources/main.swift b/Examples/MultiTenant/Sources/MultiTenant/main.swift similarity index 100% rename from Examples/MultiTenant/Sources/main.swift rename to Examples/MultiTenant/Sources/MultiTenant/main.swift diff --git a/Examples/MultiTenant/Sources/MultiTenantLocal/main.swift b/Examples/MultiTenant/Sources/MultiTenantLocal/main.swift new file mode 100644 index 00000000..3a2d93e7 --- /dev/null +++ b/Examples/MultiTenant/Sources/MultiTenantLocal/main.swift @@ -0,0 +1,117 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftAWSLambdaRuntime open source project +// +// Copyright SwiftAWSLambdaRuntime project authors +// Copyright (c) Amazon.com, Inc. or its affiliates. +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +import AWSLambdaRuntime + +#if canImport(FoundationEssentials) +import FoundationEssentials +#else +import Foundation +#endif + +let tenants = TenantDataStore() + +struct TenantMessage: Codable { + let message: String +} + +enum TenantError: Error { + case missingTenantID +} + +struct SimpleTenantHandler: LambdaHandler { + typealias Event = TenantMessage + typealias Output = TenantData + + let tenants: TenantDataStore + + func handle(_ event: TenantMessage, context: LambdaContext) async throws -> TenantData { + guard let tenantID = context.tenantID else { + throw TenantError.missingTenantID + } + + let currentData = await tenants[tenantID] ?? TenantData(tenantID: tenantID) + let updatedData = currentData.addingRequest(message: event.message) + await tenants.update(id: tenantID, data: updatedData) + return updatedData + } +} + +let handler = SimpleTenantHandler(tenants: tenants) +let runtime = LambdaRuntime(handler: LambdaCodableAdapter(handler: LambdaHandlerAdapter(handler: handler))) + +try await runtime.run() + +actor TenantDataStore { + private var tenants: [String: TenantData] = [:] + + subscript(id: String) -> TenantData? { + tenants[id] + } + + // subscript setters can't be called from outside of the actor + func update(id: String, data: TenantData) { + tenants[id] = data + } +} + +struct TenantData: Codable { + struct TenantRequest: Codable { + let requestNumber: Int + let timestamp: String + let message: String + } + + let tenantID: String + let requestCount: Int + let firstRequest: String + let requests: [TenantRequest] + + init(tenantID: String) { + self.init( + tenantID: tenantID, + requestCount: 0, + firstRequest: "\(Date().timeIntervalSince1970)", + requests: [] + ) + } + + func addingRequest(message: String) -> TenantData { + let newCount = requestCount + 1 + let newRequest = TenantRequest( + requestNumber: newCount, + timestamp: "\(Date().timeIntervalSince1970)", + message: message + ) + return TenantData( + tenantID: tenantID, + requestCount: newCount, + firstRequest: firstRequest, + requests: requests + [newRequest] + ) + } + + private init( + tenantID: String, + requestCount: Int, + firstRequest: String, + requests: [TenantRequest] + ) { + self.tenantID = tenantID + self.requestCount = requestCount + self.firstRequest = firstRequest + self.requests = requests + } +} diff --git a/Sources/AWSLambdaRuntime/HTTPServer/Lambda+LocalServer.swift b/Sources/AWSLambdaRuntime/HTTPServer/Lambda+LocalServer.swift index 8f040c72..97305ab7 100644 --- a/Sources/AWSLambdaRuntime/HTTPServer/Lambda+LocalServer.swift +++ b/Sources/AWSLambdaRuntime/HTTPServer/Lambda+LocalServer.swift @@ -383,7 +383,9 @@ internal struct LambdaHTTPServer { logger[metadataKey: "requestId"] = "\(requestId)" logger.trace("/invoke received invocation, pushing it to the pool and wait for a lambda response") - self.invocationPool.push(LocalServerInvocation(requestId: requestId, request: body)) + self.invocationPool.push( + LocalServerInvocation(requestId: requestId, request: body, headers: head.headers) + ) // wait for the lambda function to process the request // Handle streaming responses by collecting all chunks for this requestId @@ -518,7 +520,10 @@ internal struct LambdaHTTPServer { var headers = response.headers ?? HTTPHeaders() if let body = response.body { - headers.add(name: "Content-Length", value: "\(body.readableBytes)") + // Avoid adding Content-Length if already provided (e.g. forwarded from /invoke) + if !headers.contains(name: "Content-Length") && !headers.contains(name: "Transfer-Encoding") { + headers.add(name: "Content-Length", value: "\(body.readableBytes)") + } } if let status = response.status { @@ -569,11 +574,12 @@ internal struct LambdaHTTPServer { struct LocalServerInvocation: Sendable { let requestId: String let request: ByteBuffer + let headers: HTTPHeaders func acceptedResponse() -> LocalServerResponse { // required headers - let headers = HTTPHeaders([ + var headers = HTTPHeaders([ (AmazonHeaders.requestID, self.requestId), ( AmazonHeaders.invokedFunctionARN, @@ -583,6 +589,9 @@ internal struct LambdaHTTPServer { (AmazonHeaders.deadline, "\(LambdaClock.maxLambdaDeadline)"), ]) + // Forward all invocation headers + headers.add(contentsOf: self.headers) + return LocalServerResponse( id: self.requestId, status: .accepted, diff --git a/Tests/AWSLambdaRuntimeTests/LambdaLocalServerTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaLocalServerTests.swift index 844bca68..a6c576a7 100644 --- a/Tests/AWSLambdaRuntimeTests/LambdaLocalServerTests.swift +++ b/Tests/AWSLambdaRuntimeTests/LambdaLocalServerTests.swift @@ -181,11 +181,62 @@ struct LambdaLocalServerTest { } } - private func makeInvokeRequest(host: String, port: Int, payload: String) async throws -> (Data, HTTPURLResponse) { + @Test("Local server forwards invocation headers to LambdaContext") + @available(LambdaSwift 2.0, *) + func testLocalServerForwardsInvocationHeaders() async throws { + let customPort = 8082 + + setenv("LOCAL_LAMBDA_PORT", "\(customPort)", 1) + defer { unsetenv("LOCAL_LAMBDA_PORT") } + + let logger = Logger(label: "test", factory: { _ in SwiftLogNoOpLogHandler() }) + + let tenantId = try await withTimeout(deadline: .seconds(5)) { + async let serverTask: String? = LambdaHTTPServer.withLocalServer( + host: "127.0.0.1", + port: customPort, + invocationEndpoint: nil, + logger: logger + ) { + try await LambdaRuntimeClient.withRuntimeClient( + configuration: .init(ip: "127.0.0.1", port: customPort), + eventLoop: Lambda.defaultEventLoop, + logger: logger + ) { runtimeClient in + let (invocation, writer) = try await runtimeClient.nextInvocation() + try await writer.writeAndFinish(ByteBuffer(string: "\"ok\"")) + return invocation.metadata.tenantID + } + } + + try await Task.sleep(for: .milliseconds(200)) + + _ = try await self.makeInvokeRequest( + host: "127.0.0.1", + port: customPort, + payload: "\"ping\"", + headers: ["Lambda-Runtime-Aws-Tenant-Id": "123"] + ) + + return try await serverTask + } + + #expect(tenantId == "123") + } + + private func makeInvokeRequest( + host: String, + port: Int, + payload: String, + headers: [String: String] = [:] + ) async throws -> (Data, HTTPURLResponse) { let url = URL(string: "http://\(host):\(port)/invoke")! var request = URLRequest(url: url) request.httpMethod = "POST" request.setValue("application/json", forHTTPHeaderField: "Content-Type") + for (key, value) in headers { + request.setValue(value, forHTTPHeaderField: key) + } request.httpBody = payload.data(using: .utf8) request.timeoutInterval = 10.0