fixes

DataDog · Mar 14, 2024 · 00d7309 · 00d7309
1 parent 5184c05
commit 00d7309
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 14 deletions.
diff --git a/src/trace/context/extractors/sqs.ts b/src/trace/context/extractors/sqs.ts
@@ -10,20 +10,20 @@ export class SQSEventTraceExtractor implements EventTraceExtractor {
 
   extract(event: SQSEvent): SpanContextWrapper | null {
     try {
-      var prepared_headers;
+      let preparedHeaders;
       const headers = event?.Records?.[0]?.messageAttributes?._datadog?.stringValue;
       if (headers !== undefined) {
-        prepared_headers = JSON.parse(headers);
+        preparedHeaders = JSON.parse(headers);
       } else {
         if (event?.Records?.[0]?.attributes?.AWSTraceHeader) {
-          prepared_headers = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader);
+          preparedHeaders = XrayService.extraceDDContextFromAWSTraceHeader(event.Records[0].attributes.AWSTraceHeader);
         }
       }
 
-      if (!prepared_headers) return null;
-      const traceContext = this.tracerWrapper.extract(prepared_headers);
+      if (!preparedHeaders) return null;
+      const traceContext = this.tracerWrapper.extract(preparedHeaders);
       if (traceContext === null) {
-        logDebug("Unable to extract the injected trace context from event");
+        logDebug(`Failed to extract trace context from prepared headers: ${preparedHeaders}`);
         return null;
       }
       logDebug(`Extracted trace context from SQS event`, { traceContext, event });

diff --git a/src/trace/xray-service.spec.ts b/src/trace/xray-service.spec.ts
@@ -1,3 +1,8 @@
+import {
+  DATADOG_SAMPLING_PRIORITY_HEADER,
+  DATADOG_TRACE_ID_HEADER,
+  DATADOG_PARENT_ID_HEADER,
+} from "./context/extractor";
 import { SampleMode } from "./trace-context-service";
 import { XrayService } from "./xray-service";
 
@@ -340,4 +345,46 @@ describe("XrayService", () => {
       expect(traceId).toBeUndefined();
     });
   });
+
+  describe("parseAWSTraceHeader", () => {
+    it("parses AWS trace header correctly", () => {
+      const awsTraceHeader = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
+      const xrayHeaders = XrayService.parseAWSTraceHeader(awsTraceHeader);
+      expect(xrayHeaders).toEqual({
+        parentId: "94ae789b969f1cc5",
+        sampled: "1",
+        traceId: "1-5e272390-8c398be037738dc042009320",
+      });
+    });
+    it.each(["Root=1-5e272390-8c398be037738dc042009320", "Root=1-65f2f78c-0000000008addb5405b376c0;Parent;Sampled"])(
+      "returns undefined when AWS trace header is malformatted",
+      (awsTraceHeader) => {
+        const xrayHeaders = XrayService.parseAWSTraceHeader(awsTraceHeader);
+        expect(xrayHeaders).toBeUndefined();
+      },
+    );
+  });
+  describe("extraceDDContextFromAWSTraceHeader", () => {
+    it("extracts Datadog trace context from AWS trace header", () => {
+      const awsTraceId = "Root=1-65f2f78c-0000000008addb5405b376c0;Parent=5abcb7ed643995c7;Sampled=1";
+      const ddTraceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceId);
+
+      expect(ddTraceContext).toEqual({
+        [DATADOG_TRACE_ID_HEADER]: "625397077193750208",
+        [DATADOG_PARENT_ID_HEADER]: "6538302989251745223",
+        [DATADOG_SAMPLING_PRIORITY_HEADER]: "1",
+      });
+    });
+
+    it("returns null when AWS trace header is NOT injected by dd-trace", () => {
+      const awsTraceId = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
+      const ddTraceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceId);
+      expect(ddTraceContext).toBeNull();
+    });
+    it("returns null when AWS trace header cannot be parsed", () => {
+      const awsTraceId = "Root=1-5e272390-8c398be037738dc042009320;;";
+      const ddTraceContext = XrayService.extraceDDContextFromAWSTraceHeader(awsTraceId);
+      expect(ddTraceContext).toBeNull();
+    });
+  });
 });
diff --git a/src/trace/xray-service.ts b/src/trace/xray-service.ts
@@ -197,24 +197,25 @@ export class XrayService {
 
     // We want to turn the last 63 bits into a decimal number in a string representation
     try {
-      return (BigInt("0x" + lastPart) % BigInt("0x8000000000000000")).toString(10);
+      return (BigInt("0x" + lastPart) % BigInt("0x8000000000000000")).toString(10); // mod by 2^63 will leave us with the last 63 bits
     } catch (_) {
+      logDebug(`Faied to convert trace id ${lastPart}`);
       return undefined;
     }
   }
 
   public static extraceDDContextFromAWSTraceHeader(amznTraceId: string): DatadogTraceHeaders | null {
-    var aws_context = XrayService.parseAWSTraceHeader(amznTraceId);
-    if (!aws_context) {
+    const awsContext = XrayService.parseAWSTraceHeader(amznTraceId);
+    if (!awsContext) {
       return null;
     }
-    const trace_id_parts = aws_context.traceId.split("-");
-    if (trace_id_parts && trace_id_parts.length > 2 && trace_id_parts[2].startsWith("00000000")) {
+    const traceIdParts = awsContext.traceId.split("-");
+    if (traceIdParts && traceIdParts.length > 2 && traceIdParts[2].startsWith("00000000")) {
       // This AWSTraceHeader contains Datadog injected trace context
       return {
-        [DATADOG_TRACE_ID_HEADER]: hexStrToDecimalStr(trace_id_parts[2].substring(8)),
-        [DATADOG_PARENT_ID_HEADER]: hexStrToDecimalStr(aws_context.parentId),
-        [DATADOG_SAMPLING_PRIORITY_HEADER]: aws_context.sampled,
+        [DATADOG_TRACE_ID_HEADER]: hexStrToDecimalStr(traceIdParts[2].substring(8)),
+        [DATADOG_PARENT_ID_HEADER]: hexStrToDecimalStr(awsContext.parentId),
+        [DATADOG_SAMPLING_PRIORITY_HEADER]: awsContext.sampled,
       };
     }
     return null;