adding validation for tracing annotation keys

Author: jdoherty

powertools-tracing/pom.xml

@@ -48,6 +48,10 @@
         </developer>
     </developers>
 
+    <properties>
+        <logback.version>1.2.11</logback.version>
+    </properties>
+
     <distributionManagement>
         <snapshotRepository>
             <id>ossrh</id>
@@ -115,6 +119,13 @@
             <artifactId>assertj-core</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>ch.qos.logback</groupId>
+            <artifactId>logback-classic</artifactId>
+            <version>${logback.version}</version>
+            <scope>test</scope>
+        </dependency>
+
     </dependencies>
 
     <build>

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,7 @@ public final class TracingUtils {
      * @param value the value of the annotation
      */
     public static void putAnnotation(String key, String value) {
+        validateAnnotationKey(key);
         AWSXRay.getCurrentSubsegmentOptional()
                 .ifPresent(segment -> segment.putAnnotation(key, value));
     }
@@ -47,10 +51,17 @@ public static void putAnnotation(String key, String value) {
      * @param value the value of the annotation
      */
     public static void putAnnotation(String key, Number value) {
+        validateAnnotationKey(key);
         AWSXRay.getCurrentSubsegmentOptional()
                 .ifPresent(segment -> segment.putAnnotation(key, value));
     }
 
+    private static void validateAnnotationKey(String key) {
+        if (!key.matches("^[a-zA-Z0-9_]+$")) {
+            LOG.warn("ignoring annotation with unsupported characters in key: {}", key);
+        }
+    }
+
     /**
      * Put an annotation to the current subsegment with a Boolean value.
      *

powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsTest.java

@@ -20,12 +20,18 @@
 import static org.mockito.Mockito.verify;
 import static software.amazon.lambda.powertools.tracing.TracingUtils.withEntitySubsegment;
 
+import ch.qos.logback.classic.Level;
+import ch.qos.logback.classic.Logger;
+import ch.qos.logback.classic.spi.ILoggingEvent;
+import ch.qos.logback.core.read.ListAppender;
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Entity;
+import java.util.List;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.slf4j.LoggerFactory;
 
 class TracingUtilsTest {
 
@@ -123,6 +129,48 @@ void shouldInvokeCodeBlockWrappedWithinSubsegment() {
                 });
     }
 
+    @Test
+    void shouldEmitNoLogWarnIfValidCharacterInKey() {
+        AWSXRay.beginSubsegment("subSegment");
+        Logger logger = (Logger) LoggerFactory.getLogger(TracingUtils.class);
+
+        // create and start a ListAppender
+        ListAppender<ILoggingEvent> listAppender = new ListAppender<>();
+        listAppender.start();
+
+        // add the appender to the logger
+        logger.addAppender(listAppender);
+
+        TracingUtils.putAnnotation("stringKey", "val");
+
+        List<ILoggingEvent> logsList = listAppender.list;
+        assertThat(AWSXRay.getTraceEntity().getAnnotations())
+                .hasSize(1)
+                .contains(
+                        entry("stringKey", "val")
+                );
+        assertThat(logsList.size()).isZero();
+    }
+
+    @Test
+    void shouldEmitLogWarnIfInvalidCharacterInKey() {
+        AWSXRay.beginSubsegment("subSegment");
+        Logger logger = (Logger) LoggerFactory.getLogger(TracingUtils.class);
+
+        // create and start a ListAppender
+        ListAppender<ILoggingEvent> listAppender = new ListAppender<>();
+        listAppender.start();
+
+        // add the appender to the logger
+        logger.addAppender(listAppender);
+        String inputKey = "stringKey with spaces";
+        TracingUtils.putAnnotation(inputKey, "val");
+
+        List<ILoggingEvent> logsList = listAppender.list;
+        assertThat(logsList.get(0).getLevel()).isEqualTo(Level.WARN);
+        assertThat(logsList.get(0).getMessage()).isEqualTo("ignoring annotation with unsupported characters in key: {}",inputKey);
+    }
+
     @Test
     void shouldInvokeCodeBlockWrappedWithinNamespacedSubsegment() {
         Context test = mock(Context.class);