From 1787df8bb6a7ba9a98db8bb174862968ea272ea5 Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Wed, 1 May 2024 11:45:29 +0200
Subject: [PATCH] GH-1135 Fix AWS data conversion

Resolves #1135
---
 .../function/adapter/aws/AWSLambdaUtils.java  |  2 +-
 .../adapter/aws/AWSTypesMessageConverter.java |  6 +-
 .../adapter/aws/FunctionInvokerTests.java     | 56 +++++++++++++++++++
 3 files changed, 60 insertions(+), 4 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 4f27ad980..0fb568dfd 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
@@ -126,7 +126,7 @@ public final class AWSLambdaUtils {
 
 		MessageBuilder builder = MessageBuilder
 				.withPayload(structMessage instanceof Map msg && msg.containsKey("payload")
-						? ((String) msg.get("payload")).getBytes(StandardCharsets.UTF_8)
+						? (msg.get("payload"))
 						: payload);
 		if (isApiGateway) {
 			builder.setHeader(AWSLambdaUtils.AWS_API_GATEWAY, true);
diff --git a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/AWSTypesMessageConverter.java b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/AWSTypesMessageConverter.java
index f90f80dee..9c5c0bd39 100644
--- a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/AWSTypesMessageConverter.java
+++ b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/AWSTypesMessageConverter.java
@@ -89,11 +89,11 @@ class AWSTypesMessageConverter extends JsonMessageConverter {
 			}
 			else {
 				Object body;
-				if (message.getHeaders().containsKey("payload")) {
-					body = message.getPayload();
+				if (structMessage.containsKey("body")) {
+					body = structMessage.get("body");
 				}
 				else {
-					body = structMessage.get("body");
+					body = message.getPayload();
 				}
 				Object convertedResult = this.jsonMapper.fromJson(body, targetClass);
 				return convertedResult;
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 b91d88496..3db5bb6e4 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
@@ -78,6 +78,35 @@ public class FunctionInvokerTests {
 
 	String jsonPojoCollection = "[{\"name\":\"Ricky\"},{\"name\":\"Julien\"},{\"name\":\"Julien\"}]";
 
+	String someEvent = "{\n"
+			+ "    \"payload\": {\n"
+			+ "        \"headers\": {\n"
+			+ "            \"businessUnit\": \"1\"\n"
+			+ "        }\n"
+			+ "    },\n"
+			+ "    \"headers\": {\n"
+			+ "        \"aws-context\": {\n"
+			+ "            \"memoryLimit\": 1024,\n"
+			+ "            \"awsRequestId\": \"87a211bf-540f-4f9f-a218-d096a0099999\",\n"
+			+ "            \"functionName\": \"myfunction\",\n"
+			+ "            \"functionVersion\": \"278\",\n"
+			+ "            \"invokedFunctionArn\": \"arn:aws:lambda:us-east-1:xxxxxxx:function:xxxxx:snapstart\",\n"
+			+ "            \"deadlineTimeInMs\": 1712717704761,\n"
+			+ "            \"logger\": {\n"
+			+ "                \"logFiltering\": {\n"
+			+ "                    \"minimumLogLevel\": \"UNDEFINED\"\n"
+			+ "                },\n"
+			+ "                \"logFormatter\": {},\n"
+			+ "                \"logFormat\": \"TEXT\"\n"
+			+ "            }\n"
+			+ "        },\n"
+			+ "        \"businessUnit\": \"1\",\n"
+			+ "        \"id\": \"xxxx\",\n"
+			+ "        \"aws-event\": true,\n"
+			+ "        \"timestamp\": 1712716805129\n"
+			+ "    }\n"
+			+ "}";
+
 	String dynamoDbEvent = "{\n"
 			+ "  \"Records\": [\n"
 			+ "    {\n"
@@ -706,6 +735,22 @@ public class FunctionInvokerTests {
 		//this.getEnvironment().clear();
 	}
 
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	@Test
+	public void testConversionWhenPayloadExists() throws Exception {
+		System.setProperty("MAIN_CLASS", BasicConfiguration.class.getName());
+		System.setProperty("spring.cloud.function.definition", "uppercase");
+		FunctionInvoker invoker = new FunctionInvoker();
+
+		InputStream targetStream = new ByteArrayInputStream(this.someEvent.getBytes());
+		ByteArrayOutputStream output = new ByteArrayOutputStream();
+		invoker.handleRequest(targetStream, output, null);
+
+		Map result = mapper.readValue(output.toByteArray(), Map.class);
+		assertThat(result).containsKey("HEADERS");
+
+	}
+
 	@Test
 	public void testAPIGatewayCustomAuthorizerEvent() throws Exception {
 		System.setProperty("MAIN_CLASS", AuthorizerConfiguration.class.getName());
@@ -1441,6 +1486,17 @@ public class FunctionInvokerTests {
 		assertThat(result).isEqualTo(testString);
 	}
 
+	@EnableAutoConfiguration
+	@Configuration
+	public static class BasicConfiguration {
+		@Bean
+		public Function<Message<String>, Message<String>> uppercase() {
+			return v -> {
+				return MessageBuilder.withPayload(v.getPayload().toUpperCase()).build();
+			};
+		}
+	}
+
 	@EnableAutoConfiguration
 	@Configuration
 	public static class AuthorizerConfiguration {