fix(core): Fix tracestate encoding when Datadog and other members are…

… present using W3CPtags
DataDog · Nov 6, 2023 · 8dd68f2 · 8dd68f2
1 parent 737c08c
commit 8dd68f2
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 6 deletions.
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/W3CPTagsCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/W3CPTagsCodec.java
@@ -23,7 +23,7 @@ public class W3CPTagsCodec extends PTagsCodec {
   private static final char KEY_VALUE_SEPARATOR = ':';
   private static final int MIN_ALLOWED_CHAR = 32;
   private static final int MAX_ALLOWED_CHAR = 126;
-  private static final int MAX_LIST_MEMBER_COUNT = 32;
+  private static final int MAX_MEMBER_COUNT = 32;
 
   @Override
   PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
@@ -46,7 +46,7 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
     int memberIndex = 0;
     int ddMemberIndex = -1;
     while (memberStart < len) {
-      if (memberIndex == MAX_LIST_MEMBER_COUNT) {
+      if (memberIndex == MAX_MEMBER_COUNT) {
         // TODO should we return one with an error?
         // TODO should we try to pick up the `dd` member anyway?
         return tagsFactory.empty();
@@ -634,19 +634,22 @@ private static int cleanUpAndAppendSuffix(StringBuilder sb, PTags ptags, int siz
       return size;
     }
     int ddMemberStart = (ptags instanceof W3CPTags) ? ((W3CPTags) ptags).ddMemberStart : -1;
-    int remainingListMemberAllowed = size == 0 ? MAX_LIST_MEMBER_COUNT : MAX_LIST_MEMBER_COUNT - 1;
+    int remainingMemberAllowed = size == 0 ? MAX_MEMBER_COUNT : MAX_MEMBER_COUNT - 1;
     int len = original.length();
     int memberStart = findNextMember(original, 0);
     while (memberStart < len) {
+      // Look for member end position
       int memberEnd = original.indexOf(MEMBER_SEPARATOR, memberStart);
       if (memberEnd < 0) {
         memberEnd = len;
       }
+      // Try to define Datadog member start if not already found
       if (ddMemberStart == -1) {
         if (original.startsWith(DATADOG_MEMBER_KEY, memberStart)) {
           ddMemberStart = memberStart;
         }
       }
+      // Skip Datadog member (already added with prefix and tags)
       if (memberStart != ddMemberStart) {
         if (sb.length() > 0) {
           sb.append(MEMBER_SEPARATOR);
@@ -655,9 +658,10 @@ private static int cleanUpAndAppendSuffix(StringBuilder sb, PTags ptags, int siz
         int end = stripTrailingOWC(original, memberStart, memberEnd);
         sb.append(original, memberStart, end);
         size += (end - memberStart);
+        remainingMemberAllowed--;
       }
-      remainingListMemberAllowed--;
-      if (remainingListMemberAllowed == 0) {
+      // Check if remaining members are allowed
+      if (remainingMemberAllowed == 0) {
         memberStart = len;
       } else {
         memberStart = findNextMember(original, memberEnd + 1);

diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/W3CPropagationTagsTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/W3CPropagationTagsTest.groovy
@@ -145,7 +145,7 @@ class W3CPropagationTagsTest extends DDCoreSpecification {
     valueChar << (' '..'ÿ') - ((' '..'~') - [',', '='])
   }
 
-  def "validate tracestate header number of members #memberCount"() {
+  def "validate tracestate header number of members #memberCount without Datadog member"() {
     setup:
     def config = Mock(Config)
     config.getxDatadogTagsMaxLength() >> 512
@@ -168,6 +168,29 @@ class W3CPropagationTagsTest extends DDCoreSpecification {
     memberCount << (1..37) // some arbitrary number larger than 32
   }
 
+  def "validate tracestate header number of members #memberCount with Datadog member"() {
+    setup:
+    def config = Mock(Config)
+    config.getxDatadogTagsMaxLength() >> 512
+    def propagationTagsFactory = PropagationTags.factory(config)
+    def header = 'dd=s:1,'+(1..memberCount).collect { "k$it=v$it" }.join(',')
+
+    when:
+    def headerPT = propagationTagsFactory.fromHeaderValue(HeaderType.W3C, header)
+
+    then:
+    if (memberCount + 1 <= 32) {
+      assert headerPT.headerValue(HeaderType.W3C) == header
+    } else {
+      assert headerPT.headerValue(HeaderType.W3C) == null
+    }
+    // we're not using any dd members in the tests
+    headerPT.createTagMap() == [:]
+
+    where:
+    memberCount << (1..37) // some arbitrary number larger than 32
+  }
+
   def "validate tracestate header number of members #memberCount when propagating original tracestate"() {
     setup:
     def config = Mock(Config)