From 3255ae7455102c0282ae71897e180cbd067d0db3 Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Wed, 10 Mar 2021 15:25:50 +0100
Subject: [PATCH] Fix header propagation for rsocket message

---
 .../cloud/function/rsocket/RSocketListenerFunction.java         | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/spring-cloud-function-rsocket/src/main/java/org/springframework/cloud/function/rsocket/RSocketListenerFunction.java b/spring-cloud-function-rsocket/src/main/java/org/springframework/cloud/function/rsocket/RSocketListenerFunction.java
index 09c0af0df..b9a137346 100644
--- a/spring-cloud-function-rsocket/src/main/java/org/springframework/cloud/function/rsocket/RSocketListenerFunction.java
+++ b/spring-cloud-function-rsocket/src/main/java/org/springframework/cloud/function/rsocket/RSocketListenerFunction.java
@@ -98,7 +98,7 @@ class RSocketListenerFunction implements Function<Message<Flux<Object>>, Publish
 			messageToProcess.getPayload()
 				.map((payload) -> {
 					return payload instanceof Message
-							? (Message) payload
+							? MessageBuilder.fromMessage((Message<?>) payload).copyHeadersIfAbsent(messageToProcess.getHeaders()).build()
 							: MessageBuilder.withPayload(payload).copyHeadersIfAbsent(messageToProcess.getHeaders()).build();
 				});
 		if (this.targetFunction.getInputType() != null && FunctionTypeUtils.isPublisher(this.targetFunction.getInputType())) {