From 8cae6f13b27dc9b750b7e64da4912d37c2a9fcd7 Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Tue, 29 Jun 2021 15:07:40 +0200
Subject: [PATCH] GH-2101 Fix StreamBridge when sending Message with native
 encoding

There was an always cast to FunctionInvocationWrapper even though in the case of nativeEncoding the function was not of that instance. This fixes it.

Resolves #2101
---
 .../PartitionAwareFunctionWrapper.java        | 20 +++++++++++--------
 .../cloud/stream/function/StreamBridge.java   |  5 ++---
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
index 28c8efce6..6cf0fb41f 100644
--- a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
+++ b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
@@ -41,16 +41,14 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 
 	protected final Log logger = LogFactory.getLog(PartitionAwareFunctionWrapper.class);
 
-	private final FunctionInvocationWrapper function;
-
-	private Function<Object, Message> enhancerFunction;
+	@SuppressWarnings("rawtypes")
+	private final Function function;
 
 	@SuppressWarnings("rawtypes")
 	private final Function<Object, Message> outputMessageEnricher;
 
-	PartitionAwareFunctionWrapper(FunctionInvocationWrapper function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
+	PartitionAwareFunctionWrapper(Function<?, ?> function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
 		this.function = function;
-		this.enhancerFunction = function.getEnhancer();
 
 		if (producerProperties != null && producerProperties.isPartitioned()) {
 			StandardEvaluationContext evaluationContext = ExpressionUtils.createStandardEvaluationContext(context.getBeanFactory());
@@ -71,6 +69,7 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 		}
 	}
 
+	@SuppressWarnings("unchecked")
 	@Override
 	public Object apply(Object input) {
 		this.setEnhancerIfNecessary();
@@ -79,11 +78,16 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 
 	@Override
 	public Object get() {
-		this.setEnhancerIfNecessary();
-		return this.function.get();
+		if (this.function instanceof FunctionInvocationWrapper) {
+			this.setEnhancerIfNecessary();
+			return ((FunctionInvocationWrapper) this.function).get();
+		}
+		throw new IllegalStateException("Call to get() is not allowed since this function is not a Supplier.");
 	}
 
 	private void setEnhancerIfNecessary() {
-		this.function.setEnhancer(this.outputMessageEnricher);
+		if (this.function instanceof FunctionInvocationWrapper) {
+			((FunctionInvocationWrapper) this.function).setEnhancer(this.outputMessageEnricher);
+		}
 	}
 }
diff --git a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java
index 882ce6cf7..0e4ef8e13 100644
--- a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java
+++ b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java
@@ -29,7 +29,6 @@ import org.springframework.cloud.function.context.FunctionCatalog;
 import org.springframework.cloud.function.context.FunctionRegistration;
 import org.springframework.cloud.function.context.FunctionRegistry;
 import org.springframework.cloud.function.context.FunctionType;
-import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
 import org.springframework.cloud.stream.binder.Binder;
 import org.springframework.cloud.stream.binder.BinderFactory;
 import org.springframework.cloud.stream.binder.ProducerProperties;
@@ -203,11 +202,11 @@ public final class StreamBridge implements SmartInitializingSingleton {
 		boolean skipConversion = producerProperties.isUseNativeEncoding();
 
 		Function<Object, Object> functionToInvoke = skipConversion
-				? v -> v instanceof Message ? v :  MessageBuilder.withPayload(v).build()
+				? (v -> v instanceof Message ? v :  MessageBuilder.withPayload(v).build())
 						: this.functionCatalog.lookup(STREAM_BRIDGE_FUNC_NAME, outputContentType.toString());
 
 		if (producerProperties != null && producerProperties.isPartitioned()) {
-			functionToInvoke = new PartitionAwareFunctionWrapper((FunctionInvocationWrapper) functionToInvoke, this.applicationContext, producerProperties);
+			functionToInvoke = new PartitionAwareFunctionWrapper(functionToInvoke, this.applicationContext, producerProperties);
 		}
 		// this function is a pass through and is only required to force output conversion if necessary on SCF side.
 		Message<byte[]> resultMessage = (Message<byte[]>) functionToInvoke.apply(data);