From 390245e973d2e220b4dc8eef5046a9a526051b68 Mon Sep 17 00:00:00 2001
From: "chris.agocs" <chris.agocs@datadoghq.com>
Date: Mon, 21 Jun 2021 11:47:37 -0400
Subject: [PATCH 1/4] Check for the new trace propagation format

---
 src/trace/context.ts | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/trace/context.ts b/src/trace/context.ts
index 918fa1810..0fee0f2f5 100644
--- a/src/trace/context.ts
+++ b/src/trace/context.ts
@@ -223,9 +223,28 @@ export function readTraceFromLambdaContext(context: any): TraceContext | undefin
     return;
   }
 
-  const traceData = context.clientContext?.custom?._datadog;
+  const custom = context.clientContext?.custom;
 
-  if (!traceData || typeof traceData !== "object") {
+  if (!custom || typeof custom !== "object") {
+    return;
+  }
+  var traceData = null;
+
+  if (
+    custom._datadog &&
+    typeof custom._datadog === "object" &&
+    custom._datadog.hasOwnProperty(traceIDHeader) &&
+    custom._datadog.hasOwnProperty(parentIDHeader) &&
+    custom._datadog.hasOwnProperty(samplingPriorityHeader)
+  ) {
+    traceData = custom._datadog;
+  } else if (
+    custom.hasOwnProperty(traceIDHeader) &&
+    custom.hasOwnProperty(parentIDHeader) &&
+    custom.hasOwnProperty(samplingPriorityHeader)
+  ) {
+    traceData = custom;
+  } else {
     return;
   }
 

From a6158655c4fa83a563fb145b746c7785f2e6ad2c Mon Sep 17 00:00:00 2001
From: "chris.agocs" <chris.agocs@datadoghq.com>
Date: Mon, 21 Jun 2021 11:55:04 -0400
Subject: [PATCH 2/4] var -> let

---
 src/trace/context.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/trace/context.ts b/src/trace/context.ts
index 0fee0f2f5..df605a6f9 100644
--- a/src/trace/context.ts
+++ b/src/trace/context.ts
@@ -228,7 +228,7 @@ export function readTraceFromLambdaContext(context: any): TraceContext | undefin
   if (!custom || typeof custom !== "object") {
     return;
   }
-  var traceData = null;
+  let traceData = null;
 
   if (
     custom._datadog &&

From 1d4fea00581f0063acbcf1b8292295ae0ebdb7f3 Mon Sep 17 00:00:00 2001
From: "chris.agocs" <chris.agocs@datadoghq.com>
Date: Mon, 21 Jun 2021 15:20:29 -0400
Subject: [PATCH 3/4] Add a unit test

---
 src/trace/context.spec.ts | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/trace/context.spec.ts b/src/trace/context.spec.ts
index e415a9331..5ddbba9f5 100644
--- a/src/trace/context.spec.ts
+++ b/src/trace/context.spec.ts
@@ -393,7 +393,7 @@ describe("readTraceFromSQSEvent", () => {
 });
 
 describe("readTraceFromLambdaContext", () => {
-  it("can read from lambda context source", () => {
+  it("can read from lambda context source, legacy style", () => {
     const result = readTraceFromLambdaContext({
       clientContext: {
         custom: {
@@ -413,6 +413,24 @@ describe("readTraceFromLambdaContext", () => {
       source: Source.Event,
     });
   });
+  it("can read from lambda context source, new style", () => {
+    const result = readTraceFromLambdaContext({
+      clientContext: {
+        custom: {
+          "x-datadog-trace-id": "666",
+          "x-datadog-parent-id": "777",
+          "x-datadog-sampled": "1",
+          "x-datadog-sampling-priority": "1",
+        },
+      },
+    });
+    expect(result).toEqual({
+      parentID: "777",
+      sampleMode: SampleMode.AUTO_KEEP,
+      traceID: "666",
+      source: Source.Event,
+    });
+  });
   it("can handle no `custom` key", () => {
     const result = readTraceFromLambdaContext({
       clientContext: {

From c8f9f90f44d5c049aea15ed5c48bad4b4e242369 Mon Sep 17 00:00:00 2001
From: "chris.agocs" <chris.agocs@datadoghq.com>
Date: Mon, 21 Jun 2021 16:22:44 -0400
Subject: [PATCH 4/4] Incorporate PR feedback. an extra hasOwnProperty and an
 extra test

---
 src/trace/context.spec.ts | 8 ++++++++
 src/trace/context.ts      | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/trace/context.spec.ts b/src/trace/context.spec.ts
index 5ddbba9f5..fece383b5 100644
--- a/src/trace/context.spec.ts
+++ b/src/trace/context.spec.ts
@@ -439,6 +439,14 @@ describe("readTraceFromLambdaContext", () => {
     });
     expect(result).toBeUndefined();
   });
+  it("can handle a string `custom` key", () => {
+    const result = readTraceFromLambdaContext({
+      clientContext: {
+        custom: "bar",
+      },
+    });
+    expect(result).toBeUndefined();
+  });
   it("can handle no context", () => {
     const result = readTraceFromLambdaContext(undefined);
     expect(result).toBeUndefined();
diff --git a/src/trace/context.ts b/src/trace/context.ts
index df605a6f9..c38df2138 100644
--- a/src/trace/context.ts
+++ b/src/trace/context.ts
@@ -231,7 +231,7 @@ export function readTraceFromLambdaContext(context: any): TraceContext | undefin
   let traceData = null;
 
   if (
-    custom._datadog &&
+    custom.hasOwnProperty("_datadog") &&
     typeof custom._datadog === "object" &&
     custom._datadog.hasOwnProperty(traceIDHeader) &&
     custom._datadog.hasOwnProperty(parentIDHeader) &&