aj/authorizer trace context

event_samples/api-gateway-traced-authorizer.json

@@ -0,0 +1,89 @@
+{
+  "resource": "/hello",
+  "path": "/hello",
+  "httpMethod": "POST",
+  "headers": {
+    "Accept": "*/*",
+    "Accept-Encoding": "gzip, deflate, br",
+    "Authorization": "password",
+    "CloudFront-Forwarded-Proto": "https",
+    "CloudFront-Is-Desktop-Viewer": "true",
+    "CloudFront-Is-Mobile-Viewer": "false",
+    "CloudFront-Is-SmartTV-Viewer": "false",
+    "CloudFront-Is-Tablet-Viewer": "false",
+    "CloudFront-Viewer-ASN": "174",
+    "CloudFront-Viewer-Country": "US",
+    "Host": "3gsxz7lha4.execute-api.sa-east-1.amazonaws.com",
+    "Postman-Token": "62ccb3d9-a44f-427c-9952-418c0a2eb1c3",
+    "User-Agent": "PostmanRuntime/7.29.0",
+    "Via": "1.1 685fd458a839bc725280d7e9380648c6.cloudfront.net (CloudFront)",
+    "X-Amz-Cf-Id": "90JXZEr6stVabQV78Zwn5EADW0evkpWINdmt3jzkuHQh9KtqowKejw==",
+    "X-Amzn-Trace-Id": "Root=1-62ffee4f-373bdfda15f09a065a39ac73",
+    "X-Forwarded-For": "38.142.177.195, 64.252.135.71",
+    "X-Forwarded-Port": "443",
+    "X-Forwarded-Proto": "https"
+  },
+  "multiValueHeaders": {
+    "Accept": ["*/*"],
+    "Accept-Encoding": ["gzip, deflate, br"],
+    "Authorization": ["password"],
+    "CloudFront-Forwarded-Proto": ["https"],
+    "CloudFront-Is-Desktop-Viewer": ["true"],
+    "CloudFront-Is-Mobile-Viewer": ["false"],
+    "CloudFront-Is-SmartTV-Viewer": ["false"],
+    "CloudFront-Is-Tablet-Viewer": ["false"],
+    "CloudFront-Viewer-ASN": ["174"],
+    "CloudFront-Viewer-Country": ["US"],
+    "Host": ["3gsxz7lha4.execute-api.sa-east-1.amazonaws.com"],
+    "Postman-Token": ["62ccb3d9-a44f-427c-9952-418c0a2eb1c3"],
+    "User-Agent": ["PostmanRuntime/7.29.0"],
+    "Via": ["1.1 685fd458a839bc725280d7e9380648c6.cloudfront.net (CloudFront)"],
+    "X-Amz-Cf-Id": ["90JXZEr6stVabQV78Zwn5EADW0evkpWINdmt3jzkuHQh9KtqowKejw=="],
+    "X-Amzn-Trace-Id": ["Root=1-62ffee4f-373bdfda15f09a065a39ac73"],
+    "X-Forwarded-For": ["38.142.177.195, 64.252.135.71"],
+    "X-Forwarded-Port": ["443"],
+    "X-Forwarded-Proto": ["https"]
+  },
+  "queryStringParameters": null,
+  "multiValueQueryStringParameters": null,
+  "pathParameters": null,
+  "stageVariables": null,
+  "requestContext": {
+    "resourceId": "oozq9u",
+    "authorizer": {
+      "_datadog": "{\"x-datadog-trace-id\":\"2095319761084710747\",\"x-datadog-parent-id\":\"2095319761084710747\",\"x-datadog-sampling-priority\":\"1\",\"x-datadog-parent-span-finish-time\":1660939857052}",
+      "principalId": "foo",
+      "integrationLatency": 1419,
+      "preserve": "this key set by a customer"
+    },
+    "resourcePath": "/hello",
+    "httpMethod": "POST",
+    "extendedRequestId": "XIIseElXGjQFvXg=",
+    "requestTime": "19/Aug/2022:20:10:55 +0000",
+    "path": "/dev/hello",
+    "accountId": "601427279990",
+    "protocol": "HTTP/1.1",
+    "stage": "dev",
+    "domainPrefix": "3gsxz7lha4",
+    "requestTimeEpoch": 1660939855656,
+    "requestId": "f1f8d46d-ef6f-46af-a3ed-7a20a26e2651",
+    "identity": {
+      "cognitoIdentityPoolId": null,
+      "accountId": null,
+      "cognitoIdentityId": null,
+      "caller": null,
+      "sourceIp": "38.142.177.195",
+      "principalOrgId": null,
+      "accessKey": null,
+      "cognitoAuthenticationType": null,
+      "cognitoAuthenticationProvider": null,
+      "userArn": null,
+      "userAgent": "PostmanRuntime/7.29.0",
+      "user": null
+    },
+    "domainName": "3gsxz7lha4.execute-api.sa-east-1.amazonaws.com",
+    "apiId": "3gsxz7lha4"
+  },
+  "body": null,
+  "isBase64Encoded": false
+}

src/index.ts

