feat: add aws xray context propagation

packages/dd-trace/src/opentracing/propagation/text_map.js

@@ -46,6 +46,16 @@ const tracestateTagKeyFilter = /[^\x21-\x2b\x2d-\x3c\x3e-\x7e]/g
 const tracestateTagValueFilter = /[^\x20-\x2b\x2d-\x3a\x3c-\x7d]/g
 const invalidSegment = /^0+$/
 const zeroTraceId = '0000000000000000'
+// AWS X-Ray specific constants
+const xrayHeaderKey = 'x-amzn-trace-id'
+const xrayRootKey = 'root'
+const xrayRootPrefix = '1-'
+const xrayParentKey = 'parent'
+const xraySampledKey = 'sampled'
+const xrayE2EStartTimeKey = 't0'
+const xraySelfKey = 'self'
+const xrayOriginKey = '_dd.origin'
+const xrayMaxAdditionalBytes = 256
 
 class TextMapPropagator {
   constructor (config) {
@@ -60,6 +70,7 @@ class TextMapPropagator {
     this._injectB3MultipleHeaders(spanContext, carrier)
     this._injectB3SingleHeader(spanContext, carrier)
     this._injectTraceparent(spanContext, carrier)
+    this._injectAwsXrayContext(spanContext, carrier)
 
     if (injectCh.hasSubscribers) {
       injectCh.publish({ spanContext, carrier })
@@ -252,6 +263,66 @@ class TextMapPropagator {
     carrier.tracestate = ts.toString()
   }
 
+  _injectAwsXrayContext (spanContext, carrier) {
+    // injects AWS Trace Header (X-Amzn-Trace-Id) to carrier
+    // ex: 'Root=1-00000000-00000000fffffffffffffffe;Parent=ffffffffffffffff;Sampled=1;_dd.origin=fakeOrigin;
+
+    // based off: https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-tracingheader
+    if (!this._hasPropagationStyle('inject', 'xray')) return
+
+    // TODO: Do we have access to a start time? Java does
+    const e2eStart = this._getEndToEndStartTime(spanContext.start_ms)
+
+    let str = (
+      xrayRootKey + '=' +
+      xrayRootPrefix +
+      e2eStart + '-' +
+      spanContext._traceId.toString().padStart(24, '0') +
+      ';' + xrayParentKey + '=' +
+      spanContext._spanId.toString().padStart(16, '0')
+    )
+
+    if (spanContext?._sampling?.priority) {
+      str += ';' + xraySampledKey + '=' + (spanContext._sampling.priority > 0 ? '1' : '0')
+    }
+
+    const maxAdditionalCapacity = xrayMaxAdditionalBytes - str.length
+
+    const origin = spanContext._trace.origin
+    if (origin) {
+      this._addXrayBaggage(str, xrayOriginKey, origin, maxAdditionalCapacity)
+    }
+    if (e2eStart !== '00000000') {
+      this._addXrayBaggage(str, xrayE2EStartTimeKey, e2eStart.toString(), maxAdditionalCapacity)
+    }
+
+    for (const [key, value] of Object.entries(spanContext._baggageItems)) {
+      if (!this._isReservedXrayKey(key)) {
+        str = this._addXrayBaggage(str, key, value, maxAdditionalCapacity)
+      }
+    }
+
+    carrier[xrayHeaderKey] = str
+  }
+
+  _getEndToEndStartTime (start) {
+    if (!start) return '00000000'
+
+    const e2eStart = start > 0 ? start : Date.now() / 1000
+    return e2eStart.toString().padStart(8, '0')
+  }
+
+  _isReservedXrayKey (key) {
+    return [xrayRootKey, xrayParentKey, xrayOriginKey, xrayE2EStartTimeKey, xraySampledKey].includes(key.toLowerCase())
+  }
+
+  _addXrayBaggage (str, key, value, maxCapacity) {
+    if (str.length + key.length + value.toString().length + 2 <= maxCapacity) {
+      str += `;${key}=${value}`
+    }
+    return str
+  }
+
   _hasPropagationStyle (mode, name) {
     return this._config.tracePropagationStyle[mode].includes(name)
   }
@@ -303,6 +374,9 @@ class TextMapPropagator {
         case 'b3 single header': // TODO: delete in major after singular "b3"
           extractedContext = this._extractB3SingleContext(carrier)
           break
+        case 'xray':
+          extractedContext = this._extractAwsXrayContext(carrier)
+          break
         case 'b3':
           if (this._config.tracePropagationStyle.otelPropagators) {
             // TODO: should match "b3 single header" in next major
@@ -692,6 +766,93 @@ class TextMapPropagator {
     return spanContext._traceId.toString(16)
   }
 
+  _extractAwsXrayContext (carrier) {
+    if (carrier[xrayHeaderKey]) {
+      const parsedHeader = this._parseAWSTraceHeader(carrier[xrayHeaderKey])
+
+      let traceId
+      let spanId
+      let samplingPriority
+      let ddOrigin
+      const baggage = {}
+
+      if (!(
+        xrayRootKey in parsedHeader &&
+        // Regex check to ensure received header is in the same format as expected
+        // Format:
+        //   'Root=1-'
+        //   8 hexadecimal characters representing start time
+        //   '-'
+        //   24 hexadecimal characters representing trace id
+        //   ';'
+        //   'Parent='
+        //   16 hexadecimal characters representing parent span id
+        /^(?=.*Root=1-[0-9a-f]{8}-[0-9a-f]{24})(?=.*Parent=[0-9a-f]{16}).*$/i.test(carrier[xrayHeaderKey])
+      )) {
+        // header doesn't match formatting
+        return null
+      }
+
+      for (const key in parsedHeader) {
+        const value = parsedHeader[key]
+        if (key === xrayRootKey) {
+          const awsTraceIdSegments = value.split('-')
+          for (let i = 0; i < awsTraceIdSegments.length; i++) {
+            if (awsTraceIdSegments[i] === '1') {
+              continue
+            } else if (i === 1 && awsTraceIdSegments.length === 3) {
+              continue
+            } else if (i === 2 && awsTraceIdSegments.length === 3) {
+              // trace id padded by: "00000000"
+              traceId = awsTraceIdSegments[i].substring(8)
+            }
+          }
+        } else if (key === xrayParentKey) {
+          spanId = value
+        } else if (key === xraySampledKey) {
+          samplingPriority = parseInt(value)
+        } else if (key === xrayOriginKey) {
+          ddOrigin = String(value)
+        } else if (key === xraySelfKey) {
+          // self is added by load balancers and should be ignored
+          continue
+        } else if (key === xrayE2EStartTimeKey) {
+          baggage.startMs = parseInt(value)
+        } else {
+          baggage[key] = value
+        }
+      }
+
+      if (traceId && spanId) {
+        const spanContext = new DatadogSpanContext({
+          traceId: id(traceId, 16),
+          spanId: id(spanId, 16),
+          sampling: { samplingPriority },
+          baggageItems: baggage
+        })
+        if (ddOrigin) {
+          spanContext._trace.origin = ddOrigin
+        }
+        return spanContext
+      }
+      return null
+    } else {
+      return null
+    }
+  }
+
+  _parseAWSTraceHeader (header) {
+    // parses AWSTraceHeader string to object
+    // ex: 'Root=1-00000000-00000000fffffffffffffffe;Parent=ffffffffffffffff;Sampled=1;_dd.origin=fakeOrigin;
+    const obj = {}
+    const keyValuePairs = header.split(';')
+    keyValuePairs.forEach(pair => {
+      const [key, value] = pair.split('=')
+      obj[key.toLowerCase()] = value.toLowerCase()
+    })
+    return obj
+  }
+
   static _convertOtelContextToDatadog (traceId, spanId, traceFlag, ts, meta = {}) {
     const origin = null
     let samplingPriority = traceFlag

packages/dd-trace/test/opentracing/propagation/text_map.spec.js

@@ -391,6 +391,45 @@ describe('TextMapPropagator', () => {
         injectCh.unsubscribe(onSpanInject)
       }
     })
+
+    it('should skip injection of aws xray header without the feature flag', () => {
+      const carrier = {}
+      const spanContext = createContext({
+        traceId: id('0000000000000123'),
+        spanId: id('0000000000000456')
+      })
+
+      config.tracePropagationStyle.inject = []
+
+      propagator.inject(spanContext, carrier)
+
+      expect(carrier).to.not.have.property('x-amzn-trace-id')
+    })
+
+    it('should inject AWSTraceHeader X-Amzn-Trace-Id when configured', () => {
+      const baggageItems = {
+        bool: true,
+        a: 'b'
+      }
+      const spanContext = createContext({
+        baggageItems,
+        sampling: {
+          priority: 1
+        }
+      })
+
+      config.tracePropagationStyle.inject = ['xray']
+
+      const traceId = spanContext._traceId.toString().padStart(24, '0')
+      const spanId = spanContext._spanId.toString().padStart(16, '0')
+      const expectedHeader = `root=1-00000000-${traceId};parent=${spanId};sampled=1`
+      const additionalParts = ';bool=true;a=b'
+      const carrier = {}
+
+      propagator.inject(spanContext, carrier)
+
+      expect(carrier['x-amzn-trace-id']).to.equal(expectedHeader + additionalParts)
+    })
   })
 
   describe('extract', () => {
@@ -662,6 +701,100 @@ describe('TextMapPropagator', () => {
       expect(spanContext._tracestate).to.be.undefined
     })
 
+    it('should not extract AWSTraceHeader X-Amzn-Trace-Id when not configured', () => {
+      const traceId = '4ef684dbd03d632e'
+      const spanId = '7e8d56262375628a'
+      const sampled = 1
+      const unparsedHeader = `Root=1-6583199d-00000000${traceId};Parent=${spanId};Sampled=${sampled}`
+      const carrier = {
+        'x-amzn-trace-id': unparsedHeader
+      }
+
+      config.tracePropagationStyle.extract = []
+
+      const spanContext = propagator.extract(carrier)
+
+      expect(spanContext).to.be.null
+    })
+
+    it('should extract AWSTraceHeader X-Amzn-Trace-Id when configured', () => {
+      const traceId = '4ef684dbd03d632e'
+      const spanId = '7e8d56262375628a'
+      const sampled = 1
+      const unparsedHeader = `Root=1-6583199d-00000000${traceId};Parent=${spanId};Sampled=${sampled}`
+      const carrier = {
+        'x-amzn-trace-id': unparsedHeader
+      }
+
+      config.tracePropagationStyle.extract = ['xray']
+
+      const spanContext = propagator.extract(carrier)
+
+      expect(spanContext.toTraceId()).to.equal(id(traceId, 16).toString(10))
+      expect(spanContext.toSpanId()).to.equal(id(spanId, 16).toString(10))
+      expect(spanContext._sampling.samplingPriority).to.equal(sampled)
+    })
+
+    it('should not extract AWSTraceHeader X-Amzn-Trace-Id if in an unexpected format or missing data', () => {
+      const traceId = '4ef684dbd03d632e'
+      const spanId = '7e8d56262375628a'
+      const sampled = 1
+
+      // All necessary fields included but trace id root segment is not formatted properly
+      let unparsedHeader = `Root=${traceId};Parent=${spanId};Sampled=${sampled}`
+      const carrier = {
+        'x-amzn-trace-id': unparsedHeader
+      }
+
+      config.tracePropagationStyle.extract = ['xray']
+
+      let spanContext = propagator.extract(carrier)
+
+      expect(spanContext).to.be.null
+
+      // Missing necessary parent field
+      unparsedHeader = `Root=1-6583199d-00000000${traceId};Sampled=${sampled}`
+
+      config.tracePropagationStyle.extract = ['xray']
+
+      spanContext = propagator.extract(carrier)
+
+      expect(spanContext).to.be.null
+
+      // Missing necessary root field
+      unparsedHeader = `Parent=${spanId};Sampled=${sampled}`
+
+      config.tracePropagationStyle.extract = ['xray']
+
+      spanContext = propagator.extract(carrier)
+
+      expect(spanContext).to.be.null
+    })
+
+    it('should extract AWSTraceHeader X-Amzn-Trace-Id with origin and baggage', () => {
+      const traceId = '4ef684dbd03d632e'
+      const spanId = '7e8d56262375628a'
+      const sampled = 1
+      const unparsedHeader = `Root=1-6583199d-00000000${traceId};Parent=${spanId};Sampled=${sampled}`
+      const additionalParts = ';_dd.origin=localhost;baggage_key=baggage_value;foo=bar'
+      const carrier = {
+        'x-amzn-trace-id': unparsedHeader + additionalParts
+      }
+
+      config.tracePropagationStyle.extract = ['xray']
+
+      const spanContext = propagator.extract(carrier)
+
+      expect(spanContext.toTraceId()).to.equal(id(traceId, 16).toString(10))
+      expect(spanContext.toSpanId()).to.equal(id(spanId, 16).toString(10))
+      expect(spanContext._sampling.samplingPriority).to.equal(sampled)
+
+      expect(spanContext._baggageItems).to.deep.equal({
+        baggage_key: 'baggage_value', foo: 'bar'
+      })
+      expect(spanContext._trace.origin).to.equal('localhost')
+    })
+
     it('extracts span_id from tracecontext headers and stores datadog parent-id in trace_distributed_tags', () => {
       textMap['x-datadog-trace-id'] = '61185'
       textMap['x-datadog-parent-id'] = '15'