From 0af24dfd6b25bc7d9e3c90654c22bfa03bac633b Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Wed, 24 Aug 2022 16:19:22 +0200
Subject: [PATCH] GH-919 Fix regression that was modifying result message

Resolves #919
---
 .../catalog/SimpleFunctionRegistry.java       |  5 -----
 ...BeanFactoryAwareFunctionRegistryTests.java | 20 +++++++++++++++++++
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java
index 4211b0af2..fe538ac30 100644
--- a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java
+++ b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java
@@ -782,11 +782,6 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 					if (functionInvocationHelper != null && CloudEventMessageUtils.isCloudEvent(((Message) input))) {
 						result = functionInvocationHelper.postProcessResult(result, (Message) input);
 					}
-					else {
-						Map<String, Object> headersMap = new HashMap<>(((Message) result).getHeaders());
-						this.sanitizeHeaders(((Message) input).getHeaders()).forEach((k, v) -> headersMap.putIfAbsent(k, v));
-						result = MessageBuilder.withPayload(((Message) result).getPayload()).copyHeaders(headersMap).build();
-					}
 				}
 				else {
 					if (functionInvocationHelper != null && CloudEventMessageUtils.isCloudEvent(((Message) input))) {
diff --git a/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java b/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java
index 2c1ac4beb..f8511d525 100644
--- a/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java
+++ b/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java
@@ -157,6 +157,17 @@ public class BeanFactoryAwareFunctionRegistryTests {
 		assertThat(c.size()).isEqualTo(2);
 	}
 
+	@Test
+	public void testReturnedMessageIsUnmodified() throws Exception {
+		FunctionCatalog catalog = this.configureCatalog();
+		Function<Message<String>, Message<byte[]>> function = catalog.lookup("uppercaseMessage", "application/json");
+		assertThat(function).isNotNull();
+
+		Message<byte[]> result = function.apply(MessageBuilder.withPayload("bob").setHeader("foo", "foo").build());
+		assertThat(result.getHeaders().containsKey("foo")).isFalse();
+		assertThat(result.getHeaders().containsKey("bar")).isTrue();
+	}
+
 	@SuppressWarnings("unchecked")
 	@Test
 	public void testDefaultLookup() throws Exception {
@@ -1074,6 +1085,15 @@ public class BeanFactoryAwareFunctionRegistryTests {
 			return v -> v.toUpperCase();
 		}
 
+		@Bean
+		public Function<Message<String>, Message<String>> uppercaseMessage() {
+			return message -> {
+				Message<String> result = MessageBuilder.fromMessage(message)
+					.removeHeader("foo").setHeader("bar", "bar").build();
+				return result;
+			};
+		}
+
 		@Bean
 		public Function<String, Void> consumerFunction() {
 			return v -> {