From 11f6c4e7ab6b63f363aca63230aef111f476484e Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Thu, 26 Nov 2020 21:51:25 +0000
Subject: [PATCH] Matching update for copyHeadersIfAbsent

The change to copyHeaders from d46091 also applied to copyHeadersIfPresent.

Closes gh-26155
---
 .../support/NativeMessageHeaderAccessor.java  | 15 +++++++++
 .../NativeMessageHeaderAccessorTests.java     | 33 ++++++++++++++-----
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/spring-messaging/src/main/java/org/springframework/messaging/support/NativeMessageHeaderAccessor.java b/spring-messaging/src/main/java/org/springframework/messaging/support/NativeMessageHeaderAccessor.java
index 25fdb60fa73..2a3607fde2b 100644
--- a/spring-messaging/src/main/java/org/springframework/messaging/support/NativeMessageHeaderAccessor.java
+++ b/spring-messaging/src/main/java/org/springframework/messaging/support/NativeMessageHeaderAccessor.java
@@ -129,6 +129,21 @@ public class NativeMessageHeaderAccessor extends MessageHeaderAccessor {
 		super.copyHeaders(headersToCopy);
 	}
 
+	@Override
+	public void copyHeadersIfAbsent(@Nullable Map<String, ?> headersToCopy) {
+		if (headersToCopy == null) {
+			return;
+		}
+
+		@SuppressWarnings("unchecked")
+		Map<String, List<String>> map = (Map<String, List<String>>) headersToCopy.get(NATIVE_HEADERS);
+		if (map != null && getNativeHeaders() == null) {
+			map.forEach(this::setNativeHeaderValues);
+		}
+
+		super.copyHeadersIfAbsent(headersToCopy);
+	}
+
 	/**
 	 * Whether the native header map contains the give header name.
 	 * @param headerName the name of the header
diff --git a/spring-messaging/src/test/java/org/springframework/messaging/support/NativeMessageHeaderAccessorTests.java b/spring-messaging/src/test/java/org/springframework/messaging/support/NativeMessageHeaderAccessorTests.java
index 4d0e7b4b672..73eba06caa4 100644
--- a/spring-messaging/src/test/java/org/springframework/messaging/support/NativeMessageHeaderAccessorTests.java
+++ b/spring-messaging/src/test/java/org/springframework/messaging/support/NativeMessageHeaderAccessorTests.java
@@ -226,16 +226,33 @@ public class NativeMessageHeaderAccessorTests {
 
 	@Test // gh-25821
 	void copyImmutableToMutable() {
-		NativeMessageHeaderAccessor source = new NativeMessageHeaderAccessor();
-		source.addNativeHeader("foo", "bar");
-		Message<String> message = MessageBuilder.createMessage("payload", source.getMessageHeaders());
+		NativeMessageHeaderAccessor sourceAccessor = new NativeMessageHeaderAccessor();
+		sourceAccessor.addNativeHeader("foo", "bar");
+		Message<String> source = MessageBuilder.createMessage("payload", sourceAccessor.getMessageHeaders());
 
-		NativeMessageHeaderAccessor target = new NativeMessageHeaderAccessor();
-		target.copyHeaders(message.getHeaders());
-		target.setLeaveMutable(true);
-		message = MessageBuilder.createMessage(message.getPayload(), target.getMessageHeaders());
+		NativeMessageHeaderAccessor targetAccessor = new NativeMessageHeaderAccessor();
+		targetAccessor.copyHeaders(source.getHeaders());
+		targetAccessor.setLeaveMutable(true);
+		Message<?> target = MessageBuilder.createMessage(source.getPayload(), targetAccessor.getMessageHeaders());
 
-		MessageHeaderAccessor accessor = MessageHeaderAccessor.getMutableAccessor(message);
+		MessageHeaderAccessor accessor = MessageHeaderAccessor.getMutableAccessor(target);
+		assertThat(accessor.isMutable());
+		((NativeMessageHeaderAccessor) accessor).addNativeHeader("foo", "baz");
+		assertThat(((NativeMessageHeaderAccessor) accessor).getNativeHeader("foo")).containsExactly("bar", "baz");
+	}
+
+	@Test // gh-25821
+	void copyIfAbsentImmutableToMutable() {
+		NativeMessageHeaderAccessor sourceAccessor = new NativeMessageHeaderAccessor();
+		sourceAccessor.addNativeHeader("foo", "bar");
+		Message<String> source = MessageBuilder.createMessage("payload", sourceAccessor.getMessageHeaders());
+
+		MessageHeaderAccessor targetAccessor = new NativeMessageHeaderAccessor();
+		targetAccessor.copyHeadersIfAbsent(source.getHeaders());
+		targetAccessor.setLeaveMutable(true);
+		Message<?> target = MessageBuilder.createMessage(source.getPayload(), targetAccessor.getMessageHeaders());
+
+		MessageHeaderAccessor accessor = MessageHeaderAccessor.getMutableAccessor(target);
 		assertThat(accessor.isMutable());
 		((NativeMessageHeaderAccessor) accessor).addNativeHeader("foo", "baz");
 		assertThat(((NativeMessageHeaderAccessor) accessor).getNativeHeader("foo")).containsExactly("bar", "baz");