diff --git a/powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java b/powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java index 9fb021548..47c4e7422 100644 --- a/powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java +++ b/powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java @@ -21,12 +21,15 @@ import com.amazonaws.xray.entities.Subsegment; import com.fasterxml.jackson.databind.ObjectMapper; import java.util.function.Consumer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * A class of helper functions to add additional functionality and ease * of use. */ public final class TracingUtils { + private static final Logger LOG = LoggerFactory.getLogger(TracingUtils.class); private static ObjectMapper objectMapper; /** @@ -36,6 +39,10 @@ public final class TracingUtils { * @param value the value of the annotation */ public static void putAnnotation(String key, String value) { + if (!isValidAnnotationKey(key)) { + LOG.warn("Ignoring annotation with unsupported characters in key: {}", key); + return; + } AWSXRay.getCurrentSubsegmentOptional() .ifPresent(segment -> segment.putAnnotation(key, value)); } @@ -47,10 +54,24 @@ public static void putAnnotation(String key, String value) { * @param value the value of the annotation */ public static void putAnnotation(String key, Number value) { + if (!isValidAnnotationKey(key)) { + LOG.warn("Ignoring annotation with unsupported characters in key: {}", key); + return; + } AWSXRay.getCurrentSubsegmentOptional() .ifPresent(segment -> segment.putAnnotation(key, value)); } + /** + Make sure that the annotation key is valid according to + the documentation. + + Annotation keys that are added that are invalid are ignored by x-ray. + **/ + private static boolean isValidAnnotationKey(String key) { + return key.matches("^[a-zA-Z0-9_]+$"); + } + /** * Put an annotation to the current subsegment with a Boolean value. * diff --git a/powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java b/powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java index 78283fbc2..01f25f37a 100644 --- a/powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java +++ b/powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java @@ -28,7 +28,6 @@ import org.junit.jupiter.api.Test; class TracingUtilsTest { - @BeforeEach void setUp() { AWSXRay.beginSegment("test"); @@ -123,6 +122,24 @@ void shouldInvokeCodeBlockWrappedWithinSubsegment() { }); } + @Test + void shouldNotAddAnnotationIfInvalidCharacterInKey() { + AWSXRay.beginSubsegment("subSegment"); + String inputKey = "stringKey with spaces"; + TracingUtils.putAnnotation(inputKey, "val"); + AWSXRay.getCurrentSubsegmentOptional() + .ifPresent(segment -> assertThat(segment.getAnnotations()).size().isEqualTo(0)); + } + + @Test + void shouldAddAnnotationIfValidCharactersInKey() { + AWSXRay.beginSubsegment("subSegment"); + String inputKey = "validKey"; + TracingUtils.putAnnotation(inputKey, "val"); + AWSXRay.getCurrentSubsegmentOptional() + .ifPresent(segment -> assertThat(segment.getAnnotations()).size().isEqualTo(1)); + } + @Test void shouldInvokeCodeBlockWrappedWithinNamespacedSubsegment() { Context test = mock(Context.class); @@ -221,4 +238,4 @@ void shouldInvokeCodeBlockWrappedWithinNamespacedEntitySubsegment() throws Inter .containsEntry("key", "val"); }); } -} \ No newline at end of file +}