From 17ac6352e4458366d9d030b95c7d569e6ca4ba6a Mon Sep 17 00:00:00 2001
From: Artem Bilan <artem.bilan@broadcom.com>
Date: Mon, 4 Nov 2024 16:25:24 -0500
Subject: [PATCH] Fix `FunctionConfiguration.sanitize()` for generics

The `FunctionConfiguration.sanitize()` use a `rawtypes` trying to hide
generics routines.

Spring Integration 6.4 has introduced a `BaseMessageBuilder` which does not work
with `rawtypes`.

* Fix `FunctionConfiguration.sanitize()` to expose a `<P>` generic argument
to satisfy a new `MessageBuilder` byte code, plus to meet all the expectations
of this method consumers
* In addition perform some code clean to optimize a usage of `MessageBuilder`:
 - `setHeader()` does set the value overriding existing one
 - `setHeader()` removes the entry if provided value is `null`
 - This way we just can go ahead with `MessageBuilder` method chain avoiding extra conditions
---
 .../stream/function/FunctionConfiguration.java | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/FunctionConfiguration.java b/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/FunctionConfiguration.java
index 7d23f7d75..98c10ebb6 100644
--- a/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/FunctionConfiguration.java
+++ b/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/FunctionConfiguration.java
@@ -389,17 +389,13 @@ public class FunctionConfiguration {
 						: MessageBuilder.withPayload(value).build();
 	}
 
-	@SuppressWarnings({ "unchecked", "rawtypes" })
-	private static Message sanitize(Message inputMessage) {
-		MessageBuilder builder = MessageBuilder
-				.fromMessage(inputMessage)
-				.removeHeader("spring.cloud.stream.sendto.destination")
-				.removeHeader(MessageUtils.SOURCE_TYPE);
-		if (builder.getHeaders().containsKey(MessageUtils.TARGET_PROTOCOL)) {
-			builder = builder.setHeader(MessageUtils.SOURCE_TYPE, builder.getHeaders().get(MessageUtils.TARGET_PROTOCOL));
-		}
-		builder = builder.removeHeader(MessageUtils.TARGET_PROTOCOL);
-		return builder.build();
+	private static <P> Message<P> sanitize(Message<P> inputMessage) {
+		return MessageBuilder
+			.fromMessage(inputMessage)
+			.removeHeader("spring.cloud.stream.sendto.destination")
+			.setHeader(MessageUtils.SOURCE_TYPE, inputMessage.getHeaders().get(MessageUtils.TARGET_PROTOCOL))
+			.removeHeader(MessageUtils.TARGET_PROTOCOL)
+			.build();
 	}
 
 	private static class FunctionToDestinationBinder implements InitializingBean, ApplicationContextAware {