diff --git a/package.json b/package.json index 8c607366..5945096d 100644 --- a/package.json +++ b/package.json @@ -39,7 +39,8 @@ "hot-shots": "8.5.0", "promise-retry": "^2.0.1", "serialize-error": "^8.1.0", - "shimmer": "^1.2.1" + "shimmer": "^1.2.1", + "ts-md5": "1.3.1" }, "jest": { "verbose": true, diff --git a/src/trace/context.spec.ts b/src/trace/context.spec.ts index 296045ea..c09fe0c3 100644 --- a/src/trace/context.spec.ts +++ b/src/trace/context.spec.ts @@ -22,6 +22,11 @@ import { readTraceFromSQSEvent, readTraceFromHTTPEvent, readTraceFromLambdaContext, + hexToBinary, + deterministicMd5HashInBinary, + deterministicMd5HashToBigIntString, + readTraceFromStepFunctionsContext, + StepFunctionContext, } from "./context"; let sentSegment: any; @@ -1052,6 +1057,38 @@ describe("extractTraceContext", () => { expect(sentSegment).toBeUndefined(); }); + it("returns trace read from step functions event with the extractor as the highest priority", () => { + 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 = 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", () => { jest.spyOn(Date, "now").mockImplementation(() => 1487076708000); process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0"; @@ -1109,3 +1146,83 @@ describe("extractTraceContext", () => { `); }); }); + +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.ts b/src/trace/context.ts index 94cb11b8..1a259225 100644 --- a/src/trace/context.ts +++ b/src/trace/context.ts @@ -28,6 +28,7 @@ import { } from "./constants"; import { TraceExtractor } from "./listener"; import { eventSubTypes, parseEventSourceSubType } from "./trigger"; +import { Md5 } from "ts-md5"; export interface XRayTraceHeader { traceID: string; @@ -55,6 +56,48 @@ export interface StepFunctionContext { "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. @@ -95,6 +138,12 @@ export function extractTraceContext( 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) { diff --git a/tsconfig.json b/tsconfig.json index bcb941eb..256e1a0f 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -4,7 +4,11 @@ // "incremental": true, /* Enable incremental compilation */ "target": "es5" /* Specify ECMAScript target version: 'ES3' (default), 'ES5', 'ES2015', 'ES2016', 'ES2017', 'ES2018', 'ES2019' or 'ESNEXT'. */, "module": "commonjs" /* Specify module code generation: 'none', 'commonjs', 'amd', 'system', 'umd', 'es2015', or 'ESNext'. */, - "lib": ["es2015"] /* Specify library files to be included in the compilation. */, + "lib": [ + "es2015", + "es2020.bigint" + ] + /* Specify library files to be included in the compilation. */, // "allowJs": true, /* Allow javascript files to be compiled. */ // "checkJs": true, /* Report errors in .js files. */ // "jsx": "preserve", /* Specify JSX code generation: 'preserve', 'react-native', or 'react'. */ diff --git a/yarn.lock b/yarn.lock index 895cd035..5df985ab 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3378,6 +3378,11 @@ ts-jest@^27.0.1: semver "7.x" yargs-parser "20.x" +ts-md5@1.3.1: + version "1.3.1" + resolved "https://registry.yarnpkg.com/ts-md5/-/ts-md5-1.3.1.tgz#f5b860c0d5241dd9bb4e909dd73991166403f511" + integrity sha512-DiwiXfwvcTeZ5wCE0z+2A9EseZsztaiZtGrtSaY5JOD7ekPnR/GoIVD5gXZAlK9Na9Kvpo9Waz5rW64WKAWApg== + tslib@^1.13.0, tslib@^1.8.1: version "1.14.1" resolved "https://registry.yarnpkg.com/tslib/-/tslib-1.14.1.tgz#cf2d38bdc34a134bcaf1091c41f6619e2f672d00"