From e02cea7bc00dc769203f67613614671f9065838a Mon Sep 17 00:00:00 2001 From: Andreas Berger Date: Tue, 4 Sep 2018 11:06:53 +0200 Subject: [PATCH] Fixes #1081: add instanceOf checks (#1082) * Fixes #1081: add instance of checks --- .../messaging/MessageHeaderPropagation.java | 37 +++++++++---------- .../MessageHeaderPropagationTest.java | 28 ++++++++++++++ 2 files changed, 46 insertions(+), 19 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 3feedf52e..3096fb865 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 @@ -79,13 +79,14 @@ enum MessageHeaderPropagation nativeAccessor.setNativeHeader(key, value); } else { - Map> nativeHeaders = (Map) accessor - .getHeader(NATIVE_HEADERS); + Object nativeHeaders = accessor.getHeader(NATIVE_HEADERS); if (nativeHeaders == null) { accessor.setHeader(NATIVE_HEADERS, nativeHeaders = new LinkedMultiValueMap<>()); } - nativeHeaders.put(key, Collections.singletonList(value)); + if (nativeHeaders instanceof Map) { + ((Map) nativeHeaders).put(key, Collections.singletonList(value)); + } } } @@ -115,16 +116,16 @@ enum MessageHeaderPropagation if (accessor instanceof NativeMessageHeaderAccessor) { NativeMessageHeaderAccessor nativeAccessor = (NativeMessageHeaderAccessor) accessor; String result = nativeAccessor.getFirstNativeHeader(key); - if (result != null) + if (result != null) { return result; - } - else { - Map> nativeHeaders = (Map) accessor - .getHeader(NATIVE_HEADERS); - if (nativeHeaders != null) { - List result = nativeHeaders.get(key); - if (result != null && !result.isEmpty()) - return result.get(0); + } + } else { + Object nativeHeaders = accessor.getHeader(NATIVE_HEADERS); + if (nativeHeaders instanceof Map) { + Object result = ((Map) nativeHeaders).get(key); + if (result instanceof List && !((List) result).isEmpty()) { + return String.valueOf(((List) result).get(0)); + } } } Object result = accessor.getHeader(key); @@ -155,13 +156,11 @@ enum MessageHeaderPropagation if (accessor instanceof NativeMessageHeaderAccessor) { NativeMessageHeaderAccessor nativeAccessor = (NativeMessageHeaderAccessor) accessor; nativeAccessor.removeNativeHeader(keyToRemove); - } - else { - Map> nativeHeaders = (Map) accessor - .getHeader(NATIVE_HEADERS); - if (nativeHeaders == null) - continue; - nativeHeaders.remove(keyToRemove); + } else { + Object nativeHeaders = accessor.getHeader(NATIVE_HEADERS); + if (nativeHeaders instanceof Map) { + ((Map) nativeHeaders).remove(keyToRemove); + } } } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagationTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagationTest.java index 0f212a33e..939552afc 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagationTest.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagationTest.java @@ -21,6 +21,7 @@ import java.util.Collections; import brave.propagation.Propagation; import org.junit.Test; import org.springframework.messaging.support.MessageHeaderAccessor; +import org.springframework.messaging.support.NativeMessageHeaderAccessor; import static org.junit.Assert.*; @@ -73,4 +74,31 @@ public class MessageHeaderPropagationTest String value = MessageHeaderPropagation.INSTANCE.get(carrier, "non existent key"); assertNull(value); } + + @Test + public void testSkipWrongValueTypeForGet() { + MessageHeaderAccessor carrier = carrier(); + carrier.setHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS, + "{spanTraceId=[123], spanId=[456], spanSampled=[0]}" + ); + MessageHeaderPropagation.INSTANCE.get(carrier, "X-B3-SpanId"); + } + + @Test + public void testSkipWrongValueTypeForRemoval() { + MessageHeaderAccessor carrier = carrier(); + carrier.setHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS, + "{spanTraceId=[123], spanId=[456], spanSampled=[0]}" + ); + MessageHeaderPropagation.removeAnyTraceHeaders(carrier, Collections.singletonList("X-B3-SpanId")); + } + + @Test + public void testSkipWrongValueTypeForPut() { + MessageHeaderAccessor carrier = carrier(); + carrier.setHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS, + "{spanTraceId=[123], spanId=[456], spanSampled=[0]}" + ); + MessageHeaderPropagation.INSTANCE.put(carrier, "X-B3-SpanId", "1234"); + } }