From b7c13acceca2b95717b91c7909e62efbd1138389 Mon Sep 17 00:00:00 2001 From: Ramisa Alam Date: Wed, 16 Apr 2025 15:23:01 +0000 Subject: [PATCH] feat: add tenant id to lambda context and structured log message --- .../api/client/EventHandlerLoader.java | 1 + .../runtime/api/client/api/LambdaContext.java | 7 ++ .../api/client/logging/JsonLogFormatter.java | 1 + .../client/logging/StructuredLogMessage.java | 1 + .../runtimeapi/dto/InvocationRequest.java | 13 ++++ ...ime_api_client_runtimeapi_NativeClient.cpp | 6 ++ .../include/aws/lambda-runtime/runtime.h | 5 ++ .../deps/aws-lambda-cpp-0.2.7/src/runtime.cpp | 5 ++ .../api/client/api/LambdaContextTest.java | 3 +- .../client/logging/JsonLogFormatterTest.java | 20 +++++ .../LambdaRuntimeApiClientImplTest.java | 78 ++++++++++++++++--- 11 files changed, 127 insertions(+), 13 deletions(-) diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java index 096bb862..db6ceceb 100644 --- a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java +++ b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java @@ -581,6 +581,7 @@ public ByteArrayOutputStream call(InvocationRequest request) throws Error, Excep cognitoIdentity, LambdaEnvironment.FUNCTION_VERSION, request.getInvokedFunctionArn(), + request.getTenantId(), clientContext ); diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java index 2ce3b844..bd1463db 100644 --- a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java +++ b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java @@ -22,6 +22,7 @@ public class LambdaContext implements Context { private final long deadlineTimeInMs; private final CognitoIdentity cognitoIdentity; private final ClientContext clientContext; + private final String tenantId; private final LambdaLogger logger; public LambdaContext( @@ -34,6 +35,7 @@ public LambdaContext( CognitoIdentity identity, String functionVersion, String invokedFunctionArn, + String tenantId, ClientContext clientContext ) { this.memoryLimit = memoryLimit; @@ -46,6 +48,7 @@ public LambdaContext( this.clientContext = clientContext; this.functionVersion = functionVersion; this.invokedFunctionArn = invokedFunctionArn; + this.tenantId = tenantId; this.logger = com.amazonaws.services.lambda.runtime.LambdaRuntime.getLogger(); } @@ -91,6 +94,10 @@ public int getRemainingTimeInMillis() { return delta > 0 ? delta : 0; } + public String getTenantId() { + return tenantId; + } + public LambdaLogger getLogger() { return logger; } diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatter.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatter.java index b98721eb..f463e7ee 100644 --- a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatter.java +++ b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatter.java @@ -41,6 +41,7 @@ private StructuredLogMessage createLogMessage(String message, LogLevel logLevel) if (lambdaContext != null) { msg.AWSRequestId = lambdaContext.getAwsRequestId(); + msg.tenantId = lambdaContext.getTenantId(); } return msg; } diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/StructuredLogMessage.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/StructuredLogMessage.java index 5299bffa..0ae19961 100644 --- a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/StructuredLogMessage.java +++ b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/StructuredLogMessage.java @@ -12,4 +12,5 @@ class StructuredLogMessage { public String message; public LogLevel level; public String AWSRequestId; + public String tenantId; } diff --git a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java index 7bdc2500..656945b4 100644 --- a/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java +++ b/aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java @@ -40,6 +40,11 @@ public class InvocationRequest { */ private String cognitoIdentity; + /** + * The tenant ID associated with the request. + */ + private String tenantId; + private byte[] content; public String getId() { @@ -94,6 +99,14 @@ public void setCognitoIdentity(String cognitoIdentity) { this.cognitoIdentity = cognitoIdentity; } + public String getTenantId() { + return tenantId; + } + + public void setTenantId(String tenantId) { + this.tenantId = tenantId; + } + public byte[] getContent() { return content; } diff --git a/aws-lambda-java-runtime-interface-client/src/main/jni/com_amazonaws_services_lambda_runtime_api_client_runtimeapi_NativeClient.cpp b/aws-lambda-java-runtime-interface-client/src/main/jni/com_amazonaws_services_lambda_runtime_api_client_runtimeapi_NativeClient.cpp index 7fe47aa4..f0679661 100644 --- a/aws-lambda-java-runtime-interface-client/src/main/jni/com_amazonaws_services_lambda_runtime_api_client_runtimeapi_NativeClient.cpp +++ b/aws-lambda-java-runtime-interface-client/src/main/jni/com_amazonaws_services_lambda_runtime_api_client_runtimeapi_NativeClient.cpp @@ -20,6 +20,7 @@ static jfieldID contentField; static jfieldID clientContextField; static jfieldID cognitoIdentityField; static jfieldID xrayTraceIdField; +static jfieldID tenantIdField; jint JNI_OnLoad(JavaVM* vm, void* reserved) { @@ -41,6 +42,7 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) { xrayTraceIdField = env->GetFieldID(invocationRequestClass , "xrayTraceId", "Ljava/lang/String;"); clientContextField = env->GetFieldID(invocationRequestClass , "clientContext", "Ljava/lang/String;"); cognitoIdentityField = env->GetFieldID(invocationRequestClass , "cognitoIdentity", "Ljava/lang/String;"); + tenantIdField = env->GetFieldID(invocationRequestClass, "tenantId", "Ljava/lang/String;"); return JNI_VERSION; } @@ -106,6 +108,10 @@ JNIEXPORT jobject JNICALL Java_com_amazonaws_services_lambda_runtime_api_client_ CHECK_EXCEPTION(env, env->SetObjectField(invocationRequest, cognitoIdentityField, env->NewStringUTF(response.cognito_identity.c_str()))); } + if(response.tenant_id != ""){ + CHECK_EXCEPTION(env, env->SetObjectField(invocationRequest, tenantIdField, env->NewStringUTF(response.tenant_id.c_str()))); + } + bytes = reinterpret_cast(response.payload.c_str()); CHECK_EXCEPTION(env, jArray = env->NewByteArray(response.payload.length())); CHECK_EXCEPTION(env, env->SetByteArrayRegion(jArray, 0, response.payload.length(), bytes)); diff --git a/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/include/aws/lambda-runtime/runtime.h b/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/include/aws/lambda-runtime/runtime.h index 94e1e22c..d7db5f18 100644 --- a/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/include/aws/lambda-runtime/runtime.h +++ b/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/include/aws/lambda-runtime/runtime.h @@ -61,6 +61,11 @@ struct invocation_request { */ std::chrono::time_point deadline; + /** + * Tenant ID of the current invocation. + */ + std::string tenant_id; + /** * The number of milliseconds left before lambda terminates the current execution. */ diff --git a/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/src/runtime.cpp b/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/src/runtime.cpp index 91750840..eeaf0e7b 100644 --- a/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/src/runtime.cpp +++ b/aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/src/runtime.cpp @@ -40,6 +40,7 @@ static constexpr auto CLIENT_CONTEXT_HEADER = "lambda-runtime-client-context"; static constexpr auto COGNITO_IDENTITY_HEADER = "lambda-runtime-cognito-identity"; static constexpr auto DEADLINE_MS_HEADER = "lambda-runtime-deadline-ms"; static constexpr auto FUNCTION_ARN_HEADER = "lambda-runtime-invoked-function-arn"; +static constexpr auto TENANT_ID_HEADER = "lambda-runtime-aws-tenant-id"; enum Endpoints { INIT, @@ -301,6 +302,10 @@ runtime::next_outcome runtime::get_next() req.payload.c_str(), static_cast(req.get_time_remaining().count())); } + + if (resp.has_header(TENANT_ID_HEADER)) { + req.tenant_id = resp.get_header(TENANT_ID_HEADER); + } return next_outcome(req); } diff --git a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContextTest.java b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContextTest.java index 19744dd5..58880be4 100644 --- a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContextTest.java +++ b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContextTest.java @@ -18,6 +18,7 @@ public class LambdaContextTest { private static final String INVOKED_FUNCTION_ARN = "invoked-function-arn"; private static final LambdaClientContext CLIENT_CONTEXT = new LambdaClientContext(); public static final int MEMORY_LIMIT = 128; + public static final String TENANT_ID = "tenant-id"; @Test public void getRemainingTimeInMillis() { @@ -54,6 +55,6 @@ public void getRemainingTimeInMillis_Deadline() throws InterruptedException { private LambdaContext createContextWithDeadline(long deadlineTimeInMs) { return new LambdaContext(MEMORY_LIMIT, deadlineTimeInMs, REQUEST_ID, LOG_GROUP_NAME, LOG_STREAM_NAME, - FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, CLIENT_CONTEXT); + FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, TENANT_ID, CLIENT_CONTEXT); } } diff --git a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatterTest.java b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatterTest.java index 8630d5fe..531e9ca9 100644 --- a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatterTest.java +++ b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatterTest.java @@ -29,6 +29,25 @@ void testFormattingWithLambdaContext() { null, null, "function-arn", + null, + null + ); + assertFormatsString("test log", LogLevel.WARN, context); + } + + @Test + void testFormattingWithTenantIdInLambdaContext() { + LambdaContext context = new LambdaContext( + 0, + 0, + "request-id", + null, + null, + "function-name", + null, + null, + "function-arn", + "tenant-id", null ); assertFormatsString("test log", LogLevel.WARN, context); @@ -52,6 +71,7 @@ void assert_expected_log_message(StructuredLogMessage result, String message, Lo if (context != null) { assertEquals(context.getAwsRequestId(), result.AWSRequestId); + assertEquals(context.getTenantId(), result.tenantId); } } } diff --git a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/LambdaRuntimeApiClientImplTest.java b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/LambdaRuntimeApiClientImplTest.java index 1b6f3136..473e2aef 100644 --- a/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/LambdaRuntimeApiClientImplTest.java +++ b/aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/LambdaRuntimeApiClientImplTest.java @@ -15,6 +15,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.ErrorRequest; +import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.InvocationRequest; import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.StackElement; import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.XRayErrorCause; import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.XRayException; @@ -312,27 +313,62 @@ public void restoreNextWrongStatusCodeTest() { } @Test - public void nextTest() { + public void nextWithoutTenantIdHeaderTest() { try { - MockResponse mockResponse = new MockResponse(); - mockResponse.setResponseCode(HTTP_ACCEPTED); - mockResponse.setHeader("lambda-runtime-aws-request-id", "1234567890"); - mockResponse.setHeader("Content-Type", "application/json"); + MockResponse mockResponse = buildMockResponseForNextInvocation(); mockWebServer.enqueue(mockResponse); - lambdaRuntimeApiClientImpl.nextInvocation(); - RecordedRequest recordedRequest = mockWebServer.takeRequest(); - HttpUrl actualUrl = recordedRequest.getRequestUrl(); - String expectedUrl = "http://" + getHostnamePort() + "/2018-06-01/runtime/invocation/next"; - assertEquals(expectedUrl, actualUrl.toString()); + InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation(); + verifyNextInvocationRequest(); + assertNull(invocationRequest.getTenantId()); + } catch(Exception e) { + fail(); + } + } + + @Test + public void nextWithTenantIdHeaderTest() { + try { + MockResponse mockResponse = buildMockResponseForNextInvocation(); + String expectedTenantId = "my-tenant-id"; + mockResponse.setHeader("lambda-runtime-aws-tenant-id", expectedTenantId); + mockWebServer.enqueue(mockResponse); + + InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation(); + verifyNextInvocationRequest(); + assertEquals(expectedTenantId, invocationRequest.getTenantId()); - String actualBody = recordedRequest.getBody().readUtf8(); - assertEquals("", actualBody); } catch(Exception e) { fail(); } } + @Test + public void nextWithEmptyTenantIdHeaderTest() { + try { + MockResponse mockResponse = buildMockResponseForNextInvocation(); + mockResponse.setHeader("lambda-runtime-aws-tenant-id", ""); + mockWebServer.enqueue(mockResponse); + + InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation(); + verifyNextInvocationRequest(); + assertNull(invocationRequest.getTenantId()); + } catch(Exception e) { + fail(); + } + } + + @Test + public void nextWithNullTenantIdHeaderTest() { + try { + MockResponse mockResponse = buildMockResponseForNextInvocation(); + assertThrows(NullPointerException.class, () -> { + mockResponse.setHeader("lambda-runtime-aws-tenant-id", null); + }); + } catch(Exception e) { + fail(); + } + } @Test public void createUrlMalformedTest() { @@ -376,6 +412,24 @@ public void lambdaReportErrorXRayHeaderTooLongTest() { } } + private MockResponse buildMockResponseForNextInvocation() { + MockResponse mockResponse = new MockResponse(); + mockResponse.setResponseCode(HTTP_ACCEPTED); + mockResponse.setHeader("lambda-runtime-aws-request-id", "1234567890"); + mockResponse.setHeader("Content-Type", "application/json"); + return mockResponse; + } + + private void verifyNextInvocationRequest() throws Exception { + RecordedRequest recordedRequest = mockWebServer.takeRequest(); + HttpUrl actualUrl = recordedRequest.getRequestUrl(); + String expectedUrl = "http://" + getHostnamePort() + "/2018-06-01/runtime/invocation/next"; + assertEquals(expectedUrl, actualUrl.toString()); + + String actualBody = recordedRequest.getBody().readUtf8(); + assertEquals("", actualBody); + } + private String getHostnamePort() { return mockWebServer.getHostName() + ":" + mockWebServer.getPort(); }