diff --git a/src/trace/constants.ts b/src/trace/constants.ts index 4caaafec..46758933 100644 --- a/src/trace/constants.ts +++ b/src/trace/constants.ts @@ -1,28 +1,4 @@ -export enum SampleMode { - USER_REJECT = -1, - AUTO_REJECT = 0, - AUTO_KEEP = 1, - USER_KEEP = 2, -} -export enum Source { - Xray = "xray", - Event = "event", - DDTrace = "ddtrace", -} - -export const traceIDHeader = "x-datadog-trace-id"; -export const parentIDHeader = "x-datadog-parent-id"; -export const samplingPriorityHeader = "x-datadog-sampling-priority"; -export const xraySubsegmentName = "datadog-metadata"; -export const xraySubsegmentKey = "trace"; -export const xrayBaggageSubsegmentKey = "root_span_metadata"; -export const xrayLambdaFunctionTagsKey = "lambda_function_tags"; -export const xraySubsegmentNamespace = "datadog"; -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"; -export const authorizingRequestIdHeader = "x-datadog-authorizing-requestid"; - export const DD_SERVICE_ENV_VAR = "DD_SERVICE"; diff --git a/src/trace/context.ts b/src/trace/context.ts deleted file mode 100644 index 0df9f53e..00000000 --- a/src/trace/context.ts +++ /dev/null @@ -1,699 +0,0 @@ -import { Context, EventBridgeEvent, KinesisStreamEvent, SNSEvent, SNSMessage, SQSEvent } from "aws-lambda"; -import { BigNumber } from "bignumber.js"; -import { randomBytes } from "crypto"; -import { createSocket, Socket } from "dgram"; -import { logDebug, logError } from "../utils"; -import { - isAppSyncResolverEvent, - isEventBridgeEvent, - isKinesisStreamEvent, - isSNSEvent, - isSNSSQSEvent, - isSQSEvent, - isEBSQSEvent, -} from "../utils/event-type-guards"; -import { - authorizingRequestIdHeader, - awsXrayDaemonAddressEnvVar, - parentIDHeader, - SampleMode, - samplingPriorityHeader, - Source, - traceIDHeader, - xrayBaggageSubsegmentKey, - xrayLambdaFunctionTagsKey, - xraySubsegmentKey, - xraySubsegmentName, - xraySubsegmentNamespace, - xrayTraceEnvVar, -} from "./constants"; -import { TraceExtractor } from "./listener"; -import { eventSubTypes, parseEventSourceSubType } from "./trigger"; -import { Md5 } from "ts-md5"; - -export interface XRayTraceHeader { - traceID: string; - parentID: string; - sampled: number; -} - -export interface TraceContext { - traceID: string; - parentID: string; - sampleMode: SampleMode; - source: Source; -} - -export interface StepFunctionContext { - "step_function.execution_name": string; - "step_function.execution_id": string; - "step_function.execution_input": object; - "step_function.execution_role_arn": string; - "step_function.execution_start_time": string; - "step_function.state_machine_name": string; - "step_function.state_machine_arn": string; - "step_function.state_entered_time": string; - "step_function.state_name": string; - "step_function.state_retry_count": number; -} - -export function readTraceFromStepFunctionsContext(stepFunctionContext: StepFunctionContext): TraceContext | undefined { - return { - traceID: deterministicMd5HashToBigIntString(stepFunctionContext["step_function.execution_id"]), - parentID: deterministicMd5HashToBigIntString( - stepFunctionContext["step_function.execution_id"] + - "#" + - stepFunctionContext["step_function.state_name"] + - "#" + - stepFunctionContext["step_function.state_entered_time"], - ), - sampleMode: SampleMode.AUTO_KEEP.valueOf(), - source: Source.Event, - }; -} - -export function hexToBinary(hex: string) { - // convert hex to binary and padding with 0 in the front to fill 128 bits - return parseInt(hex, 16).toString(2).padStart(4, "0"); -} - -export function deterministicMd5HashInBinary(s: string): string { - // Md5 here is used here because we don't need a cryptographically secure hashing method but to generate the same trace/span ids as the backend does - const hex = Md5.hashStr(s); - - let binary = ""; - for (let i = 0; i < hex.length; i++) { - const ch = hex.charAt(i); - binary = binary + hexToBinary(ch); - } - - const res = "0" + binary.substring(1, 64); - if (res === "0".repeat(64)) { - return "1"; - } - return res; -} - -export function deterministicMd5HashToBigIntString(s: string): string { - const binaryString = deterministicMd5HashInBinary(s); - return BigInt("0b" + binaryString).toString(); -} - -/** - * Reads the trace context from either an incoming lambda event, or the current xray segment. - * @param event An incoming lambda event. This must have incoming trace headers in order to be read. - */ -export async function extractTraceContext( - event: any, - context: Context, - extractor?: TraceExtractor, - decodeAuthorizerContext: boolean = true, -): Promise { - let trace; - - if (extractor) { - try { - trace = await extractor(event, context); - logDebug(`extracted trace context from the custom extractor`, { trace }); - } catch (error) { - if (error instanceof Error) { - logError("custom extractor function failed", error as Error); - } - } - } - - if (!trace) { - trace = readTraceFromEvent(event, decodeAuthorizerContext); - } - - if (!trace) { - trace = readTraceFromLambdaContext(context); - } - - const stepFuncContext = readStepFunctionContextFromEvent(event); - if (stepFuncContext) { - try { - addStepFunctionContextToXray(stepFuncContext); - } catch (error) { - if (error instanceof Error) { - logError("couldn't add step function metadata to xray", error as Error); - } - } - if (trace === undefined) { - trace = readTraceFromStepFunctionsContext(stepFuncContext); - if (trace !== undefined) { - return trace; - } - } - } - - if (trace !== undefined) { - try { - addTraceContextToXray(trace); - logDebug(`added trace context to xray metadata`, { trace }); - } catch (error) { - // This might fail if running in an environment where xray isn't set up, (like for local development). - if (error instanceof Error) { - logError("couldn't add trace context to xray metadata", error as Error); - } - } - return trace; - } - return readTraceContextFromXray(); -} - -export function addTraceContextToXray(traceContext: TraceContext) { - const val = { - "parent-id": traceContext.parentID, - "sampling-priority": traceContext.sampleMode.toString(10), - "trace-id": traceContext.traceID, - }; - - addXrayMetadata(xraySubsegmentKey, val); -} - -export function addStepFunctionContextToXray(context: StepFunctionContext) { - addXrayMetadata(xrayBaggageSubsegmentKey, context); -} - -export function addLambdaFunctionTagsToXray(triggerTags: { [key: string]: string }) { - addXrayMetadata(xrayLambdaFunctionTagsKey, triggerTags); -} - -export function addXrayMetadata(key: string, metadata: Record) { - const segment = generateXraySubsegment(key, metadata); - if (segment === undefined) { - return; - } - sendXraySubsegment(segment); -} - -export function generateXraySubsegment(key: string, metadata: Record) { - const header = process.env[xrayTraceEnvVar]; - if (header === undefined) { - logDebug("couldn't read xray trace header from env"); - return; - } - const context = parseXrayTraceContextHeader(header); - if (context === undefined) { - logDebug("couldn't parse xray trace header from env"); - return; - } - const sampled = convertToSampleMode(parseInt(context.xraySampled, 10)); - if (sampled === SampleMode.USER_REJECT || sampled === SampleMode.AUTO_REJECT) { - logDebug("discarding xray metadata subsegment due to sampling"); - return; - } - - // Convert from milliseconds to seconds - const time = Date.now() * 0.001; - - return JSON.stringify({ - id: randomBytes(8).toString("hex"), - trace_id: context.xrayTraceID, - parent_id: context.xrayParentID, - name: xraySubsegmentName, - start_time: time, - end_time: time, - type: "subsegment", - metadata: { - [xraySubsegmentNamespace]: { - [key]: metadata, - }, - }, - }); -} - -export function sendXraySubsegment(segment: string) { - const xrayDaemonEnv = process.env[awsXrayDaemonAddressEnvVar]; - if (xrayDaemonEnv === undefined) { - logDebug("X-Ray daemon env var not set, not sending sub-segment"); - return; - } - const parts = xrayDaemonEnv.split(":"); - if (parts.length <= 1) { - logDebug("X-Ray daemon env var has invalid format, not sending sub-segment"); - return; - } - const port = parseInt(parts[1], 10); - const address = parts[0]; - - const message = Buffer.from(`{\"format\": \"json\", \"version\": 1}\n${segment}`); - let client: Socket | undefined; - try { - client = createSocket("udp4"); - // Send segment asynchronously to xray daemon - client.send(message, 0, message.length, port, address, (error, bytes) => { - client?.close(); - logDebug(`Xray daemon received metadata payload`, { error, bytes }); - }); - } catch (error) { - if (error instanceof Error) { - client?.close(); - logDebug("Error occurred submitting to xray daemon", error); - } - } -} - -export function readTraceFromAppSyncEvent(event: any): TraceContext | undefined { - event.headers = event.request.headers; - return readTraceFromHTTPEvent(event, false); -} - -export function readTraceFromSQSEvent(event: SQSEvent): TraceContext | undefined { - if (event?.Records?.[0]?.messageAttributes?._datadog?.stringValue) { - const traceHeaders = event.Records[0].messageAttributes._datadog.stringValue; - - try { - const trace = exportTraceData(JSON.parse(traceHeaders)); - - logDebug(`extracted trace context from sqs event`, { trace, event }); - return trace; - } catch (err) { - if (err instanceof Error) { - logDebug("Error parsing SQS message trace data", err as Error); - } - return; - } - } - - return; -} - -export function readTraceFromSNSSQSEvent(event: SQSEvent): TraceContext | undefined { - if (event?.Records?.[0]?.body) { - try { - const parsedBody = JSON.parse(event.Records[0].body) as SNSMessage; - if ( - parsedBody.MessageAttributes && - parsedBody.MessageAttributes._datadog && - parsedBody.MessageAttributes._datadog.Value - ) { - let traceData; - if (parsedBody.MessageAttributes._datadog.Type === "String") { - traceData = JSON.parse(parsedBody.MessageAttributes._datadog.Value); - } else { - const b64Decoded = Buffer.from(parsedBody.MessageAttributes._datadog.Value, "base64").toString("ascii"); - traceData = JSON.parse(b64Decoded); - } - const trace = exportTraceData(traceData); - - logDebug(`extracted trace context from SNS SQS event`, { trace, event }); - return trace; - } - } catch (err) { - if (err instanceof Error) { - logDebug("Error parsing SNS SQS message trace data", err as Error); - } - return; - } - } -} - -export function readTraceFromEBSQSEvent(event: SQSEvent): TraceContext | undefined { - if (event?.Records?.[0]?.body) { - try { - const parsedBody = JSON.parse(event.Records[0].body) as EventBridgeEvent; - if (parsedBody?.detail?._datadog) { - const trace = exportTraceData(parsedBody.detail._datadog); - - logDebug(`extracted trace context from EventBridge SQS event`, { trace, event }); - return trace; - } - } catch (err) { - if (err instanceof Error) { - logDebug("Error parsing EventBridge SQS message trace data", err as Error); - } - return; - } - } -} - -export function readTraceFromKinesisEvent(event: KinesisStreamEvent): TraceContext | undefined { - if (event?.Records?.[0]?.kinesis?.data) { - try { - const parsedBody = JSON.parse(Buffer.from(event.Records[0].kinesis.data, "base64").toString("ascii")) as any; - if (parsedBody && parsedBody._datadog) { - const trace = exportTraceData(parsedBody._datadog); - logDebug(`extracted trace context from Kinesis event`, { trace }); - return trace; - } - } catch (err) { - if (err instanceof Error) { - logDebug("Error parsing Kinesis message trace data", err as Error); - } - return; - } - } -} - -export function readTraceFromEventbridgeEvent(event: EventBridgeEvent): TraceContext | undefined { - if (event?.detail?._datadog) { - try { - const trace = exportTraceData(event.detail._datadog); - logDebug(`extracted trace context from Eventbridge event`, { trace, event }); - return trace; - } catch (err) { - if (err instanceof Error) { - logDebug("Error parsing Eventbridge trace data", err as Error); - } - return; - } - } -} - -export function readTraceFromSNSEvent(event: SNSEvent): TraceContext | undefined { - if (event?.Records?.[0]?.Sns?.MessageAttributes?._datadog?.Value) { - try { - let traceData; - if (event.Records[0].Sns.MessageAttributes._datadog.Type === "String") { - traceData = JSON.parse(event.Records[0].Sns.MessageAttributes._datadog.Value); - } else { - const b64Decoded = Buffer.from(event.Records[0].Sns.MessageAttributes._datadog.Value, "base64").toString( - "ascii", - ); - traceData = JSON.parse(b64Decoded); - } - const trace = exportTraceData(traceData); - logDebug(`extracted trace context from SNS event`, { trace, event }); - return trace; - } catch (err) { - if (err instanceof Error) { - logDebug("Error parsing SNS SQS message trace data", err as Error); - } - return; - } - } -} - -export function readTraceFromLambdaContext(context: any): TraceContext | undefined { - if (!context || typeof context !== "object") { - return; - } - - const custom = context.clientContext?.custom; - - if (!custom || typeof custom !== "object") { - return; - } - let traceData = null; - - if ( - custom.hasOwnProperty("_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; - } - - const trace = exportTraceData(traceData); - logDebug(`extracted trace context from lambda context`, { trace, context }); - return trace; -} - -export function getInjectedAuthorizerData(event: any, eventSourceSubType: eventSubTypes) { - const authorizerHeaders = event?.requestContext?.authorizer; - if (!authorizerHeaders) return null; - const rawDatadogData = - eventSourceSubType === eventSubTypes.apiGatewayV2 ? authorizerHeaders.lambda._datadog : authorizerHeaders._datadog; - if (!rawDatadogData) return null; - const injectedData = JSON.parse(Buffer.from(rawDatadogData, "base64").toString()); - // use the injected requestId to tell if it's the authorizing invocation (not cached) - if ( - authorizerHeaders.integrationLatency > 0 || - event.requestContext.requestId === injectedData[authorizingRequestIdHeader] - ) { - return injectedData; - } else { - return null; - } -} - -export function readTraceFromHTTPEvent(event: any, decodeAuthorizerContext: boolean = true): TraceContext | undefined { - if (decodeAuthorizerContext) { - // need to set the trace context if using authorizer lambda in authorizing (non-cached) cases - try { - const eventSourceSubType: eventSubTypes = parseEventSourceSubType(event); - const injectedAuthorizerData = getInjectedAuthorizerData(event, eventSourceSubType); - if (injectedAuthorizerData !== null) { - return exportTraceData(injectedAuthorizerData); - } - } catch (error) { - logDebug(`unable to extract trace context from authorizer event.`, { error }); - } - } - - const headers = event.headers; - const lowerCaseHeaders: { [key: string]: string } = {}; - - for (const key of Object.keys(headers)) { - lowerCaseHeaders[key.toLowerCase()] = headers[key]; - } - - const trace = exportTraceData(lowerCaseHeaders); - - logDebug(`extracted trace context from http event`, { trace, event }); - return trace; -} - -export function readTraceFromEvent(event: any, decodeAuthorizerContext: boolean = true): TraceContext | undefined { - if (!event || typeof event !== "object") { - return; - } - - if (event.headers !== null && typeof event.headers === "object") { - return readTraceFromHTTPEvent(event, decodeAuthorizerContext); - } - - if (isSNSEvent(event)) { - return readTraceFromSNSEvent(event); - } - - if (isSNSSQSEvent(event)) { - return readTraceFromSNSSQSEvent(event); - } - - if (isEBSQSEvent(event)) { - return readTraceFromEBSQSEvent(event); - } - - if (isAppSyncResolverEvent(event)) { - return readTraceFromAppSyncEvent(event); - } - - if (isSQSEvent(event)) { - return readTraceFromSQSEvent(event); - } - if (isKinesisStreamEvent(event)) { - return readTraceFromKinesisEvent(event); - } - - if (isEventBridgeEvent(event)) { - return readTraceFromEventbridgeEvent(event); - } - - return; -} - -export function readTraceContextFromXray(): TraceContext | undefined { - const header = process.env[xrayTraceEnvVar]; - if (header === undefined) { - logDebug("couldn't read xray trace header from env"); - return; - } - const context = parseXrayTraceContextHeader(header); - - if (context === undefined) { - logError("couldn't read xray trace context from env, variable had invalid format"); - return undefined; - } - const parentID = convertToAPMParentID(context.xrayParentID); - if (parentID === undefined) { - logDebug("couldn't parse xray parent ID", context); - return; - } - const traceID = convertToAPMTraceID(context.xrayTraceID); - if (traceID === undefined) { - logDebug("couldn't parse xray trace ID", context); - return; - } - const sampleMode = convertToSampleMode(parseInt(context.xraySampled, 10)); - - const trace = { - parentID, - sampleMode, - source: Source.Xray, - traceID, - }; - logDebug(`extracted trace context from xray context`, { trace, header }); - return trace; -} - -function parseXrayTraceContextHeader(header: string) { - // Example: Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1 - logDebug(`Reading trace context from env var ${header}`); - const [root, parent, sampled] = header.split(";"); - if (parent === undefined || sampled === undefined) { - return; - } - const [, xrayTraceID] = root.split("="); - const [, xrayParentID] = parent.split("="); - const [, xraySampled] = sampled.split("="); - if (xraySampled === undefined || xrayParentID === undefined || xrayTraceID === undefined) { - return; - } - return { - xrayTraceID, - xraySampled, - xrayParentID, - }; -} - -export function readStepFunctionContextFromEvent(event: any): StepFunctionContext | undefined { - if (typeof event !== "object") { - return; - } - - const execution = event.Execution; - if (typeof execution !== "object") { - logDebug("event.Execution is not an object."); - return; - } - const executionID = execution.Id; - if (typeof executionID !== "string") { - logDebug("event.Execution.Id is not a string."); - return; - } - const executionInput = execution.Input; - const executionName = execution.Name; - if (typeof executionName !== "string") { - logDebug("event.Execution.Name is not a string."); - return; - } - const executionRoleArn = execution.RoleArn; - if (typeof executionRoleArn !== "string") { - logDebug("event.Execution.RoleArn is not a string."); - return; - } - const executionStartTime = execution.StartTime; - if (typeof executionStartTime !== "string") { - logDebug("event.Execution.StartTime is not a string."); - return; - } - - const state = event.State; - if (typeof state !== "object") { - logDebug("event.State is not an object."); - return; - } - const stateRetryCount = state.RetryCount; - if (typeof stateRetryCount !== "number") { - logDebug("event.State.RetryCount is not a string."); - return; - } - const stateEnteredTime = state.EnteredTime; - if (typeof stateEnteredTime !== "string") { - logDebug("event.State.EnteredTime is not a string."); - return; - } - const stateName = state.Name; - if (typeof stateName !== "string") { - logDebug("event.State.Name is not a string."); - return; - } - - const stateMachine = event.StateMachine; - if (typeof stateMachine !== "object") { - logDebug("event.StateMachine is not an object."); - return; - } - const stateMachineArn = stateMachine.Id; - if (typeof stateMachineArn !== "string") { - logDebug("event.StateMachine.Id is not a string."); - return; - } - const stateMachineName = stateMachine.Name; - if (typeof stateMachineName !== "string") { - logDebug("event.StateMachine.Name is not a string."); - return; - } - - return { - "step_function.execution_name": executionName, - "step_function.execution_id": executionID, - "step_function.execution_input": executionInput ?? {}, - "step_function.execution_role_arn": executionRoleArn, - "step_function.execution_start_time": executionStartTime, - "step_function.state_entered_time": stateEnteredTime, - "step_function.state_machine_arn": stateMachineArn, - "step_function.state_machine_name": stateMachineName, - "step_function.state_name": stateName, - "step_function.state_retry_count": stateRetryCount, - }; -} - -export function convertToSampleMode(xraySampled: number): SampleMode { - return xraySampled === 1 ? SampleMode.USER_KEEP : SampleMode.USER_REJECT; -} - -export function convertToAPMTraceID(xrayTraceID: string): string | undefined { - const parts = xrayTraceID.split("-"); - if (parts.length < 3) { - return; - } - const lastPart = parts[2]; - if (lastPart.length !== 24) { - return; - } - - // We want to turn the last 63 bits into a decimal number in a string representation - // Unfortunately, all numbers in javascript are represented by float64 bit numbers, which - // means we can't parse 64 bit integers accurately. - const hex = new BigNumber(lastPart, 16); - if (hex.isNaN()) { - return; - } - // Toggle off the 64th bit - const last63Bits = hex.mod(new BigNumber("8000000000000000", 16)); - return last63Bits.toString(10); -} - -export function convertToAPMParentID(xrayParentID: string): string | undefined { - if (xrayParentID.length !== 16) { - return; - } - const hex = new BigNumber(xrayParentID, 16); - if (hex.isNaN()) { - return; - } - 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, - }; -} diff --git a/src/trace/context.spec.ts b/src/trace/context/extractor.spec.ts similarity index 65% rename from src/trace/context.spec.ts rename to src/trace/context/extractor.spec.ts index 5a94d5e1..5bc4a9fd 100644 --- a/src/trace/context.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -1,33 +1,16 @@ -import { Context, SQSEvent } from "aws-lambda"; -import { LogLevel, setLogLevel } from "../utils"; +import { Context } from "aws-lambda"; +import { awsXrayDaemonAddressEnvVar, xrayTraceEnvVar } from "../xray-service"; import { SampleMode, - xrayBaggageSubsegmentKey, - xraySubsegmentNamespace, Source, - xrayTraceEnvVar, - awsXrayDaemonAddressEnvVar, - traceIDHeader, - parentIDHeader, - samplingPriorityHeader, -} from "./constants"; -import { - convertToAPMParentID, - convertToAPMTraceID, - convertToSampleMode, extractTraceContext, - readTraceContextFromXray, + parentIDHeader, readTraceFromEvent, - readStepFunctionContextFromEvent, - readTraceFromSQSEvent, - readTraceFromHTTPEvent, - readTraceFromLambdaContext, - hexToBinary, - deterministicMd5HashInBinary, - deterministicMd5HashToBigIntString, - readTraceFromStepFunctionsContext, - StepFunctionContext, -} from "./context"; + samplingPriorityHeader, + traceIDHeader, +} from "./extractor"; +import { readTraceFromHTTPEvent } from "./extractors/http"; +import { LogLevel, setLogLevel } from "../../utils"; let sentSegment: any; let closedSocket = false; @@ -66,158 +49,169 @@ beforeEach(() => { setLogLevel(LogLevel.NONE); }); -describe("convertToAPMTraceID", () => { - it("converts an xray trace id to a Datadog trace ID", () => { - const xrayTraceID = "1-5ce31dc2-2c779014b90ce44db5e03875"; - const traceID = convertToAPMTraceID(xrayTraceID); - expect(traceID).toEqual("4110911582297405557"); - }); - it("converts an xray trace id to a Datadog trace ID removing first bit", () => { - const xrayTraceID = "1-5ce31dc2-ac779014b90ce44db5e03875"; // Number with 64bit toggled on - const traceID = convertToAPMTraceID(xrayTraceID); - expect(traceID).toEqual("4110911582297405557"); - }); - it("returns undefined when xray trace id is too short", () => { - const xrayTraceID = "1-5ce31dc2-5e03875"; - const traceID = convertToAPMTraceID(xrayTraceID); - expect(traceID).toBeUndefined(); +describe("extractTraceContext", () => { + afterEach(() => { + process.env["_X_AMZN_TRACE_ID"] = undefined; + process.env[awsXrayDaemonAddressEnvVar] = undefined; }); + it("returns trace read from header as highest priority with no extractor", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - it("returns undefined when xray trace id is in an invalid format", () => { - const xrayTraceID = "1-2c779014b90ce44db5e03875"; - const traceID = convertToAPMTraceID(xrayTraceID); - expect(traceID).toBeUndefined(); + const result = await extractTraceContext( + { + headers: { + "x-datadog-parent-id": "797643193680388251", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405551", + }, + }, + {} as Context, + ); + expect(result).toEqual({ + parentID: "797643193680388251", + sampleMode: SampleMode.USER_KEEP, + traceID: "4110911582297405551", + source: Source.Event, + }); }); - it("returns undefined when xray trace id uses invalid characters", () => { - const xrayTraceID = "1-5ce31dc2-c779014b90ce44db5e03875;"; - const traceID = convertToAPMTraceID(xrayTraceID); - expect(traceID).toBeUndefined(); + it("returns an empty context when headers are null", async () => { + const result = await extractTraceContext( + { + headers: null, + }, + {} as Context, + ); + expect(result).toEqual(undefined); }); -}); -describe("convertToAPMParentID", () => { - it("converts an xray parent ID to an APM parent ID", () => { - const xrayParentID = "0b11cc4230d3e09e"; - const parentID = convertToAPMParentID(xrayParentID); - expect(parentID).toEqual("797643193680388254"); - }); - it("returns undefined when parent ID uses invalid characters", () => { - const xrayParentID = ";79014b90ce44db5e0;875"; - const parentID = convertToAPMParentID(xrayParentID); - expect(parentID).toBeUndefined(); - }); - it("returns undefined when parent ID is wrong size", () => { - const xrayParentID = "5e03875"; - const parentID = convertToAPMParentID(xrayParentID); - expect(parentID).toBeUndefined(); - }); -}); + it("returns trace read from event with an async extractor as the highest priority", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; -describe("convertToSampleMode", () => { - it("returns USER_KEEP if xray was sampled", () => { - const result = convertToSampleMode(1); - expect(result).toBe(SampleMode.USER_KEEP); - }); - it("returns USER_REJECT if xray wasn't sampled", () => { - const result = convertToSampleMode(0); - expect(result).toBe(SampleMode.USER_REJECT); - }); -}); + const extractor = async (event: any, context: Context) => { + const traceID = event.foo[traceIDHeader]; + const parentID = event.foo[parentIDHeader]; + const sampledHeader = event.foo[samplingPriorityHeader]; + const sampleMode = parseInt(sampledHeader, 10); -describe("readTraceContextFromXray", () => { - afterEach(() => { - process.env["_X_AMZN_TRACE_ID"] = undefined; - }); - it("returns a trace context from a valid env var", () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; - const context = readTraceContextFromXray(); - expect(context).toEqual({ - parentID: "10713633173203262661", - sampleMode: 2, - source: "xray", - traceID: "3995693151288333088", + return { + parentID, + sampleMode, + source: Source.Event, + traceID, + }; + }; + + const result = await extractTraceContext( + { + foo: { + "x-datadog-parent-id": "797643193680388251", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405551", + }, + }, + {} as Context, + extractor, + ); + expect(result).toEqual({ + parentID: "797643193680388251", + sampleMode: SampleMode.USER_KEEP, + traceID: "4110911582297405551", + source: Source.Event, }); }); - it("returns undefined when given an invalid env var", () => { - const badCases = [ - "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5", - "Root=1-5e272390-8c398be037738dc042009320", - "1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1", - "Root=1-5e272390-8c398be037738dc042009320;94ae789b969f1cc5;Sampled=1", - "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;1", - "Root=a;Parent=94ae789b969f1cc5;Sampled=1", - "Root=1-5e272390-8c398be037738dc042009320;Parent=b;Sampled=1", - undefined, - ]; - for (const badCase of badCases) { - process.env["_X_AMZN_TRACE_ID"] = badCase; - expect(readTraceContextFromXray()).toBeUndefined(); - } - }); -}); -describe("readTraceFromEvent", () => { - it("can read well formed event with headers", () => { - const result = readTraceFromEvent({ - headers: { - "x-datadog-parent-id": "797643193680388254", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405557", + it("returns trace read from event with a synchronous extractor as the highest priority", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; + + const extractor = (event: any, context: Context) => { + const traceID = event.foo[traceIDHeader]; + const parentID = event.foo[parentIDHeader]; + const sampledHeader = event.foo[samplingPriorityHeader]; + const sampleMode = parseInt(sampledHeader, 10); + + return { + parentID, + sampleMode, + source: Source.Event, + traceID, + }; + }; + + const result = await extractTraceContext( + { + foo: { + "x-datadog-parent-id": "797643193680388251", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405551", + }, }, - }); + {} as Context, + extractor, + ); expect(result).toEqual({ - parentID: "797643193680388254", + parentID: "797643193680388251", sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405557", + traceID: "4110911582297405551", source: Source.Event, }); }); - it("can read from appsync source", () => { - const result = readTraceFromEvent({ - info: { - selectionSetGraphQL: "{ items }", - }, - request: { - headers: { - "x-datadog-parent-id": "797643193680388254", + + it("handles gracefully errors in extractors", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; + + const extractor = (event: any, context: Context) => { + throw new Error("test"); + }; + + const result = await extractTraceContext( + { + foo: { + "x-datadog-parent-id": "797643193680388251", "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405557", + "x-datadog-trace-id": "4110911582297405551", }, }, - }); + {} as Context, + extractor, + ); expect(result).toEqual({ parentID: "797643193680388254", sampleMode: SampleMode.USER_KEEP, traceID: "4110911582297405557", - source: Source.Event, + source: "xray", }); }); - it("can read from sqs source", () => { - const result = readTraceFromEvent({ - Records: [ - { - body: "Hello world", - attributes: { - ApproximateReceiveCount: "1", - SentTimestamp: "1605544528092", - SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", - ApproximateFirstReceiveTimestamp: "1605544528094", - }, - messageAttributes: { - _datadog: { - stringValue: - '{"x-datadog-trace-id":"4555236104497098341","x-datadog-parent-id":"3369753143434738315","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', - stringListValues: [], - binaryListValues: [], - dataType: "String", + it("returns trace read from SQS metadata as second highest priority", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; + + const result = await extractTraceContext( + { + Records: [ + { + body: "Hello world", + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1605544528092", + SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", + ApproximateFirstReceiveTimestamp: "1605544528094", }, + messageAttributes: { + _datadog: { + stringValue: + '{"x-datadog-trace-id":"4555236104497098341","x-datadog-parent-id":"3369753143434738315","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', + stringListValues: [], + binaryListValues: [], + dataType: "String", + }, + }, + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", + awsRegion: "eu-west-1", }, - eventSource: "aws:sqs", - eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", - awsRegion: "eu-west-1", - }, - ], - }); + ], + }, + {} as Context, + ); expect(result).toEqual({ parentID: "3369753143434738315", sampleMode: SampleMode.AUTO_KEEP, @@ -225,28 +219,280 @@ describe("readTraceFromEvent", () => { source: Source.Event, }); }); - - it("can parse a traced authorizer source", () => { - const result = readTraceFromHTTPEvent({ - requestContext: { - resourceId: "oozq9u", - authorizer: { - _datadog: - "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiIyMzg5NTg5OTU0MDI2MDkwMjk2IiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjIzODk1ODk5NTQwMjYwOTAyOTYiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIiwieC1kYXRhZG9nLXBhcmVudC1zcGFuLWZpbmlzaC10aW1lIjoxNjYwOTM5ODk5MjMzLCJ4LWRhdGFkb2ctYXV0aG9yaXppbmctcmVxdWVzdGlkIjoicmFuZG9tLWlkIn0==", - principalId: "foo", - integrationLatency: 71, - preserve: "this key set by a customer", + it("returns trace read from Lambda Context as third highest priority", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; + const lambdaContext: Context = { + clientContext: { + custom: { + _datadog: { + "x-datadog-trace-id": "4555236104497098341", + "x-datadog-parent-id": "3369753143434738315", + "x-datadog-sampled": "1", + "x-datadog-sampling-priority": "1", + }, }, - stage: "dev", - requestId: "random-id", }, - httpMethod: "GET", - resource: "/hello", - }); + } as any; + const result = await extractTraceContext( + { + Records: [ + { + body: "Hello world", + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1605544528092", + SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", + ApproximateFirstReceiveTimestamp: "1605544528094", + }, + messageAttributes: { + _datadog: { + stringValue: '{"x-datadog-parent-id":"666","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', + stringListValues: [], + binaryListValues: [], + dataType: "String", + }, + }, + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", + awsRegion: "eu-west-1", + }, + ], + }, + lambdaContext, + ); expect(result).toEqual({ - parentID: "2389589954026090296", - sampleMode: 1, - source: "event", + parentID: "3369753143434738315", + sampleMode: SampleMode.AUTO_KEEP, + traceID: "4555236104497098341", + source: Source.Event, + }); + }); + it("returns trace read from env if no headers present", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; + + const result = await extractTraceContext({}, {} as Context); + expect(result).toEqual({ + parentID: "797643193680388254", + sampleMode: SampleMode.USER_KEEP, + traceID: "4110911582297405557", + source: "xray", + }); + }); + it("returns trace read from env if no headers present", async () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; + + const result = await extractTraceContext({}, {} as Context); + expect(result).toEqual({ + parentID: "797643193680388254", + sampleMode: SampleMode.USER_KEEP, + traceID: "4110911582297405557", + source: "xray", + }); + }); + it("adds datadog metadata segment to xray when trace context is in event", async () => { + jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); + process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; + process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000"; + + const result = await extractTraceContext( + { + headers: { + "x-datadog-parent-id": "797643193680388251", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405551", + }, + }, + {} as Context, + ); + + expect(sentSegment instanceof Buffer).toBeTruthy(); + expect(closedSocket).toBeTruthy(); + const sentMessage = sentSegment.toString(); + expect(sentMessage).toMatchInlineSnapshot(` + "{\\"format\\": \\"json\\", \\"version\\": 1} + {\\"id\\":\\"11111\\",\\"trace_id\\":\\"1-5e272390-8c398be037738dc042009320\\",\\"parent_id\\":\\"94ae789b969f1cc5\\",\\"name\\":\\"datadog-metadata\\",\\"start_time\\":1487076708,\\"end_time\\":1487076708,\\"type\\":\\"subsegment\\",\\"metadata\\":{\\"datadog\\":{\\"trace\\":{\\"parent-id\\":\\"797643193680388251\\",\\"sampling-priority\\":\\"2\\",\\"trace-id\\":\\"4110911582297405551\\"}}}}" + `); + }); + it("skips adding datadog metadata to x-ray when daemon isn't present", async () => { + jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); + process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; + + const result = await extractTraceContext( + { + headers: { + "x-datadog-parent-id": "797643193680388251", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405551", + }, + }, + {} as Context, + ); + + expect(sentSegment).toBeUndefined(); + }); + + it("returns trace read from step functions event with the extractor as the highest priority", async () => { + const stepFunctionEvent = { + MyInput: "MyValue", + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + Input: { + MyInput: "MyValue", + }, + Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", + RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + StartTime: "2022-12-08T21:08:17.924Z", + }, + State: { + Name: "step-one", + EnteredTime: "2022-12-08T21:08:19.224Z", + RetryCount: 2, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + Name: "my-state-machine", + }, + }; + + const result = await extractTraceContext(stepFunctionEvent, {} as Context, undefined); + expect(result).toEqual({ + parentID: "4602916161841036335", + sampleMode: 1, + traceID: "947965466153612645", + source: "event", + }); + }); + + it("skips adding datadog metadata to x-ray when x-ray trace isn't sampled", async () => { + jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); + process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0"; + process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000"; + + const result = await extractTraceContext( + { + headers: { + "x-datadog-parent-id": "797643193680388251", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405551", + }, + }, + {} as Context, + ); + + expect(sentSegment).toBeUndefined(); + }); + + it("adds step function metadata to xray", async () => { + const stepFunctionEvent = { + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", + RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + StartTime: "2022-12-08T21:08:17.924Z", + Input: { + MyInput: "MyValue", + }, + }, + State: { + Name: "step-one", + EnteredTime: "2022-12-08T21:08:19.224Z", + RetryCount: 2, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + Name: "my-state-machine", + }, + } as const; + + jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); + process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; + process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000"; + + await extractTraceContext(stepFunctionEvent, {} as Context); + expect(sentSegment instanceof Buffer).toBeTruthy(); + + expect(closedSocket).toBeTruthy(); + + const sentMessage = sentSegment.toString(); + expect(sentMessage).toMatchInlineSnapshot(` + "{\\"format\\": \\"json\\", \\"version\\": 1} + {\\"id\\":\\"11111\\",\\"trace_id\\":\\"1-5e272390-8c398be037738dc042009320\\",\\"parent_id\\":\\"94ae789b969f1cc5\\",\\"name\\":\\"datadog-metadata\\",\\"start_time\\":1487076708,\\"end_time\\":1487076708,\\"type\\":\\"subsegment\\",\\"metadata\\":{\\"datadog\\":{\\"root_span_metadata\\":{\\"step_function.execution_name\\":\\"85a9933e-9e11-83dc-6a61-b92367b6c3be\\",\\"step_function.execution_id\\":\\"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf\\",\\"step_function.execution_input\\":{\\"MyInput\\":\\"MyValue\\"},\\"step_function.execution_role_arn\\":\\"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03\\",\\"step_function.execution_start_time\\":\\"2022-12-08T21:08:17.924Z\\",\\"step_function.state_entered_time\\":\\"2022-12-08T21:08:19.224Z\\",\\"step_function.state_machine_arn\\":\\"arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential\\",\\"step_function.state_machine_name\\":\\"my-state-machine\\",\\"step_function.state_name\\":\\"step-one\\",\\"step_function.state_retry_count\\":2}}}}" + `); + }); +}); + +describe("readTraceFromEvent", () => { + it("can read well formed event with headers", () => { + const result = readTraceFromEvent({ + headers: { + "x-datadog-parent-id": "797643193680388254", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405557", + }, + }); + expect(result).toEqual({ + parentID: "797643193680388254", + sampleMode: SampleMode.USER_KEEP, + traceID: "4110911582297405557", + source: Source.Event, + }); + }); + + it("can read from sqs source", () => { + const result = readTraceFromEvent({ + Records: [ + { + body: "Hello world", + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1605544528092", + SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", + ApproximateFirstReceiveTimestamp: "1605544528094", + }, + messageAttributes: { + _datadog: { + stringValue: + '{"x-datadog-trace-id":"4555236104497098341","x-datadog-parent-id":"3369753143434738315","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', + stringListValues: [], + binaryListValues: [], + dataType: "String", + }, + }, + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", + awsRegion: "eu-west-1", + }, + ], + }); + expect(result).toEqual({ + parentID: "3369753143434738315", + sampleMode: SampleMode.AUTO_KEEP, + traceID: "4555236104497098341", + source: Source.Event, + }); + }); + + it("can parse a traced authorizer source", () => { + const result = readTraceFromHTTPEvent({ + requestContext: { + resourceId: "oozq9u", + authorizer: { + _datadog: + "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiIyMzg5NTg5OTU0MDI2MDkwMjk2IiwieC1kYXRhZG9nLXBhcmVudC1pZCI6IjIzODk1ODk5NTQwMjYwOTAyOTYiLCJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiIxIiwieC1kYXRhZG9nLXBhcmVudC1zcGFuLWZpbmlzaC10aW1lIjoxNjYwOTM5ODk5MjMzLCJ4LWRhdGFkb2ctYXV0aG9yaXppbmctcmVxdWVzdGlkIjoicmFuZG9tLWlkIn0==", + principalId: "foo", + integrationLatency: 71, + preserve: "this key set by a customer", + }, + stage: "dev", + requestId: "random-id", + }, + httpMethod: "GET", + resource: "/hello", + }); + expect(result).toEqual({ + parentID: "2389589954026090296", + sampleMode: 1, + source: "event", traceID: "2389589954026090296", }); }); @@ -538,760 +784,3 @@ describe("readTraceFromEvent", () => { expect(result).toBeUndefined(); }); }); - -describe("readTraceFromHTTPEvent", () => { - it("can read well formed event with headers", () => { - const result = readTraceFromHTTPEvent({ - headers: { - "x-datadog-parent-id": "797643193680388254", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405557", - }, - }); - expect(result).toEqual({ - parentID: "797643193680388254", - sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405557", - source: Source.Event, - }); - }); - it("can read well formed headers with mixed casing", () => { - const result = readTraceFromHTTPEvent({ - headers: { - "X-Datadog-Parent-Id": "797643193680388254", - "X-Datadog-Sampling-Priority": "2", - "X-Datadog-Trace-Id": "4110911582297405557", - }, - }); - expect(result).toEqual({ - parentID: "797643193680388254", - sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405557", - source: Source.Event, - }); - }); -}); - -describe("readTraceFromSQSEvent", () => { - it("can read from sqs source", () => { - const result = readTraceFromSQSEvent({ - Records: [ - { - body: "Hello world", - attributes: { - ApproximateReceiveCount: "1", - SentTimestamp: "1605544528092", - SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", - ApproximateFirstReceiveTimestamp: "1605544528094", - }, - messageAttributes: { - _datadog: { - stringValue: - '{"x-datadog-trace-id":"4555236104497098341","x-datadog-parent-id":"3369753143434738315","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', - stringListValues: undefined, - binaryListValues: undefined, - dataType: "String", - }, - }, - eventSource: "aws:sqs", - eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", - awsRegion: "eu-west-1", - messageId: "foo", - md5OfBody: "x", - receiptHandle: "x", - }, - ], - } as unknown as SQSEvent); - expect(result).toEqual({ - parentID: "3369753143434738315", - sampleMode: SampleMode.AUTO_KEEP, - traceID: "4555236104497098341", - source: Source.Event, - }); - }); - it("can handle malformed JSON", () => { - const result = readTraceFromSQSEvent({ - Records: [ - { - body: "Hello world", - attributes: { - ApproximateReceiveCount: "1", - SentTimestamp: "1605544528092", - SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", - ApproximateFirstReceiveTimestamp: "1605544528094", - }, - messageAttributes: { - _datadog: { - stringValue: - '{asdasdasd"x-datadog-trace-id":"4555236104497098341","x-datadog-parent-id":"3369753143434738315","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', - stringListValues: undefined, - binaryListValues: undefined, - dataType: "String", - }, - }, - eventSource: "aws:sqs", - eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", - awsRegion: "eu-west-1", - messageId: "foo", - md5OfBody: "x", - receiptHandle: "x", - }, - ], - } as unknown as SQSEvent); - expect(result).toBeUndefined(); - }); -}); - -describe("readTraceFromLambdaContext", () => { - it("can read from lambda context source, legacy style", () => { - const result = readTraceFromLambdaContext({ - clientContext: { - custom: { - _datadog: { - "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 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: { - foo: "bar", - }, - }); - 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(); - }); -}); - -describe("readStepFunctionContextFromEvent", () => { - const stepFunctionEvent = { - Execution: { - Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", - Input: { - MyInput: "MyValue", - }, - Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", - RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", - StartTime: "2022-12-08T21:08:17.924Z", - }, - State: { - Name: "step-one", - EnteredTime: "2022-12-08T21:08:19.224Z", - RetryCount: 2, - }, - StateMachine: { - Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", - Name: "my-state-machine", - }, - } as const; - - it("reads a step function context from event with Execution.Input", () => { - const result = readStepFunctionContextFromEvent(stepFunctionEvent); - expect(result).toEqual({ - "step_function.execution_id": - "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", - "step_function.execution_input": { MyInput: "MyValue" }, - "step_function.execution_name": "85a9933e-9e11-83dc-6a61-b92367b6c3be", - "step_function.execution_role_arn": - "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", - "step_function.execution_start_time": "2022-12-08T21:08:17.924Z", - "step_function.state_entered_time": "2022-12-08T21:08:19.224Z", - "step_function.state_machine_arn": "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", - "step_function.state_machine_name": "my-state-machine", - "step_function.state_name": "step-one", - "step_function.state_retry_count": 2, - }); - }); - it("returns undefined when event isn't an object", () => { - const result = readStepFunctionContextFromEvent("event"); - expect(result).toBeUndefined(); - }); - it("returns undefined when event is missing datadogContext property", () => { - const result = readStepFunctionContextFromEvent({}); - expect(result).toBeUndefined(); - }); - it("returns undefined when datadogContext is missing Execution property", () => { - const result = readStepFunctionContextFromEvent({ - dd: {}, - }); - expect(result).toBeUndefined(); - }); - it("returns undefined when Execution is missing Name field", () => { - const result = readStepFunctionContextFromEvent({ - dd: { - ...stepFunctionEvent, - Execution: {}, - }, - }); - expect(result).toBeUndefined(); - }); - it("returns undefined when Name isn't a string", () => { - const result = readStepFunctionContextFromEvent({ - dd: { - ...stepFunctionEvent, - Execution: { - Name: 12345, - }, - }, - }); - expect(result).toBeUndefined(); - }); - it("returns undefined when State isn't defined", () => { - const result = readStepFunctionContextFromEvent({ - dd: { - ...stepFunctionEvent, - State: undefined, - }, - }); - expect(result).toBeUndefined(); - }); - it("returns undefined when try retry count isn't a number", () => { - const result = readStepFunctionContextFromEvent({ - dd: { - ...stepFunctionEvent, - State: { - ...stepFunctionEvent.State, - RetryCount: "1", - }, - }, - }); - expect(result).toBeUndefined(); - }); - it("returns undefined when try step name isn't a string", () => { - const result = readStepFunctionContextFromEvent({ - dd: { - ...stepFunctionEvent, - State: { - ...stepFunctionEvent.State, - Name: 1, - }, - }, - }); - expect(result).toBeUndefined(); - }); - it("returns undefined when StateMachine is undefined", () => { - const result = readStepFunctionContextFromEvent({ - dd: { - ...stepFunctionEvent, - StateMachine: undefined, - }, - }); - expect(result).toBeUndefined(); - }); - it("returns undefined when StateMachineId isn't a string", () => { - const result = readStepFunctionContextFromEvent({ - dd: { - ...stepFunctionEvent, - StateMachine: { - ...stepFunctionEvent.StateMachine, - Id: 1, - }, - }, - }); - expect(result).toBeUndefined(); - }); - it("returns undefined when StateMachineName isn't a string", () => { - const result = readStepFunctionContextFromEvent({ - dd: { - ...stepFunctionEvent, - StateMachine: { - ...stepFunctionEvent.StateMachine, - Name: 1, - }, - }, - }); - expect(result).toBeUndefined(); - }); -}); - -describe("extractTraceContext", () => { - afterEach(() => { - process.env["_X_AMZN_TRACE_ID"] = undefined; - process.env[awsXrayDaemonAddressEnvVar] = undefined; - }); - it("returns trace read from header as highest priority with no extractor", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - - const result = await extractTraceContext( - { - headers: { - "x-datadog-parent-id": "797643193680388251", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405551", - }, - }, - {} as Context, - ); - expect(result).toEqual({ - parentID: "797643193680388251", - sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405551", - source: Source.Event, - }); - }); - it("returns an empty context when headers are null", async () => { - const result = await extractTraceContext( - { - headers: null, - }, - {} as Context, - ); - expect(result).toEqual(undefined); - }); - - it("returns trace read from event with an async extractor as the highest priority", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - - const extractor = async (event: any, context: Context) => { - const traceID = event.foo[traceIDHeader]; - const parentID = event.foo[parentIDHeader]; - const sampledHeader = event.foo[samplingPriorityHeader]; - const sampleMode = parseInt(sampledHeader, 10); - - return { - parentID, - sampleMode, - source: Source.Event, - traceID, - }; - }; - - const result = await extractTraceContext( - { - foo: { - "x-datadog-parent-id": "797643193680388251", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405551", - }, - }, - {} as Context, - extractor, - ); - expect(result).toEqual({ - parentID: "797643193680388251", - sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405551", - source: Source.Event, - }); - }); - - it("returns trace read from event with a synchronous extractor as the highest priority", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - - const extractor = (event: any, context: Context) => { - const traceID = event.foo[traceIDHeader]; - const parentID = event.foo[parentIDHeader]; - const sampledHeader = event.foo[samplingPriorityHeader]; - const sampleMode = parseInt(sampledHeader, 10); - - return { - parentID, - sampleMode, - source: Source.Event, - traceID, - }; - }; - - const result = await extractTraceContext( - { - foo: { - "x-datadog-parent-id": "797643193680388251", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405551", - }, - }, - {} as Context, - extractor, - ); - expect(result).toEqual({ - parentID: "797643193680388251", - sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405551", - source: Source.Event, - }); - }); - - it("handles gracefully errors in extractors", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - - const extractor = (event: any, context: Context) => { - throw new Error("test"); - }; - - const result = await extractTraceContext( - { - foo: { - "x-datadog-parent-id": "797643193680388251", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405551", - }, - }, - {} as Context, - extractor, - ); - expect(result).toEqual({ - parentID: "797643193680388254", - sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405557", - source: "xray", - }); - }); - it("returns trace read from SQS metadata as second highest priority", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - - const result = await extractTraceContext( - { - Records: [ - { - body: "Hello world", - attributes: { - ApproximateReceiveCount: "1", - SentTimestamp: "1605544528092", - SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", - ApproximateFirstReceiveTimestamp: "1605544528094", - }, - messageAttributes: { - _datadog: { - stringValue: - '{"x-datadog-trace-id":"4555236104497098341","x-datadog-parent-id":"3369753143434738315","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', - stringListValues: [], - binaryListValues: [], - dataType: "String", - }, - }, - eventSource: "aws:sqs", - eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", - awsRegion: "eu-west-1", - }, - ], - }, - {} as Context, - ); - expect(result).toEqual({ - parentID: "3369753143434738315", - sampleMode: SampleMode.AUTO_KEEP, - traceID: "4555236104497098341", - source: Source.Event, - }); - }); - it("returns trace read from Lambda Context as third highest priority", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - const lambdaContext: Context = { - clientContext: { - custom: { - _datadog: { - "x-datadog-trace-id": "4555236104497098341", - "x-datadog-parent-id": "3369753143434738315", - "x-datadog-sampled": "1", - "x-datadog-sampling-priority": "1", - }, - }, - }, - } as any; - const result = await extractTraceContext( - { - Records: [ - { - body: "Hello world", - attributes: { - ApproximateReceiveCount: "1", - SentTimestamp: "1605544528092", - SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", - ApproximateFirstReceiveTimestamp: "1605544528094", - }, - messageAttributes: { - _datadog: { - stringValue: '{"x-datadog-parent-id":"666","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', - stringListValues: [], - binaryListValues: [], - dataType: "String", - }, - }, - eventSource: "aws:sqs", - eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", - awsRegion: "eu-west-1", - }, - ], - }, - lambdaContext, - ); - expect(result).toEqual({ - parentID: "3369753143434738315", - sampleMode: SampleMode.AUTO_KEEP, - traceID: "4555236104497098341", - source: Source.Event, - }); - }); - it("returns trace read from env if no headers present", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - - const result = await extractTraceContext({}, {} as Context); - expect(result).toEqual({ - parentID: "797643193680388254", - sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405557", - source: "xray", - }); - }); - it("returns trace read from env if no headers present", async () => { - process.env["_X_AMZN_TRACE_ID"] = "Root=1-5ce31dc2-2c779014b90ce44db5e03875;Parent=0b11cc4230d3e09e;Sampled=1"; - - const result = await extractTraceContext({}, {} as Context); - expect(result).toEqual({ - parentID: "797643193680388254", - sampleMode: SampleMode.USER_KEEP, - traceID: "4110911582297405557", - source: "xray", - }); - }); - it("adds datadog metadata segment to xray when trace context is in event", async () => { - jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); - process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; - process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000"; - - const result = await extractTraceContext( - { - headers: { - "x-datadog-parent-id": "797643193680388251", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405551", - }, - }, - {} as Context, - ); - - expect(sentSegment instanceof Buffer).toBeTruthy(); - expect(closedSocket).toBeTruthy(); - const sentMessage = sentSegment.toString(); - expect(sentMessage).toMatchInlineSnapshot(` - "{\\"format\\": \\"json\\", \\"version\\": 1} - {\\"id\\":\\"11111\\",\\"trace_id\\":\\"1-5e272390-8c398be037738dc042009320\\",\\"parent_id\\":\\"94ae789b969f1cc5\\",\\"name\\":\\"datadog-metadata\\",\\"start_time\\":1487076708,\\"end_time\\":1487076708,\\"type\\":\\"subsegment\\",\\"metadata\\":{\\"datadog\\":{\\"trace\\":{\\"parent-id\\":\\"797643193680388251\\",\\"sampling-priority\\":\\"2\\",\\"trace-id\\":\\"4110911582297405551\\"}}}}" - `); - }); - it("skips adding datadog metadata to x-ray when daemon isn't present", async () => { - jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); - process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; - - const result = await extractTraceContext( - { - headers: { - "x-datadog-parent-id": "797643193680388251", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405551", - }, - }, - {} as Context, - ); - - expect(sentSegment).toBeUndefined(); - }); - - it("returns trace read from step functions event with the extractor as the highest priority", async () => { - const stepFunctionEvent = { - MyInput: "MyValue", - Execution: { - Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", - Input: { - MyInput: "MyValue", - }, - Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", - RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", - StartTime: "2022-12-08T21:08:17.924Z", - }, - State: { - Name: "step-one", - EnteredTime: "2022-12-08T21:08:19.224Z", - RetryCount: 2, - }, - StateMachine: { - Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", - Name: "my-state-machine", - }, - }; - - const result = await extractTraceContext(stepFunctionEvent, {} as Context, undefined); - expect(result).toEqual({ - parentID: "4602916161841036335", - sampleMode: 1, - traceID: "947965466153612645", - source: "event", - }); - }); - - it("skips adding datadog metadata to x-ray when x-ray trace isn't sampled", async () => { - jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); - process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0"; - process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000"; - - const result = await extractTraceContext( - { - headers: { - "x-datadog-parent-id": "797643193680388251", - "x-datadog-sampling-priority": "2", - "x-datadog-trace-id": "4110911582297405551", - }, - }, - {} as Context, - ); - - expect(sentSegment).toBeUndefined(); - }); - - it("adds step function metadata to xray", async () => { - const stepFunctionEvent = { - Execution: { - Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", - Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", - RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", - StartTime: "2022-12-08T21:08:17.924Z", - Input: { - MyInput: "MyValue", - }, - }, - State: { - Name: "step-one", - EnteredTime: "2022-12-08T21:08:19.224Z", - RetryCount: 2, - }, - StateMachine: { - Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", - Name: "my-state-machine", - }, - } as const; - - jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); - process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; - process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000"; - - await extractTraceContext(stepFunctionEvent, {} as Context); - expect(sentSegment instanceof Buffer).toBeTruthy(); - - expect(closedSocket).toBeTruthy(); - - const sentMessage = sentSegment.toString(); - expect(sentMessage).toMatchInlineSnapshot(` - "{\\"format\\": \\"json\\", \\"version\\": 1} - {\\"id\\":\\"11111\\",\\"trace_id\\":\\"1-5e272390-8c398be037738dc042009320\\",\\"parent_id\\":\\"94ae789b969f1cc5\\",\\"name\\":\\"datadog-metadata\\",\\"start_time\\":1487076708,\\"end_time\\":1487076708,\\"type\\":\\"subsegment\\",\\"metadata\\":{\\"datadog\\":{\\"root_span_metadata\\":{\\"step_function.execution_name\\":\\"85a9933e-9e11-83dc-6a61-b92367b6c3be\\",\\"step_function.execution_id\\":\\"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf\\",\\"step_function.execution_input\\":{\\"MyInput\\":\\"MyValue\\"},\\"step_function.execution_role_arn\\":\\"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03\\",\\"step_function.execution_start_time\\":\\"2022-12-08T21:08:17.924Z\\",\\"step_function.state_entered_time\\":\\"2022-12-08T21:08:19.224Z\\",\\"step_function.state_machine_arn\\":\\"arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential\\",\\"step_function.state_machine_name\\":\\"my-state-machine\\",\\"step_function.state_name\\":\\"step-one\\",\\"step_function.state_retry_count\\":2}}}}" - `); - }); -}); - -describe.each([ - ["0", "0000"], - ["1", "0001"], - ["2", "0010"], - ["3", "0011"], - ["4", "0100"], - ["5", "0101"], - ["6", "0110"], - ["7", "0111"], - ["8", "1000"], - ["9", "1001"], - ["a", "1010"], - ["b", "1011"], - ["c", "1100"], - ["d", "1101"], - ["e", "1110"], - ["f", "1111"], -])(`test hexToBinary`, (hex, expected) => { - test(`${hex} to binary returns ${expected}`, () => { - expect(hexToBinary(hex)).toBe(expected); - }); -}); - -describe("test_deterministicMd5HashInBinary", () => { - it("test same hashing is generated as logs-backend for a random string", () => { - const actual = deterministicMd5HashInBinary("some_testing_random_string"); - expect(actual).toEqual("0001111100111110001000110110011110010111000110001001001111110001"); - }); - - it("test same hashing is generated as logs-backend for an execution id", () => { - const actual = deterministicMd5HashInBinary( - "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d041f4", - ); - expect(actual).toEqual("0010010000101100100000101011111101111100110110001110111100111101"); - }); - - it("test same hashing is generated as logs-backend for another execution id", () => { - const actual = deterministicMd5HashInBinary( - "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111", - ); - expect(actual).toEqual("0010001100110000011011011111010000100111100000110000100100101010"); - }); - - it("test same hashing is generated as logs-backend for execution id # state name # entered time", () => { - const actual = deterministicMd5HashInBinary( - "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111#step-one#2022-12-08T21:08:19.224Z", - ); - expect(actual).toEqual("0110111110000000010011011001111101110011100111000000011010100001"); - }); - - it("test hashing different strings would generate different hashes", () => { - const times = 20; - for (let i = 0; i < times; i++) { - for (let j = i + 1; j < times; j++) { - expect(deterministicMd5HashInBinary(i.toString())).not.toMatch(deterministicMd5HashInBinary(j.toString())); - } - } - }); - - it("test always leading with 0", () => { - for (let i = 0; i < 20; i++) { - expect(deterministicMd5HashInBinary(i.toString()).substring(0, 1)).toMatch("0"); - } - }); -}); - -describe("test_deterministicMd5HashToBigIntString", () => { - it("test same hashing number is generated as logs-backend for a random string", () => { - const actual = deterministicMd5HashToBigIntString("some_testing_random_string"); - expect(actual).toEqual("2251275791555400689"); - }); - - it("test same hashing number is generated as logs-backend for execution id # state name # entered time", () => { - const actual = deterministicMd5HashToBigIntString( - "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111#step-one#2022-12-08T21:08:19.224Z", - ); - expect(actual).toEqual("8034507082463708833"); - }); -}); diff --git a/src/trace/context/extractor.ts b/src/trace/context/extractor.ts new file mode 100644 index 00000000..42569d2a --- /dev/null +++ b/src/trace/context/extractor.ts @@ -0,0 +1,168 @@ +import { Context } from "aws-lambda"; +import { TraceExtractor } from "../listener"; +import { logDebug, logError } from "../../utils"; +import { readTraceFromHTTPEvent } from "./extractors/http"; +import { readTraceFromSNSEvent } from "./extractors/sns"; +import { readTraceFromSNSSQSEvent } from "./extractors/sns-sqs"; +import { readTraceFromEBSQSEvent } from "./extractors/event-bridge-sqs"; +import { readTraceFromAppSyncEvent } from "./extractors/app-sync"; +import { readTraceFromSQSEvent } from "./extractors/sqs"; +import { readTraceFromKinesisEvent } from "./extractors/kinesis"; +import { readTraceFromEventbridgeEvent } from "./extractors/event-bridge"; +import { + isAppSyncResolverEvent, + isEBSQSEvent, + isEventBridgeEvent, + isKinesisStreamEvent, + isSNSEvent, + isSNSSQSEvent, + isSQSEvent, +} from "../../utils/event-type-guards"; +import { readTraceFromLambdaContext } from "./extractors/lambda-context"; +import { readStepFunctionContextFromEvent } from "../step-function-service"; +import { addStepFunctionContextToXray, addTraceContextToXray, readTraceContextFromXray } from "../xray-service"; +import { readTraceFromStepFunctionsContext } from "./extractors/step-function"; + +export enum SampleMode { + USER_REJECT = -1, + AUTO_REJECT = 0, + AUTO_KEEP = 1, + USER_KEEP = 2, +} +export enum Source { + Xray = "xray", + Event = "event", + DDTrace = "ddtrace", +} + +export interface TraceContext { + traceID: string; + parentID: string; + sampleMode: SampleMode; + source: Source; +} + +export const traceIDHeader = "x-datadog-trace-id"; +export const parentIDHeader = "x-datadog-parent-id"; +export const samplingPriorityHeader = "x-datadog-sampling-priority"; + +/** + * Reads the trace context from either an incoming lambda event, or the current xray segment. + * @param event An incoming lambda event. This must have incoming trace headers in order to be read. + */ +export async function extractTraceContext( + event: any, + context: Context, + extractor?: TraceExtractor, + decodeAuthorizerContext: boolean = true, +): Promise { + let trace; + + if (extractor) { + try { + trace = await extractor(event, context); + logDebug(`extracted trace context from the custom extractor`, { trace }); + } catch (error) { + if (error instanceof Error) { + logError("custom extractor function failed", error as Error); + } + } + } + + if (!trace) { + trace = readTraceFromEvent(event, decodeAuthorizerContext); + } + + if (!trace) { + trace = readTraceFromLambdaContext(context); + } + + const stepFuncContext = readStepFunctionContextFromEvent(event); + if (stepFuncContext) { + try { + addStepFunctionContextToXray(stepFuncContext); + } catch (error) { + if (error instanceof Error) { + logError("couldn't add step function metadata to xray", error as Error); + } + } + if (trace === undefined) { + trace = readTraceFromStepFunctionsContext(stepFuncContext); + if (trace !== undefined) { + return trace; + } + } + } + + if (trace !== undefined) { + try { + addTraceContextToXray(trace); + logDebug(`added trace context to xray metadata`, { trace }); + } catch (error) { + // This might fail if running in an environment where xray isn't set up, (like for local development). + if (error instanceof Error) { + logError("couldn't add trace context to xray metadata", error as Error); + } + } + return trace; + } + return readTraceContextFromXray(); +} + +export function readTraceFromEvent(event: any, decodeAuthorizerContext: boolean = true): TraceContext | undefined { + if (!event || typeof event !== "object") { + return; + } + + if (event.headers !== null && typeof event.headers === "object") { + return readTraceFromHTTPEvent(event, decodeAuthorizerContext); + } + + if (isSNSEvent(event)) { + return readTraceFromSNSEvent(event); + } + + if (isSNSSQSEvent(event)) { + return readTraceFromSNSSQSEvent(event); + } + + if (isEBSQSEvent(event)) { + return readTraceFromEBSQSEvent(event); + } + + if (isAppSyncResolverEvent(event)) { + return readTraceFromAppSyncEvent(event); + } + + if (isSQSEvent(event)) { + return readTraceFromSQSEvent(event); + } + if (isKinesisStreamEvent(event)) { + return readTraceFromKinesisEvent(event); + } + + if (isEventBridgeEvent(event)) { + return readTraceFromEventbridgeEvent(event); + } + + return; +} + +export 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, + }; +} diff --git a/src/trace/context/extractors/app-sync.spec.ts b/src/trace/context/extractors/app-sync.spec.ts new file mode 100644 index 00000000..f3d3c100 --- /dev/null +++ b/src/trace/context/extractors/app-sync.spec.ts @@ -0,0 +1,24 @@ +import { SampleMode, Source, readTraceFromEvent } from "../extractor"; + +describe("readTraceFromAppSyncEvent", () => { + it("can read from appsync source", () => { + const result = readTraceFromEvent({ + info: { + selectionSetGraphQL: "{ items }", + }, + request: { + headers: { + "x-datadog-parent-id": "797643193680388254", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405557", + }, + }, + }); + expect(result).toEqual({ + parentID: "797643193680388254", + sampleMode: SampleMode.USER_KEEP, + traceID: "4110911582297405557", + source: Source.Event, + }); + }); +}); diff --git a/src/trace/context/extractors/app-sync.ts b/src/trace/context/extractors/app-sync.ts new file mode 100644 index 00000000..c071b163 --- /dev/null +++ b/src/trace/context/extractors/app-sync.ts @@ -0,0 +1,7 @@ +import { TraceContext } from "../extractor"; +import { readTraceFromHTTPEvent } from "./http"; + +export function readTraceFromAppSyncEvent(event: any): TraceContext | undefined { + event.headers = event.request.headers; + return readTraceFromHTTPEvent(event, false); +} diff --git a/src/trace/context/extractors/event-bridge-sqs.ts b/src/trace/context/extractors/event-bridge-sqs.ts new file mode 100644 index 00000000..acfb3fa0 --- /dev/null +++ b/src/trace/context/extractors/event-bridge-sqs.ts @@ -0,0 +1,22 @@ +import { EventBridgeEvent, SQSEvent } from "aws-lambda"; +import { TraceContext, exportTraceData } from "../extractor"; +import { logDebug } from "../../../utils"; + +export function readTraceFromEBSQSEvent(event: SQSEvent): TraceContext | undefined { + if (event?.Records?.[0]?.body) { + try { + const parsedBody = JSON.parse(event.Records[0].body) as EventBridgeEvent; + if (parsedBody?.detail?._datadog) { + const trace = exportTraceData(parsedBody.detail._datadog); + + logDebug(`extracted trace context from EventBridge SQS event`, { trace, event }); + return trace; + } + } catch (err) { + if (err instanceof Error) { + logDebug("Error parsing EventBridge SQS message trace data", err as Error); + } + return; + } + } +} diff --git a/src/trace/context/extractors/event-bridge.ts b/src/trace/context/extractors/event-bridge.ts new file mode 100644 index 00000000..c9dedab2 --- /dev/null +++ b/src/trace/context/extractors/event-bridge.ts @@ -0,0 +1,18 @@ +import { EventBridgeEvent } from "aws-lambda"; +import { TraceContext, exportTraceData } from "../extractor"; +import { logDebug } from "../../../utils"; + +export function readTraceFromEventbridgeEvent(event: EventBridgeEvent): TraceContext | undefined { + if (event?.detail?._datadog) { + try { + const trace = exportTraceData(event.detail._datadog); + logDebug(`extracted trace context from Eventbridge event`, { trace, event }); + return trace; + } catch (err) { + if (err instanceof Error) { + logDebug("Error parsing Eventbridge trace data", err as Error); + } + return; + } + } +} diff --git a/src/trace/context/extractors/http.spec.ts b/src/trace/context/extractors/http.spec.ts new file mode 100644 index 00000000..d0e3ed79 --- /dev/null +++ b/src/trace/context/extractors/http.spec.ts @@ -0,0 +1,35 @@ +import { SampleMode, Source } from "../extractor"; +import { readTraceFromHTTPEvent } from "./http"; + +describe("readTraceFromHTTPEvent", () => { + it("can read well formed event with headers", () => { + const result = readTraceFromHTTPEvent({ + headers: { + "x-datadog-parent-id": "797643193680388254", + "x-datadog-sampling-priority": "2", + "x-datadog-trace-id": "4110911582297405557", + }, + }); + expect(result).toEqual({ + parentID: "797643193680388254", + sampleMode: SampleMode.USER_KEEP, + traceID: "4110911582297405557", + source: Source.Event, + }); + }); + it("can read well formed headers with mixed casing", () => { + const result = readTraceFromHTTPEvent({ + headers: { + "X-Datadog-Parent-Id": "797643193680388254", + "X-Datadog-Sampling-Priority": "2", + "X-Datadog-Trace-Id": "4110911582297405557", + }, + }); + expect(result).toEqual({ + parentID: "797643193680388254", + sampleMode: SampleMode.USER_KEEP, + traceID: "4110911582297405557", + source: Source.Event, + }); + }); +}); diff --git a/src/trace/context/extractors/http.ts b/src/trace/context/extractors/http.ts new file mode 100644 index 00000000..bb2daf8d --- /dev/null +++ b/src/trace/context/extractors/http.ts @@ -0,0 +1,50 @@ +import { logDebug } from "../../../utils"; +import { eventSubTypes, parseEventSourceSubType } from "../../trigger"; +import { TraceContext, exportTraceData } from "../extractor"; + +export const authorizingRequestIdHeader = "x-datadog-authorizing-requestid"; + +export function readTraceFromHTTPEvent(event: any, decodeAuthorizerContext: boolean = true): TraceContext | undefined { + if (decodeAuthorizerContext) { + // need to set the trace context if using authorizer lambda in authorizing (non-cached) cases + try { + const eventSourceSubType: eventSubTypes = parseEventSourceSubType(event); + const injectedAuthorizerData = getInjectedAuthorizerData(event, eventSourceSubType); + if (injectedAuthorizerData !== null) { + return exportTraceData(injectedAuthorizerData); + } + } catch (error) { + logDebug(`unable to extract trace context from authorizer event.`, { error }); + } + } + + const headers = event.headers; + const lowerCaseHeaders: { [key: string]: string } = {}; + + for (const key of Object.keys(headers)) { + lowerCaseHeaders[key.toLowerCase()] = headers[key]; + } + + const trace = exportTraceData(lowerCaseHeaders); + + logDebug(`extracted trace context from http event`, { trace, event }); + return trace; +} + +export function getInjectedAuthorizerData(event: any, eventSourceSubType: eventSubTypes) { + const authorizerHeaders = event?.requestContext?.authorizer; + if (!authorizerHeaders) return null; + const rawDatadogData = + eventSourceSubType === eventSubTypes.apiGatewayV2 ? authorizerHeaders.lambda._datadog : authorizerHeaders._datadog; + if (!rawDatadogData) return null; + const injectedData = JSON.parse(Buffer.from(rawDatadogData, "base64").toString()); + // use the injected requestId to tell if it's the authorizing invocation (not cached) + if ( + authorizerHeaders.integrationLatency > 0 || + event.requestContext.requestId === injectedData[authorizingRequestIdHeader] + ) { + return injectedData; + } else { + return null; + } +} diff --git a/src/trace/context/extractors/kinesis.ts b/src/trace/context/extractors/kinesis.ts new file mode 100644 index 00000000..dce1f032 --- /dev/null +++ b/src/trace/context/extractors/kinesis.ts @@ -0,0 +1,21 @@ +import { KinesisStreamEvent } from "aws-lambda"; +import { TraceContext, exportTraceData } from "../extractor"; +import { logDebug } from "../../../utils"; + +export function readTraceFromKinesisEvent(event: KinesisStreamEvent): TraceContext | undefined { + if (event?.Records?.[0]?.kinesis?.data) { + try { + const parsedBody = JSON.parse(Buffer.from(event.Records[0].kinesis.data, "base64").toString("ascii")) as any; + if (parsedBody && parsedBody._datadog) { + const trace = exportTraceData(parsedBody._datadog); + logDebug(`extracted trace context from Kinesis event`, { trace }); + return trace; + } + } catch (err) { + if (err instanceof Error) { + logDebug("Error parsing Kinesis message trace data", err as Error); + } + return; + } + } +} diff --git a/src/trace/context/extractors/lambda-context.spec.ts b/src/trace/context/extractors/lambda-context.spec.ts new file mode 100644 index 00000000..64824c01 --- /dev/null +++ b/src/trace/context/extractors/lambda-context.spec.ts @@ -0,0 +1,63 @@ +import { SampleMode, Source } from "../extractor"; +import { readTraceFromLambdaContext } from "./lambda-context"; + +describe("readTraceFromLambdaContext", () => { + it("can read from lambda context source, legacy style", () => { + const result = readTraceFromLambdaContext({ + clientContext: { + custom: { + _datadog: { + "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 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: { + foo: "bar", + }, + }); + 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/extractors/lambda-context.ts b/src/trace/context/extractors/lambda-context.ts new file mode 100644 index 00000000..10436db8 --- /dev/null +++ b/src/trace/context/extractors/lambda-context.ts @@ -0,0 +1,37 @@ +import { logDebug } from "../../../utils"; +import { TraceContext, exportTraceData, parentIDHeader, samplingPriorityHeader, traceIDHeader } from "../extractor"; + +export function readTraceFromLambdaContext(context: any): TraceContext | undefined { + if (!context || typeof context !== "object") { + return; + } + + const custom = context.clientContext?.custom; + + if (!custom || typeof custom !== "object") { + return; + } + let traceData = null; + + if ( + custom.hasOwnProperty("_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; + } + + const trace = exportTraceData(traceData); + logDebug(`extracted trace context from lambda context`, { trace, context }); + return trace; +} diff --git a/src/trace/context/extractors/sns-sqs.ts b/src/trace/context/extractors/sns-sqs.ts new file mode 100644 index 00000000..f2516f83 --- /dev/null +++ b/src/trace/context/extractors/sns-sqs.ts @@ -0,0 +1,33 @@ +import { SNSMessage, SQSEvent } from "aws-lambda"; +import { TraceContext, exportTraceData } from "../extractor"; +import { logDebug } from "../../../utils"; + +export function readTraceFromSNSSQSEvent(event: SQSEvent): TraceContext | undefined { + if (event?.Records?.[0]?.body) { + try { + const parsedBody = JSON.parse(event.Records[0].body) as SNSMessage; + if ( + parsedBody.MessageAttributes && + parsedBody.MessageAttributes._datadog && + parsedBody.MessageAttributes._datadog.Value + ) { + let traceData; + if (parsedBody.MessageAttributes._datadog.Type === "String") { + traceData = JSON.parse(parsedBody.MessageAttributes._datadog.Value); + } else { + const b64Decoded = Buffer.from(parsedBody.MessageAttributes._datadog.Value, "base64").toString("ascii"); + traceData = JSON.parse(b64Decoded); + } + const trace = exportTraceData(traceData); + + logDebug(`extracted trace context from SNS SQS event`, { trace, event }); + return trace; + } + } catch (err) { + if (err instanceof Error) { + logDebug("Error parsing SNS SQS message trace data", err as Error); + } + return; + } + } +} diff --git a/src/trace/context/extractors/sns.ts b/src/trace/context/extractors/sns.ts new file mode 100644 index 00000000..a3815b5c --- /dev/null +++ b/src/trace/context/extractors/sns.ts @@ -0,0 +1,27 @@ +import { SNSEvent } from "aws-lambda"; +import { TraceContext, exportTraceData } from "../extractor"; +import { logDebug } from "../../../utils"; + +export function readTraceFromSNSEvent(event: SNSEvent): TraceContext | undefined { + if (event?.Records?.[0]?.Sns?.MessageAttributes?._datadog?.Value) { + try { + let traceData; + if (event.Records[0].Sns.MessageAttributes._datadog.Type === "String") { + traceData = JSON.parse(event.Records[0].Sns.MessageAttributes._datadog.Value); + } else { + const b64Decoded = Buffer.from(event.Records[0].Sns.MessageAttributes._datadog.Value, "base64").toString( + "ascii", + ); + traceData = JSON.parse(b64Decoded); + } + const trace = exportTraceData(traceData); + logDebug(`extracted trace context from SNS event`, { trace, event }); + return trace; + } catch (err) { + if (err instanceof Error) { + logDebug("Error parsing SNS SQS message trace data", err as Error); + } + return; + } + } +} diff --git a/src/trace/context/extractors/sqs.spec.ts b/src/trace/context/extractors/sqs.spec.ts new file mode 100644 index 00000000..b3a7f3c1 --- /dev/null +++ b/src/trace/context/extractors/sqs.spec.ts @@ -0,0 +1,73 @@ +import { SQSEvent } from "aws-lambda"; +import { readTraceFromSQSEvent } from "./sqs"; +import { SampleMode, Source } from "../extractor"; + +describe("readTraceFromSQSEvent", () => { + it("can read from sqs source", () => { + const result = readTraceFromSQSEvent({ + Records: [ + { + body: "Hello world", + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1605544528092", + SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", + ApproximateFirstReceiveTimestamp: "1605544528094", + }, + messageAttributes: { + _datadog: { + stringValue: + '{"x-datadog-trace-id":"4555236104497098341","x-datadog-parent-id":"3369753143434738315","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', + stringListValues: undefined, + binaryListValues: undefined, + dataType: "String", + }, + }, + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", + awsRegion: "eu-west-1", + messageId: "foo", + md5OfBody: "x", + receiptHandle: "x", + }, + ], + } as unknown as SQSEvent); + expect(result).toEqual({ + parentID: "3369753143434738315", + sampleMode: SampleMode.AUTO_KEEP, + traceID: "4555236104497098341", + source: Source.Event, + }); + }); + it("can handle malformed JSON", () => { + const result = readTraceFromSQSEvent({ + Records: [ + { + body: "Hello world", + attributes: { + ApproximateReceiveCount: "1", + SentTimestamp: "1605544528092", + SenderId: "AROAYYB64AB3JHSRKO6XR:sqs-trace-dev-producer", + ApproximateFirstReceiveTimestamp: "1605544528094", + }, + messageAttributes: { + _datadog: { + stringValue: + '{asdasdasd"x-datadog-trace-id":"4555236104497098341","x-datadog-parent-id":"3369753143434738315","x-datadog-sampled":"1","x-datadog-sampling-priority":"1"}', + stringListValues: undefined, + binaryListValues: undefined, + dataType: "String", + }, + }, + eventSource: "aws:sqs", + eventSourceARN: "arn:aws:sqs:eu-west-1:601427279990:metal-queue", + awsRegion: "eu-west-1", + messageId: "foo", + md5OfBody: "x", + receiptHandle: "x", + }, + ], + } as unknown as SQSEvent); + expect(result).toBeUndefined(); + }); +}); diff --git a/src/trace/context/extractors/sqs.ts b/src/trace/context/extractors/sqs.ts new file mode 100644 index 00000000..2acc866b --- /dev/null +++ b/src/trace/context/extractors/sqs.ts @@ -0,0 +1,23 @@ +import { SQSEvent } from "aws-lambda"; +import { TraceContext, exportTraceData } from "../extractor"; +import { logDebug } from "../../../utils"; + +export function readTraceFromSQSEvent(event: SQSEvent): TraceContext | undefined { + if (event?.Records?.[0]?.messageAttributes?._datadog?.stringValue) { + const traceHeaders = event.Records[0].messageAttributes._datadog.stringValue; + + try { + const trace = exportTraceData(JSON.parse(traceHeaders)); + + logDebug(`extracted trace context from sqs event`, { trace, event }); + return trace; + } catch (err) { + if (err instanceof Error) { + logDebug("Error parsing SQS message trace data", err as Error); + } + return; + } + } + + return; +} diff --git a/src/trace/context/extractors/step-function.spec.ts b/src/trace/context/extractors/step-function.spec.ts new file mode 100644 index 00000000..68b6e1c9 --- /dev/null +++ b/src/trace/context/extractors/step-function.spec.ts @@ -0,0 +1,142 @@ +import { readStepFunctionContextFromEvent } from "../../step-function-service"; + +describe("readStepFunctionContextFromEvent", () => { + const stepFunctionEvent = { + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + Input: { + MyInput: "MyValue", + }, + Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", + RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + StartTime: "2022-12-08T21:08:17.924Z", + }, + State: { + Name: "step-one", + EnteredTime: "2022-12-08T21:08:19.224Z", + RetryCount: 2, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + Name: "my-state-machine", + }, + } as const; + + it("reads a step function context from event with Execution.Input", () => { + const result = readStepFunctionContextFromEvent(stepFunctionEvent); + expect(result).toEqual({ + "step_function.execution_id": + "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + "step_function.execution_input": { MyInput: "MyValue" }, + "step_function.execution_name": "85a9933e-9e11-83dc-6a61-b92367b6c3be", + "step_function.execution_role_arn": + "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + "step_function.execution_start_time": "2022-12-08T21:08:17.924Z", + "step_function.state_entered_time": "2022-12-08T21:08:19.224Z", + "step_function.state_machine_arn": "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + "step_function.state_machine_name": "my-state-machine", + "step_function.state_name": "step-one", + "step_function.state_retry_count": 2, + }); + }); + it("returns undefined when event isn't an object", () => { + const result = readStepFunctionContextFromEvent("event"); + expect(result).toBeUndefined(); + }); + it("returns undefined when event is missing datadogContext property", () => { + const result = readStepFunctionContextFromEvent({}); + expect(result).toBeUndefined(); + }); + it("returns undefined when datadogContext is missing Execution property", () => { + const result = readStepFunctionContextFromEvent({ + dd: {}, + }); + expect(result).toBeUndefined(); + }); + it("returns undefined when Execution is missing Name field", () => { + const result = readStepFunctionContextFromEvent({ + dd: { + ...stepFunctionEvent, + Execution: {}, + }, + }); + expect(result).toBeUndefined(); + }); + it("returns undefined when Name isn't a string", () => { + const result = readStepFunctionContextFromEvent({ + dd: { + ...stepFunctionEvent, + Execution: { + Name: 12345, + }, + }, + }); + expect(result).toBeUndefined(); + }); + it("returns undefined when State isn't defined", () => { + const result = readStepFunctionContextFromEvent({ + dd: { + ...stepFunctionEvent, + State: undefined, + }, + }); + expect(result).toBeUndefined(); + }); + it("returns undefined when try retry count isn't a number", () => { + const result = readStepFunctionContextFromEvent({ + dd: { + ...stepFunctionEvent, + State: { + ...stepFunctionEvent.State, + RetryCount: "1", + }, + }, + }); + expect(result).toBeUndefined(); + }); + it("returns undefined when try step name isn't a string", () => { + const result = readStepFunctionContextFromEvent({ + dd: { + ...stepFunctionEvent, + State: { + ...stepFunctionEvent.State, + Name: 1, + }, + }, + }); + expect(result).toBeUndefined(); + }); + it("returns undefined when StateMachine is undefined", () => { + const result = readStepFunctionContextFromEvent({ + dd: { + ...stepFunctionEvent, + StateMachine: undefined, + }, + }); + expect(result).toBeUndefined(); + }); + it("returns undefined when StateMachineId isn't a string", () => { + const result = readStepFunctionContextFromEvent({ + dd: { + ...stepFunctionEvent, + StateMachine: { + ...stepFunctionEvent.StateMachine, + Id: 1, + }, + }, + }); + expect(result).toBeUndefined(); + }); + it("returns undefined when StateMachineName isn't a string", () => { + const result = readStepFunctionContextFromEvent({ + dd: { + ...stepFunctionEvent, + StateMachine: { + ...stepFunctionEvent.StateMachine, + Name: 1, + }, + }, + }); + expect(result).toBeUndefined(); + }); +}); diff --git a/src/trace/context/extractors/step-function.ts b/src/trace/context/extractors/step-function.ts new file mode 100644 index 00000000..21fb685a --- /dev/null +++ b/src/trace/context/extractors/step-function.ts @@ -0,0 +1,20 @@ +import { StepFunctionContext, deterministicMd5HashToBigIntString } from "../../step-function-service"; +import { SampleMode, Source, TraceContext } from "../extractor"; + +export function readTraceFromStepFunctionsContext(stepFunctionContext: StepFunctionContext): TraceContext | undefined { + const traceID = deterministicMd5HashToBigIntString(stepFunctionContext["step_function.execution_id"]); + const parentID = deterministicMd5HashToBigIntString( + stepFunctionContext["step_function.execution_id"] + + "#" + + stepFunctionContext["step_function.state_name"] + + "#" + + stepFunctionContext["step_function.state_entered_time"], + ); + + return { + parentID, + traceID, + sampleMode: SampleMode.AUTO_KEEP.valueOf(), + source: Source.Event, + }; +} diff --git a/src/trace/listener.spec.ts b/src/trace/listener.spec.ts index 022ca09e..9dcd888e 100644 --- a/src/trace/listener.spec.ts +++ b/src/trace/listener.spec.ts @@ -1,17 +1,10 @@ import { TraceExtractor, TraceListener } from "./listener"; -import { - Source, - ddtraceVersion, - parentIDHeader, - traceIDHeader, - samplingPriorityHeader, - parentSpanFinishTimeHeader, -} from "./constants"; import { datadogLambdaVersion } from "../constants"; import { Context } from "aws-lambda"; import { TraceHeaders } from "./trace-context-service"; import { SpanWrapper } from "./span-wrapper"; -import { eventSubTypes } from "./trigger"; +import { Source, parentIDHeader, samplingPriorityHeader, traceIDHeader } from "./context/extractor"; +import { ddtraceVersion, parentSpanFinishTimeHeader } from "./constants"; let mockWrap: jest.Mock; let mockExtract: jest.Mock; diff --git a/src/trace/listener.ts b/src/trace/listener.ts index be64b91a..3f9e5fd7 100644 --- a/src/trace/listener.ts +++ b/src/trace/listener.ts @@ -1,11 +1,5 @@ import { Context } from "aws-lambda"; -import { - addLambdaFunctionTagsToXray, - TraceContext, - readStepFunctionContextFromEvent, - StepFunctionContext, -} from "./context"; import { patchHttp, unpatchHttp } from "./patch-http"; import { TraceContextService } from "./trace-context-service"; import { extractTriggerTags, extractHTTPStatusCodeTag } from "./trigger"; @@ -14,12 +8,16 @@ import { ColdStartTracerConfig, ColdStartTracer } from "./cold-start-tracer"; import { logDebug, tagObject } from "../utils"; import { didFunctionColdStart, isProactiveInitialization } from "../utils/cold-start"; import { datadogLambdaVersion } from "../constants"; -import { Source, ddtraceVersion, parentSpanFinishTimeHeader, authorizingRequestIdHeader } from "./constants"; +import { ddtraceVersion, parentSpanFinishTimeHeader } from "./constants"; import { patchConsole } from "./patch-console"; import { SpanContext, TraceOptions, TracerWrapper } from "./tracer-wrapper"; import { SpanInferrer } from "./span-inferrer"; import { SpanWrapper } from "./span-wrapper"; import { getTraceTree, clearTraceTree } from "../runtime/index"; +import { Source, TraceContext } from "./context/extractor"; +import { StepFunctionContext, readStepFunctionContextFromEvent } from "./step-function-service"; +import { authorizingRequestIdHeader } from "./context/extractors/http"; +import { addLambdaFunctionTagsToXray } from "./xray-service"; export type TraceExtractor = (event: any, context: Context) => Promise | TraceContext; export interface TraceConfig { diff --git a/src/trace/patch-console.spec.ts b/src/trace/patch-console.spec.ts index 1960922e..59a0a77a 100644 --- a/src/trace/patch-console.spec.ts +++ b/src/trace/patch-console.spec.ts @@ -1,6 +1,6 @@ -import { SampleMode, Source } from "./constants"; import { TraceContextService } from "./trace-context-service"; import { patchConsole, unpatchConsole } from "./patch-console"; +import { SampleMode, Source } from "./context/extractor"; describe("patchConsole", () => { let traceWrapper = { diff --git a/src/trace/patch-http.spec.ts b/src/trace/patch-http.spec.ts index e4ed86f8..b7a23646 100644 --- a/src/trace/patch-http.spec.ts +++ b/src/trace/patch-http.spec.ts @@ -4,10 +4,10 @@ import nock from "nock"; import { parse } from "url"; import { LogLevel, setLogLevel } from "../utils"; -import { parentIDHeader, SampleMode, samplingPriorityHeader, traceIDHeader, Source } from "./constants"; import { patchHttp, unpatchHttp } from "./patch-http"; import { TraceContextService } from "./trace-context-service"; import { URL } from "url"; +import { SampleMode, Source, parentIDHeader, samplingPriorityHeader, traceIDHeader } from "./context/extractor"; describe("patchHttp", () => { let traceWrapper = { diff --git a/src/trace/span-inferrer.ts b/src/trace/span-inferrer.ts index 0198f033..41f54185 100644 --- a/src/trace/span-inferrer.ts +++ b/src/trace/span-inferrer.ts @@ -13,8 +13,8 @@ import { eventSubTypes, eventTypes, parseEventSource, parseEventSourceSubType } import { SpanWrapper } from "./span-wrapper"; import { DD_SERVICE_ENV_VAR, parentSpanFinishTimeHeader } from "./constants"; import { logDebug } from "../utils"; -import { getInjectedAuthorizerData } from "./context"; import { decodeAuthorizerContextEnvVar } from "../index"; +import { getInjectedAuthorizerData } from "./context/extractors/http"; export class SpanInferrer { private static serviceMapping: Record = {}; diff --git a/src/trace/step-function-service.spec.ts b/src/trace/step-function-service.spec.ts new file mode 100644 index 00000000..0991e973 --- /dev/null +++ b/src/trace/step-function-service.spec.ts @@ -0,0 +1,81 @@ +import { deterministicMd5HashInBinary, deterministicMd5HashToBigIntString, hexToBinary } from "./step-function-service"; + +describe("test_deterministicMd5HashToBigIntString", () => { + it("test same hashing number is generated as logs-backend for a random string", () => { + const actual = deterministicMd5HashToBigIntString("some_testing_random_string"); + expect(actual).toEqual("2251275791555400689"); + }); + + it("test same hashing number is generated as logs-backend for execution id # state name # entered time", () => { + const actual = deterministicMd5HashToBigIntString( + "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111#step-one#2022-12-08T21:08:19.224Z", + ); + expect(actual).toEqual("8034507082463708833"); + }); +}); + +describe("test_deterministicMd5HashInBinary", () => { + it("test same hashing is generated as logs-backend for a random string", () => { + const actual = deterministicMd5HashInBinary("some_testing_random_string"); + expect(actual).toEqual("0001111100111110001000110110011110010111000110001001001111110001"); + }); + + it("test same hashing is generated as logs-backend for an execution id", () => { + const actual = deterministicMd5HashInBinary( + "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d041f4", + ); + expect(actual).toEqual("0010010000101100100000101011111101111100110110001110111100111101"); + }); + + it("test same hashing is generated as logs-backend for another execution id", () => { + const actual = deterministicMd5HashInBinary( + "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111", + ); + expect(actual).toEqual("0010001100110000011011011111010000100111100000110000100100101010"); + }); + + it("test same hashing is generated as logs-backend for execution id # state name # entered time", () => { + const actual = deterministicMd5HashInBinary( + "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111#step-one#2022-12-08T21:08:19.224Z", + ); + expect(actual).toEqual("0110111110000000010011011001111101110011100111000000011010100001"); + }); + + it("test hashing different strings would generate different hashes", () => { + const times = 20; + for (let i = 0; i < times; i++) { + for (let j = i + 1; j < times; j++) { + expect(deterministicMd5HashInBinary(i.toString())).not.toMatch(deterministicMd5HashInBinary(j.toString())); + } + } + }); + + it("test always leading with 0", () => { + for (let i = 0; i < 20; i++) { + expect(deterministicMd5HashInBinary(i.toString()).substring(0, 1)).toMatch("0"); + } + }); +}); + +describe.each([ + ["0", "0000"], + ["1", "0001"], + ["2", "0010"], + ["3", "0011"], + ["4", "0100"], + ["5", "0101"], + ["6", "0110"], + ["7", "0111"], + ["8", "1000"], + ["9", "1001"], + ["a", "1010"], + ["b", "1011"], + ["c", "1100"], + ["d", "1101"], + ["e", "1110"], + ["f", "1111"], +])(`test hexToBinary`, (hex, expected) => { + test(`${hex} to binary returns ${expected}`, () => { + expect(hexToBinary(hex)).toBe(expected); + }); +}); diff --git a/src/trace/step-function-service.ts b/src/trace/step-function-service.ts new file mode 100644 index 00000000..e21cdfc3 --- /dev/null +++ b/src/trace/step-function-service.ts @@ -0,0 +1,125 @@ +import { Md5 } from "ts-md5"; +import { logDebug } from "../utils"; + +export interface StepFunctionContext { + "step_function.execution_name": string; + "step_function.execution_id": string; + "step_function.execution_input": object; + "step_function.execution_role_arn": string; + "step_function.execution_start_time": string; + "step_function.state_machine_name": string; + "step_function.state_machine_arn": string; + "step_function.state_entered_time": string; + "step_function.state_name": string; + "step_function.state_retry_count": number; +} + +export function hexToBinary(hex: string) { + // convert hex to binary and padding with 0 in the front to fill 128 bits + return parseInt(hex, 16).toString(2).padStart(4, "0"); +} + +export function deterministicMd5HashInBinary(s: string): string { + // Md5 here is used here because we don't need a cryptographically secure hashing method but to generate the same trace/span ids as the backend does + const hex = Md5.hashStr(s); + + let binary = ""; + for (let i = 0; i < hex.length; i++) { + const ch = hex.charAt(i); + binary = binary + hexToBinary(ch); + } + + const res = "0" + binary.substring(1, 64); + if (res === "0".repeat(64)) { + return "1"; + } + return res; +} + +export function deterministicMd5HashToBigIntString(s: string): string { + const binaryString = deterministicMd5HashInBinary(s); + return BigInt("0b" + binaryString).toString(); +} + +export function readStepFunctionContextFromEvent(event: any): StepFunctionContext | undefined { + if (typeof event !== "object") { + return; + } + + const execution = event.Execution; + if (typeof execution !== "object") { + logDebug("event.Execution is not an object."); + return; + } + const executionID = execution.Id; + if (typeof executionID !== "string") { + logDebug("event.Execution.Id is not a string."); + return; + } + const executionInput = execution.Input; + const executionName = execution.Name; + if (typeof executionName !== "string") { + logDebug("event.Execution.Name is not a string."); + return; + } + const executionRoleArn = execution.RoleArn; + if (typeof executionRoleArn !== "string") { + logDebug("event.Execution.RoleArn is not a string."); + return; + } + const executionStartTime = execution.StartTime; + if (typeof executionStartTime !== "string") { + logDebug("event.Execution.StartTime is not a string."); + return; + } + + const state = event.State; + if (typeof state !== "object") { + logDebug("event.State is not an object."); + return; + } + const stateRetryCount = state.RetryCount; + if (typeof stateRetryCount !== "number") { + logDebug("event.State.RetryCount is not a string."); + return; + } + const stateEnteredTime = state.EnteredTime; + if (typeof stateEnteredTime !== "string") { + logDebug("event.State.EnteredTime is not a string."); + return; + } + const stateName = state.Name; + if (typeof stateName !== "string") { + logDebug("event.State.Name is not a string."); + return; + } + + const stateMachine = event.StateMachine; + if (typeof stateMachine !== "object") { + logDebug("event.StateMachine is not an object."); + return; + } + const stateMachineArn = stateMachine.Id; + if (typeof stateMachineArn !== "string") { + logDebug("event.StateMachine.Id is not a string."); + return; + } + const stateMachineName = stateMachine.Name; + if (typeof stateMachineName !== "string") { + logDebug("event.StateMachine.Name is not a string."); + return; + } + + return { + "step_function.execution_name": executionName, + "step_function.execution_id": executionID, + "step_function.execution_input": executionInput ?? {}, + "step_function.execution_role_arn": executionRoleArn, + "step_function.execution_start_time": executionStartTime, + "step_function.state_entered_time": stateEnteredTime, + "step_function.state_machine_arn": stateMachineArn, + "step_function.state_machine_name": stateMachineName, + "step_function.state_name": stateName, + "step_function.state_retry_count": stateRetryCount, + }; +} diff --git a/src/trace/trace-context-service.spec.ts b/src/trace/trace-context-service.spec.ts index 64f03992..9e2e871d 100644 --- a/src/trace/trace-context-service.spec.ts +++ b/src/trace/trace-context-service.spec.ts @@ -1,6 +1,5 @@ +import { SampleMode, Source, TraceContext } from "./context/extractor"; import { TraceContextService } from "./trace-context-service"; -import { TraceContext } from "./context"; -import { SampleMode, Source } from "./constants"; let mockXRaySegment: any; let mockXRayShouldThrow = false; diff --git a/src/trace/trace-context-service.ts b/src/trace/trace-context-service.ts index e19db7e5..ba9560c6 100644 --- a/src/trace/trace-context-service.ts +++ b/src/trace/trace-context-service.ts @@ -1,9 +1,14 @@ import { Context } from "aws-lambda"; import { logDebug } from "../utils"; -import { parentIDHeader, samplingPriorityHeader, traceIDHeader } from "./constants"; -import { TraceContext, extractTraceContext } from "./context"; import { TracerWrapper } from "./tracer-wrapper"; import { TraceExtractor } from "./listener"; +import { + TraceContext, + extractTraceContext, + parentIDHeader, + samplingPriorityHeader, + traceIDHeader, +} from "./context/extractor"; /** * Headers that can be added to a request. */ diff --git a/src/trace/tracer-wrapper.spec.ts b/src/trace/tracer-wrapper.spec.ts index 6f9eb00b..07dedb02 100644 --- a/src/trace/tracer-wrapper.spec.ts +++ b/src/trace/tracer-wrapper.spec.ts @@ -1,5 +1,5 @@ +import { SampleMode, Source } from "./context/extractor"; import { TracerWrapper } from "./tracer-wrapper"; -import { Source, SampleMode } from "./constants"; let mockNoTracer = false; let mockTracerInitialised = false; diff --git a/src/trace/tracer-wrapper.ts b/src/trace/tracer-wrapper.ts index 896b46d7..53325b87 100644 --- a/src/trace/tracer-wrapper.ts +++ b/src/trace/tracer-wrapper.ts @@ -1,6 +1,5 @@ import { logDebug } from "../utils"; -import { SampleMode, Source } from "./constants"; -import { TraceContext } from "./context"; +import { SampleMode, Source, TraceContext } from "./context/extractor"; import { TraceHeaders } from "./trace-context-service"; export interface SpanContext { diff --git a/src/trace/xray-service.spec.ts b/src/trace/xray-service.spec.ts new file mode 100644 index 00000000..cfb4ad62 --- /dev/null +++ b/src/trace/xray-service.spec.ts @@ -0,0 +1,97 @@ +import { SampleMode } from "./context/extractor"; +import { + convertToAPMParentID, + convertToAPMTraceID, + convertToSampleMode, + readTraceContextFromXray, +} from "./xray-service"; + +describe("readTraceContextFromXray", () => { + afterEach(() => { + process.env["_X_AMZN_TRACE_ID"] = undefined; + }); + it("returns a trace context from a valid env var", () => { + process.env["_X_AMZN_TRACE_ID"] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"; + const context = readTraceContextFromXray(); + expect(context).toEqual({ + parentID: "10713633173203262661", + sampleMode: 2, + source: "xray", + traceID: "3995693151288333088", + }); + }); + it("returns undefined when given an invalid env var", () => { + const badCases = [ + "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5", + "Root=1-5e272390-8c398be037738dc042009320", + "1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1", + "Root=1-5e272390-8c398be037738dc042009320;94ae789b969f1cc5;Sampled=1", + "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;1", + "Root=a;Parent=94ae789b969f1cc5;Sampled=1", + "Root=1-5e272390-8c398be037738dc042009320;Parent=b;Sampled=1", + undefined, + ]; + for (const badCase of badCases) { + process.env["_X_AMZN_TRACE_ID"] = badCase; + expect(readTraceContextFromXray()).toBeUndefined(); + } + }); +}); + +describe("convertToAPMTraceID", () => { + it("converts an xray trace id to a Datadog trace ID", () => { + const xrayTraceID = "1-5ce31dc2-2c779014b90ce44db5e03875"; + const traceID = convertToAPMTraceID(xrayTraceID); + expect(traceID).toEqual("4110911582297405557"); + }); + it("converts an xray trace id to a Datadog trace ID removing first bit", () => { + const xrayTraceID = "1-5ce31dc2-ac779014b90ce44db5e03875"; // Number with 64bit toggled on + const traceID = convertToAPMTraceID(xrayTraceID); + expect(traceID).toEqual("4110911582297405557"); + }); + it("returns undefined when xray trace id is too short", () => { + const xrayTraceID = "1-5ce31dc2-5e03875"; + const traceID = convertToAPMTraceID(xrayTraceID); + expect(traceID).toBeUndefined(); + }); + + it("returns undefined when xray trace id is in an invalid format", () => { + const xrayTraceID = "1-2c779014b90ce44db5e03875"; + const traceID = convertToAPMTraceID(xrayTraceID); + expect(traceID).toBeUndefined(); + }); + it("returns undefined when xray trace id uses invalid characters", () => { + const xrayTraceID = "1-5ce31dc2-c779014b90ce44db5e03875;"; + const traceID = convertToAPMTraceID(xrayTraceID); + expect(traceID).toBeUndefined(); + }); +}); + +describe("convertToAPMParentID", () => { + it("converts an xray parent ID to an APM parent ID", () => { + const xrayParentID = "0b11cc4230d3e09e"; + const parentID = convertToAPMParentID(xrayParentID); + expect(parentID).toEqual("797643193680388254"); + }); + it("returns undefined when parent ID uses invalid characters", () => { + const xrayParentID = ";79014b90ce44db5e0;875"; + const parentID = convertToAPMParentID(xrayParentID); + expect(parentID).toBeUndefined(); + }); + it("returns undefined when parent ID is wrong size", () => { + const xrayParentID = "5e03875"; + const parentID = convertToAPMParentID(xrayParentID); + expect(parentID).toBeUndefined(); + }); +}); + +describe("convertToSampleMode", () => { + it("returns USER_KEEP if xray was sampled", () => { + const result = convertToSampleMode(1); + expect(result).toBe(SampleMode.USER_KEEP); + }); + it("returns USER_REJECT if xray wasn't sampled", () => { + const result = convertToSampleMode(0); + expect(result).toBe(SampleMode.USER_REJECT); + }); +}); diff --git a/src/trace/xray-service.ts b/src/trace/xray-service.ts new file mode 100644 index 00000000..ec6f2924 --- /dev/null +++ b/src/trace/xray-service.ts @@ -0,0 +1,204 @@ +import { randomBytes } from "crypto"; +import { logDebug, logError } from "../utils"; +import { SampleMode, Source, TraceContext } from "./context/extractor"; +import { StepFunctionContext } from "./step-function-service"; +import { Socket, createSocket } from "dgram"; +import BigNumber from "bignumber.js"; + +export interface XRayTraceHeader { + traceID: string; + parentID: string; + sampled: number; +} + +export const xraySubsegmentName = "datadog-metadata"; +export const xraySubsegmentKey = "trace"; +export const xrayBaggageSubsegmentKey = "root_span_metadata"; +export const xrayLambdaFunctionTagsKey = "lambda_function_tags"; +export const xraySubsegmentNamespace = "datadog"; +export const xrayTraceEnvVar = "_X_AMZN_TRACE_ID"; +export const awsXrayDaemonAddressEnvVar = "AWS_XRAY_DAEMON_ADDRESS"; + +export function addTraceContextToXray(traceContext: TraceContext) { + const val = { + "parent-id": traceContext.parentID, + "sampling-priority": traceContext.sampleMode.toString(10), + "trace-id": traceContext.traceID, + }; + + addXrayMetadata(xraySubsegmentKey, val); +} + +export function addStepFunctionContextToXray(context: StepFunctionContext) { + addXrayMetadata(xrayBaggageSubsegmentKey, context); +} + +export function addLambdaFunctionTagsToXray(triggerTags: { [key: string]: string }) { + addXrayMetadata(xrayLambdaFunctionTagsKey, triggerTags); +} + +export function addXrayMetadata(key: string, metadata: Record) { + const segment = generateXraySubsegment(key, metadata); + if (segment === undefined) { + return; + } + sendXraySubsegment(segment); +} + +export function generateXraySubsegment(key: string, metadata: Record) { + const header = process.env[xrayTraceEnvVar]; + if (header === undefined) { + logDebug("couldn't read xray trace header from env"); + return; + } + const context = parseXrayTraceContextHeader(header); + if (context === undefined) { + logDebug("couldn't parse xray trace header from env"); + return; + } + const sampled = convertToSampleMode(parseInt(context.xraySampled, 10)); + if (sampled === SampleMode.USER_REJECT || sampled === SampleMode.AUTO_REJECT) { + logDebug("discarding xray metadata subsegment due to sampling"); + return; + } + + // Convert from milliseconds to seconds + const time = Date.now() * 0.001; + + return JSON.stringify({ + id: randomBytes(8).toString("hex"), + trace_id: context.xrayTraceID, + parent_id: context.xrayParentID, + name: xraySubsegmentName, + start_time: time, + end_time: time, + type: "subsegment", + metadata: { + [xraySubsegmentNamespace]: { + [key]: metadata, + }, + }, + }); +} + +export function sendXraySubsegment(segment: string) { + const xrayDaemonEnv = process.env[awsXrayDaemonAddressEnvVar]; + if (xrayDaemonEnv === undefined) { + logDebug("X-Ray daemon env var not set, not sending sub-segment"); + return; + } + const parts = xrayDaemonEnv.split(":"); + if (parts.length <= 1) { + logDebug("X-Ray daemon env var has invalid format, not sending sub-segment"); + return; + } + const port = parseInt(parts[1], 10); + const address = parts[0]; + + const message = Buffer.from(`{\"format\": \"json\", \"version\": 1}\n${segment}`); + let client: Socket | undefined; + try { + client = createSocket("udp4"); + // Send segment asynchronously to xray daemon + client.send(message, 0, message.length, port, address, (error, bytes) => { + client?.close(); + logDebug(`Xray daemon received metadata payload`, { error, bytes }); + }); + } catch (error) { + if (error instanceof Error) { + client?.close(); + logDebug("Error occurred submitting to xray daemon", error); + } + } +} + +export function readTraceContextFromXray(): TraceContext | undefined { + const header = process.env[xrayTraceEnvVar]; + if (header === undefined) { + logDebug("couldn't read xray trace header from env"); + return; + } + const context = parseXrayTraceContextHeader(header); + + if (context === undefined) { + logError("couldn't read xray trace context from env, variable had invalid format"); + return undefined; + } + const parentID = convertToAPMParentID(context.xrayParentID); + if (parentID === undefined) { + logDebug("couldn't parse xray parent ID", context); + return; + } + const traceID = convertToAPMTraceID(context.xrayTraceID); + if (traceID === undefined) { + logDebug("couldn't parse xray trace ID", context); + return; + } + const sampleMode = convertToSampleMode(parseInt(context.xraySampled, 10)); + + const trace = { + parentID, + sampleMode, + source: Source.Xray, + traceID, + }; + logDebug(`extracted trace context from xray context`, { trace, header }); + return trace; +} + +function parseXrayTraceContextHeader(header: string) { + // Example: Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1 + logDebug(`Reading trace context from env var ${header}`); + const [root, parent, sampled] = header.split(";"); + if (parent === undefined || sampled === undefined) { + return; + } + const [, xrayTraceID] = root.split("="); + const [, xrayParentID] = parent.split("="); + const [, xraySampled] = sampled.split("="); + if (xraySampled === undefined || xrayParentID === undefined || xrayTraceID === undefined) { + return; + } + return { + xrayTraceID, + xraySampled, + xrayParentID, + }; +} + +export function convertToSampleMode(xraySampled: number): SampleMode { + return xraySampled === 1 ? SampleMode.USER_KEEP : SampleMode.USER_REJECT; +} + +export function convertToAPMTraceID(xrayTraceID: string): string | undefined { + const parts = xrayTraceID.split("-"); + if (parts.length < 3) { + return; + } + const lastPart = parts[2]; + if (lastPart.length !== 24) { + return; + } + + // We want to turn the last 63 bits into a decimal number in a string representation + // Unfortunately, all numbers in javascript are represented by float64 bit numbers, which + // means we can't parse 64 bit integers accurately. + const hex = new BigNumber(lastPart, 16); + if (hex.isNaN()) { + return; + } + // Toggle off the 64th bit + const last63Bits = hex.mod(new BigNumber("8000000000000000", 16)); + return last63Bits.toString(10); +} + +export function convertToAPMParentID(xrayParentID: string): string | undefined { + if (xrayParentID.length !== 16) { + return; + } + const hex = new BigNumber(xrayParentID, 16); + if (hex.isNaN()) { + return; + } + return hex.toString(10); +}