@@ -64,6 +64,7 @@ export const defaultConfig: Config = {
   captureLambdaPayload: false,
   createInferredSpan: true,
   debugLogging: false,
+  encodeAuthorizerContext: true,
   enhancedMetrics: true,
   forceWrap: false,
   injectLogContext: true,

src/trace/constants.ts

@@ -22,3 +22,4 @@ export const xrayTraceEnvVar = "_X_AMZN_TRACE_ID";
 export const awsXrayDaemonAddressEnvVar = "AWS_XRAY_DAEMON_ADDRESS";
 export const ddtraceVersion = "X.X.X";
 export const apiGatewayEventV2 = "2.0";
+export const parentSpanFinishTimeHeader = "x-datadog-parent-span-finish-time";

src/trace/context.spec.ts

@@ -22,6 +22,7 @@ import {
   readTraceFromSQSEvent,
   readTraceFromHTTPEvent,
   readTraceFromLambdaContext,
+  readTraceFromAuthorizerEvent,
 } from "./context";
 
 let sentSegment: any;
@@ -221,6 +222,27 @@ describe("readTraceFromEvent", () => {
     });
   });
 
+  it("can parse a traced authorizer source", () => {
+    const result = readTraceFromEvent({
+      requestContext: {
+        resourceId: "oozq9u",
+        authorizer: {
+          _datadog:
+            '{"x-datadog-trace-id":"2389589954026090296","x-datadog-parent-id":"2389589954026090296","x-datadog-sampling-priority":"1","x-datadog-parent-span-finish-time":1660939899233}',
+          principalId: "foo",
+          integrationLatency: 71,
+          preserve: "this key set by a customer",
+        },
+      },
+    });
+    expect(result).toEqual({
+      parentID: "2389589954026090296",
+      sampleMode: 1,
+      source: "event",
+      traceID: "2389589954026090296",
+    });
+  });
+
   it("can parse an SNS message source", () => {
     const result = readTraceFromEvent({
       Records: [

src/trace/context.ts

@@ -207,23 +207,8 @@ export function readTraceFromSQSEvent(event: SQSEvent): TraceContext | undefined
     const traceHeaders = event.Records[0].messageAttributes._datadog.stringValue;
 
     try {
-      const traceData = JSON.parse(traceHeaders);
-      const traceID = traceData[traceIDHeader];
-      const parentID = traceData[parentIDHeader];
-      const sampledHeader = traceData[samplingPriorityHeader];
+      const trace = exportTraceData(JSON.parse(traceHeaders));
 
-      if (typeof traceID !== "string" || typeof parentID !== "string" || typeof sampledHeader !== "string") {
-        return;
-      }
-
-      const sampleMode = parseInt(sampledHeader, 10);
-
-      const trace = {
-        parentID,
-        sampleMode,
-        source: Source.Event,
-        traceID,
-      };
       logDebug(`extracted trace context from sqs event`, { trace, event });
       return trace;
     } catch (err) {
@@ -253,21 +238,8 @@ export function readTraceFromSNSSQSEvent(event: SQSEvent): TraceContext | undefi
           const b64Decoded = Buffer.from(parsedBody.MessageAttributes._datadog.Value, "base64").toString("ascii");
           traceData = JSON.parse(b64Decoded);
         }
-        const traceID = traceData[traceIDHeader];
-        const parentID = traceData[parentIDHeader];
-        const sampledHeader = traceData[samplingPriorityHeader];
+        const trace = exportTraceData(traceData);
 
-        if (typeof traceID !== "string" || typeof parentID !== "string" || typeof sampledHeader !== "string") {
-          return;
-        }
-        const sampleMode = parseInt(sampledHeader, 10);
-
-        const trace = {
-          parentID,
-          sampleMode,
-          source: Source.Event,
-          traceID,
-        };
         logDebug(`extracted trace context from SNS SQS event`, { trace, event });
         return trace;
       }
@@ -285,22 +257,7 @@ export function readTraceFromKinesisEvent(event: KinesisStreamEvent): TraceConte
     try {
       const parsedBody = JSON.parse(Buffer.from(event.Records[0].kinesis.data, "base64").toString("ascii")) as any;
       if (parsedBody && parsedBody._datadog) {
-        const traceData = parsedBody._datadog;
-        const traceID = traceData[traceIDHeader];
-        const parentID = traceData[parentIDHeader];
-        const sampledHeader = traceData[samplingPriorityHeader];
-
-        if (typeof traceID !== "string" || typeof parentID !== "string" || typeof sampledHeader !== "string") {
-          return;
-        }
-        const sampleMode = parseInt(sampledHeader, 10);
-
-        const trace = {
-          parentID,
-          sampleMode,
-          source: Source.Event,
-          traceID,
-        };
+        const trace = exportTraceData(parsedBody._datadog);
         logDebug(`extracted trace context from Kinesis event`, { trace });
         return trace;
       }
@@ -316,22 +273,7 @@ export function readTraceFromKinesisEvent(event: KinesisStreamEvent): TraceConte
 export function readTraceFromEventbridgeEvent(event: EventBridgeEvent<any, any>): TraceContext | undefined {
   if (event?.detail?._datadog) {
     try {
-      const traceData = event.detail._datadog;
-      const traceID = traceData[traceIDHeader];
-      const parentID = traceData[parentIDHeader];
-      const sampledHeader = traceData[samplingPriorityHeader];
-
-      if (typeof traceID !== "string" || typeof parentID !== "string" || typeof sampledHeader !== "string") {
-        return;
-      }
-      const sampleMode = parseInt(sampledHeader, 10);
-
-      const trace = {
-        parentID,
-        sampleMode,
-        source: Source.Event,
-        traceID,
-      };
+      const trace = exportTraceData(event.detail._datadog);
       logDebug(`extracted trace context from Eventbridge event`, { trace, event });
       return trace;
     } catch (err) {
@@ -355,21 +297,7 @@ export function readTraceFromSNSEvent(event: SNSEvent): TraceContext | undefined
         );
         traceData = JSON.parse(b64Decoded);
       }
-      const traceID = traceData[traceIDHeader];
-      const parentID = traceData[parentIDHeader];
-      const sampledHeader = traceData[samplingPriorityHeader];
-
-      if (typeof traceID !== "string" || typeof parentID !== "string" || typeof sampledHeader !== "string") {
-        return;
-      }
-      const sampleMode = parseInt(sampledHeader, 10);
-
-      const trace = {
-        parentID,
-        sampleMode,
-        source: Source.Event,
-        traceID,
-      };
+      const trace = exportTraceData(traceData);
       logDebug(`extracted trace context from SNS event`, { trace, event });
       return trace;
     } catch (err) {
@@ -411,26 +339,7 @@ export function readTraceFromLambdaContext(context: any): TraceContext | undefin
     return;
   }
 
-  const traceID = traceData[traceIDHeader];
-  if (typeof traceID !== "string") {
-    return;
-  }
-  const parentID = traceData[parentIDHeader];
-  if (typeof parentID !== "string") {
-    return;
-  }
-  const sampledHeader = traceData[samplingPriorityHeader];
-  if (typeof sampledHeader !== "string") {
-    return;
-  }
-  const sampleMode = parseInt(sampledHeader, 10);
-
-  const trace = {
-    parentID,
-    sampleMode,
-    source: Source.Event,
-    traceID,
-  };
+  const trace = exportTraceData(traceData);
   logDebug(`extracted trace context from lambda context`, { trace, context });
   return trace;
 }
@@ -443,43 +352,40 @@ export function readTraceFromHTTPEvent(event: any): TraceContext | undefined {
     lowerCaseHeaders[key.toLowerCase()] = headers[key];
   }
 
-  const traceID = lowerCaseHeaders[traceIDHeader];
-  if (typeof traceID !== "string") {
-    return;
-  }
-  const parentID = lowerCaseHeaders[parentIDHeader];
-  if (typeof parentID !== "string") {
-    return;
-  }
-  const sampledHeader = lowerCaseHeaders[samplingPriorityHeader];
-  if (typeof sampledHeader !== "string") {
-    return;
-  }
-  const sampleMode = parseInt(sampledHeader, 10);
-
-  const trace = {
-    parentID,
-    sampleMode,
-    source: Source.Event,
-    traceID,
-  };
+  const trace = exportTraceData(lowerCaseHeaders);
 
   logDebug(`extracted trace context from http event`, { trace, event });
   return trace;
 }
 
+export function readTraceFromAuthorizerEvent(event: any): TraceContext | undefined {
+  let traceData;
+  try {
+    traceData = JSON.parse(event.requestContext.authorizer._datadog);
+  } catch (error) {
+    logDebug(`unable to extract trace context from authorizer event`, { error });
+    return;
+  }
+  return exportTraceData(traceData);
+}
+
 export function readTraceFromEvent(event: any): TraceContext | undefined {
   if (!event || typeof event !== "object") {
     return;
   }
 
+  if (event?.requestContext?.authorizer?._datadog && event.requestContext.authorizer.integrationLatency > 0) {
+    return readTraceFromAuthorizerEvent(event);
+  }
+
   if (event.headers !== null && typeof event.headers === "object") {
     return readTraceFromHTTPEvent(event);
   }
 
   if (isSNSEvent(event)) {
     return readTraceFromSNSEvent(event);
   }
+
   if (isSNSSQSEvent(event)) {
     return readTraceFromSNSSQSEvent(event);
   }
@@ -640,3 +546,22 @@ export function convertToAPMParentID(xrayParentID: string): string | undefined {
   }
   return hex.toString(10);
 }
+
+function exportTraceData(traceData: any): TraceContext | undefined {
+  const traceID = traceData[traceIDHeader];
+  const parentID = traceData[parentIDHeader];
+  const sampledHeader = traceData[samplingPriorityHeader];
+
+  if (typeof traceID !== "string" || typeof parentID !== "string" || typeof sampledHeader !== "string") {
+    return;
+  }
+
+  const sampleMode = parseInt(sampledHeader, 10);
+
+  return {
+    parentID,
+    sampleMode,
+    source: Source.Event,
+    traceID,
+  };
+}