From ab82f30623f0e577a3d3d84a6882bd762a352ecf Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin@grzejszczak.pl>
Date: Tue, 2 Apr 2019 14:10:36 +0200
Subject: [PATCH] Added a workaround for the mutability of native headers
 (#1323)

Since for some reason, the native headers sometimes are immutable even though the accessor says that the headers are mutable, then we have to ensure their
	 * mutability. We do so by first making a mutable copy of the native headers, then by removing the native headers from the headers map and replacing them with a mutable copy

fixes #1184
---
 .../messaging/MessageHeaderPropagation.java   | 28 +++++++++++++++++--
 .../client/HttpClientBeanPostProcessor.java   |  3 +-
 src/checkstyle/checkstyle-suppressions.xml    |  1 +
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagation.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagation.java
index 2ac83145c..13ec0e0f8 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagation.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagation.java
@@ -83,7 +83,11 @@ enum MessageHeaderPropagation
 			accessor.removeHeader(keyToRemove);
 			if (accessor instanceof NativeMessageHeaderAccessor) {
 				NativeMessageHeaderAccessor nativeAccessor = (NativeMessageHeaderAccessor) accessor;
-				nativeAccessor.removeNativeHeader(keyToRemove);
+				if (accessor.isMutable()) {
+					// 1184 native headers can be an immutable map
+					ensureNativeHeadersAreMutable(nativeAccessor)
+							.removeNativeHeader(keyToRemove);
+				}
 			}
 			else {
 				Object nativeHeaders = accessor
@@ -95,6 +99,26 @@ enum MessageHeaderPropagation
 		}
 	}
 
+	/**
+	 * Since for some reason, the native headers sometimes are immutable even though the
+	 * accessor says that the headers are mutable, then we have to ensure their
+	 * mutability. We do so by first making a mutable copy of the native headers, then by
+	 * removing the native headers from the headers map and replacing them with a mutable
+	 * copy. Workaround for #1184
+	 * @param nativeAccessor accessor containing (or not) native headers
+	 * @return modified accessor
+	 */
+	private static NativeMessageHeaderAccessor ensureNativeHeadersAreMutable(
+			NativeMessageHeaderAccessor nativeAccessor) {
+		Map<String, List<String>> nativeHeaderMap = nativeAccessor.toNativeHeaderMap();
+		nativeHeaderMap = nativeHeaderMap instanceof LinkedMultiValueMap ? nativeHeaderMap
+				: new LinkedMultiValueMap<>(nativeHeaderMap);
+		nativeAccessor.removeHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS);
+		nativeAccessor.setHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS,
+				nativeHeaderMap);
+		return nativeAccessor;
+	}
+
 	@Override
 	public void put(MessageHeaderAccessor accessor, String key, String value) {
 		try {
@@ -116,7 +140,7 @@ enum MessageHeaderPropagation
 		accessor.setHeader(key, value);
 		if (accessor instanceof NativeMessageHeaderAccessor) {
 			NativeMessageHeaderAccessor nativeAccessor = (NativeMessageHeaderAccessor) accessor;
-			nativeAccessor.setNativeHeader(key, value);
+			ensureNativeHeadersAreMutable(nativeAccessor).setNativeHeader(key, value);
 		}
 		else {
 			Object nativeHeaders = accessor
diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java
index 9cf640e9a..096b4c9d8 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java
@@ -161,7 +161,8 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			AtomicReference reference = req.currentContext()
 					.getOrDefault(AtomicReference.class, new AtomicReference());
 			Span span = handler().handleSend(injector(), req.requestHeaders(), req,
-					reference.get() == null ? handler().nextSpan(req) : (Span) reference.get());
+					reference.get() == null ? handler().nextSpan(req)
+							: (Span) reference.get());
 			reference.set(span);
 		}
 
diff --git a/src/checkstyle/checkstyle-suppressions.xml b/src/checkstyle/checkstyle-suppressions.xml
index 264dcd341..228a2c038 100644
--- a/src/checkstyle/checkstyle-suppressions.xml
+++ b/src/checkstyle/checkstyle-suppressions.xml
@@ -3,6 +3,7 @@
 	"-//Puppy Crawl//DTD Suppressions 1.1//EN"
 	"https://www.puppycrawl.com/dtds/suppressions_1_1.dtd">
 <suppressions>
+	<suppress files=".*/test/.*" checks="JavadocVariable"/>
 	<suppress files=".*FinishedSpanHandlerTests.*" checks="LineLengthCheck"/>
 	<suppress files=".*GrpcTracingIntegrationTests.*" checks="LineLengthCheck"/>
 	<suppress files=".*IgnoreAutoConfiguredSkipPatternsIntegrationTests.*" checks="LineLengthCheck"/>