From 8901c7b5e17880336baa0fb04b1e690ed3cfaef7 Mon Sep 17 00:00:00 2001 From: Anton Barkan Date: Tue, 26 Dec 2023 22:45:35 +0200 Subject: [PATCH] Add initial support for unwrapping reactive respponses during AWS processing --- .../function/adapter/aws/AWSLambdaUtils.java | 19 +++++----- .../adapter/aws/FunctionInvokerTests.java | 35 +++++++++++++++++++ 2 files changed, 43 insertions(+), 11 deletions(-) diff --git a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/AWSLambdaUtils.java b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/AWSLambdaUtils.java index 8bba1103a..9d8d91213 100644 --- a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/AWSLambdaUtils.java +++ b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/AWSLambdaUtils.java @@ -70,12 +70,13 @@ public final class AWSLambdaUtils { } - static boolean isSupportedAWSType(Type inputType) { - if (FunctionTypeUtils.isMessage(inputType) || FunctionTypeUtils.isPublisher(inputType)) { - inputType = FunctionTypeUtils.getImmediateGenericType(inputType, 0); + static boolean isSupportedAWSType(Type type) { + if (FunctionTypeUtils.isMessage(type) || FunctionTypeUtils.isPublisher(type)) { + type = FunctionTypeUtils.getImmediateGenericType(type, 0); } - return FunctionTypeUtils.getRawType(inputType).getPackage() != null && - FunctionTypeUtils.getRawType(inputType).getPackage().getName().startsWith( + Class rawType = FunctionTypeUtils.getRawType(type); + return rawType != null && rawType.getPackage() != null && + rawType.getPackage().getName().startsWith( "com.amazonaws.services.lambda.runtime.events"); } @@ -203,12 +204,8 @@ public final class AWSLambdaUtils { public static byte[] generateOutput(Message requestMessage, Message responseMessage, JsonMapper objectMapper, Type functionOutputType) { - Class outputClass = FunctionTypeUtils.getRawType(functionOutputType); - if (outputClass != null) { - String outputClassName = outputClass.getName(); - if (outputClassName.startsWith("com.amazonaws.services.lambda.runtime.events.")) { - return extractPayload((Message) responseMessage, objectMapper); - } + if (isSupportedAWSType(functionOutputType)) { + return extractPayload((Message) responseMessage, objectMapper); } byte[] responseBytes = responseMessage == null ? "\"OK\"".getBytes() : extractPayload((Message) responseMessage, objectMapper); diff --git a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/test/java/org/springframework/cloud/function/adapter/aws/FunctionInvokerTests.java b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/test/java/org/springframework/cloud/function/adapter/aws/FunctionInvokerTests.java index dc5a253ef..b91d88496 100644 --- a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/test/java/org/springframework/cloud/function/adapter/aws/FunctionInvokerTests.java +++ b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/test/java/org/springframework/cloud/function/adapter/aws/FunctionInvokerTests.java @@ -24,6 +24,7 @@ import java.io.OutputStream; import java.nio.charset.StandardCharsets; import java.util.Base64; import java.util.Collections; +import java.util.List; import java.util.Map; import java.util.function.Consumer; import java.util.function.Function; @@ -38,6 +39,7 @@ import com.amazonaws.services.lambda.runtime.events.APIGatewayV2HTTPResponse; import com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerRequestEvent; import com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerResponseEvent; import com.amazonaws.services.lambda.runtime.events.DynamodbEvent; +import com.amazonaws.services.lambda.runtime.events.IamPolicyResponse; import com.amazonaws.services.lambda.runtime.events.KinesisEvent; import com.amazonaws.services.lambda.runtime.events.S3Event; import com.amazonaws.services.lambda.runtime.events.SNSEvent; @@ -1347,6 +1349,21 @@ public class FunctionInvokerTests { assertThat(result.get("body")).isEqualTo("\"hello\""); } + @Test + public void testShouldNotWrapIamPolicyResponse() throws Exception { + System.setProperty("MAIN_CLASS", ApiGatewayConfiguration.class.getName()); + System.setProperty("spring.cloud.function.definition", "outputPolicyResponse"); + FunctionInvoker invoker = new FunctionInvoker(); + + InputStream targetStream = new ByteArrayInputStream(this.apiGatewayEvent.getBytes()); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + invoker.handleRequest(targetStream, output, null); + + Map result = mapper.readValue(output.toByteArray(), Map.class); + assertThat(result.get("body")).isNull(); + assertThat(result.get("principalId")).isNotNull(); + } + @SuppressWarnings("rawtypes") @Test public void testApiGatewayEventConsumer() throws Exception { @@ -1829,6 +1846,24 @@ public class FunctionInvokerTests { return body; }; } + + @Bean + public Function, Mono> outputPolicyResponse() { + return input -> + input.map(v -> IamPolicyResponse.builder() + .withPrincipalId("principalId") + .withPolicyDocument(IamPolicyResponse.PolicyDocument.builder() + .withVersion("2012-10-17") + .withStatement( + List.of( + IamPolicyResponse.Statement.builder().withAction("execute-api:Invoke") + .withResource( + List.of(v)).withEffect("Allow").build() + ) + ).build() + ).build() + ); + } } @EnableAutoConfiguration