Skip to content

Commit 7753de1

Browse files
kimi-pkimi-p
authored
breaking: Use sha256 to hash lambda traceId that are triggered by Step Functions and set _dd.p.tid (#534)
1 parent a22b5e9 commit 7753de1

File tree

6 files changed

+847
-818
lines changed

6 files changed

+847
-818
lines changed

‎package.json

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
"@types/aws-sdk": "^2.7.0",
2626
"@types/jest": "^26.0.23",
2727
"@types/mock-fs": "4.13.0",
28-
"@types/node": "^15.6.1",
28+
"@types/node": "^20.12.10",
2929
"@types/promise-retry": "^1.1.3",
3030
"@types/shimmer": "^1.0.1",
3131
"dd-trace": "^4.37.0",
@@ -38,12 +38,12 @@
3838
"typescript": "^4.3.2"
3939
},
4040
"dependencies": {
41+
"@aws-crypto/sha256-js": "5.2.0",
4142
"dc-polyfill": "^0.1.3",
4243
"hot-shots": "8.5.0",
4344
"promise-retry": "^2.0.1",
4445
"serialize-error": "^8.1.0",
45-
"shimmer": "1.2.1",
46-
"ts-md5": "1.3.1"
46+
"shimmer": "1.2.1"
4747
},
4848
"jest": {
4949
"verbose": true,

‎src/trace/context/extractor.spec.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -688,8 +688,8 @@ describe("TraceContextExtractor", () => {
688688
const traceContext = await extractor.extract(event, {} as Context);
689689
expect(traceContext).not.toBeNull();
690690

691-
expect(traceContext?.toTraceId()).toBe("947965466153612645");
692-
expect(traceContext?.toSpanId()).toBe("4602916161841036335");
691+
expect(traceContext?.toTraceId()).toBe("1139193989631387307");
692+
expect(traceContext?.toSpanId()).toBe("5892738536804826142");
693693
expect(traceContext?.sampleMode()).toBe("1");
694694
expect(traceContext?.source).toBe("event");
695695
});

‎src/trace/context/extractors/step-function.spec.ts

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,8 @@ describe("StepFunctionEventTraceExtractor", () => {
3636
const traceContext = extractor.extract(payload);
3737
expect(traceContext).not.toBeNull();
3838

39-
expect(traceContext?.toTraceId()).toBe("947965466153612645");
40-
expect(traceContext?.toSpanId()).toBe("4602916161841036335");
39+
expect(traceContext?.toTraceId()).toBe("1139193989631387307");
40+
expect(traceContext?.toSpanId()).toBe("5892738536804826142");
4141
expect(traceContext?.sampleMode()).toBe("1");
4242
expect(traceContext?.source).toBe("event");
4343
});
@@ -49,8 +49,8 @@ describe("StepFunctionEventTraceExtractor", () => {
4949
const traceContext = extractor.extract(payload);
5050
expect(traceContext).not.toBeNull();
5151

52-
expect(traceContext?.toTraceId()).toBe("947965466153612645");
53-
expect(traceContext?.toSpanId()).toBe("4602916161841036335");
52+
expect(traceContext?.toTraceId()).toBe("1139193989631387307");
53+
expect(traceContext?.toSpanId()).toBe("5892738536804826142");
5454
expect(traceContext?.sampleMode()).toBe("1");
5555
expect(traceContext?.source).toBe("event");
5656
});

‎src/trace/step-function-service.spec.ts

Lines changed: 44 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { StepFunctionContextService } from "./step-function-service";
1+
import { PARENT_ID, StepFunctionContextService } from "./step-function-service";
22

33
describe("StepFunctionContextService", () => {
44
const stepFunctionEvent = {
@@ -194,8 +194,8 @@ describe("StepFunctionContextService", () => {
194194

195195
expect(spanContext).not.toBeNull();
196196

197-
expect(spanContext?.toTraceId()).toBe("947965466153612645");
198-
expect(spanContext?.toSpanId()).toBe("4602916161841036335");
197+
expect(spanContext?.toTraceId()).toBe("1139193989631387307");
198+
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
199199
expect(spanContext?.sampleMode()).toBe("1");
200200
expect(spanContext?.source).toBe("event");
201201
});
@@ -211,54 +211,55 @@ describe("StepFunctionContextService", () => {
211211
});
212212
});
213213

214-
describe("deterministicMd5HashToBigIntString", () => {
214+
describe("deterministicSha256HashToBigIntString", () => {
215215
it("returns the same hash number generated in `logs backend` for a random string", () => {
216216
const instance = StepFunctionContextService.instance();
217-
const hash = instance["deterministicMd5HashToBigIntString"]("some_testing_random_string");
218-
expect(hash).toEqual("2251275791555400689");
217+
const hash = instance["deterministicSha256HashToBigIntString"]("some_testing_random_string", PARENT_ID);
218+
expect(hash).toEqual("4364271812988819936");
219219
});
220220

221221
it("returns the same hash number generated in `logs backend` for execution id # state name # entered time", () => {
222222
const instance = StepFunctionContextService.instance();
223-
const hash = instance["deterministicMd5HashToBigIntString"](
223+
const hash = instance["deterministicSha256HashToBigIntString"](
224224
"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",
225+
PARENT_ID,
225226
);
226-
expect(hash).toEqual("8034507082463708833");
227+
expect(hash).toEqual("4340734536022949921");
227228
});
228229
});
229230

230-
describe("deterministicMd5HashInBinary", () => {
231+
describe("deterministicSha256Hash", () => {
231232
it.each([
232233
[
233234
"a random string",
234235
"some_testing_random_string",
235-
"0001111100111110001000110110011110010111000110001001001111110001",
236+
"0011110010010001000000100001011101001100011100101101100111100000",
236237
],
237238
[
238239
"an execution id",
239240
"arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d041f4",
240-
"0010010000101100100000101011111101111100110110001110111100111101",
241+
"0100010100110010010010100001011001110100111011010100110010000100",
241242
],
242243
[
243244
"another execution id",
244245
"arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111",
245-
"0010001100110000011011011111010000100111100000110000100100101010",
246+
"0010111110001100100010000101001100110000000000010111011100101011",
246247
],
247248
[
248249
"execution id # state name # entered time",
249250
"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",
250-
"0110111110000000010011011001111101110011100111000000011010100001",
251+
"0011110000111101011000110000111111110011111010110000000000100001",
251252
],
252253
])("returns the same hash number generated in `logs backend` for %s", (_, str, expected) => {
253254
const instance = StepFunctionContextService.instance();
254-
const hash = instance["deterministicMd5HashInBinary"](str);
255+
const hash = instance["deterministicSha256Hash"](str, PARENT_ID);
255256
expect(hash).toEqual(expected);
256257
});
257258

258259
it("returns a hash always leading with 0", () => {
259260
const instance = StepFunctionContextService.instance();
260261
for (let i = 0; i < 20; i++) {
261-
const hash = instance["deterministicMd5HashInBinary"](i.toString());
262+
const hash = instance["deterministicSha256Hash"](i.toString(), PARENT_ID);
262263
expect(hash.substring(0, 1)).toMatch("0");
263264
}
264265
});
@@ -268,36 +269,44 @@ describe("StepFunctionContextService", () => {
268269
const times = 20;
269270
for (let i = 0; i < times; i++) {
270271
for (let j = i + 1; j < times; j++) {
271-
const hash1 = instance["deterministicMd5HashInBinary"](i.toString());
272-
const hash2 = instance["deterministicMd5HashInBinary"](j.toString());
272+
const hash1 = instance["deterministicSha256Hash"](i.toString(), PARENT_ID);
273+
const hash2 = instance["deterministicSha256Hash"](j.toString(), PARENT_ID);
273274
expect(hash1).not.toMatch(hash2);
274275
}
275276
}
276277
});
277278
});
278279

279-
describe("hexToBinary", () => {
280+
describe("numberToBinaryString", () => {
280281
const instance = StepFunctionContextService.instance();
281282
it.each([
282-
["0", "0000"],
283-
["1", "0001"],
284-
["2", "0010"],
285-
["3", "0011"],
286-
["4", "0100"],
287-
["5", "0101"],
288-
["6", "0110"],
289-
["7", "0111"],
290-
["8", "1000"],
291-
["9", "1001"],
292-
["a", "1010"],
293-
["b", "1011"],
294-
["c", "1100"],
295-
["d", "1101"],
296-
["e", "1110"],
297-
["f", "1111"],
283+
[0, "00000000"],
284+
[1, "00000001"],
285+
[2, "00000010"],
286+
[3, "00000011"],
287+
[4, "00000100"],
298288
])("returns the right binary number for %s => %s", (hex, expected) => {
299-
const binary = instance["hexToBinary"](hex);
289+
const binary = instance["numberToBinaryString"](hex);
300290
expect(binary).toBe(expected);
301291
});
302292
});
293+
294+
describe("test 64 bits deterministicSha256HashToBigIntString for span id", () => {
295+
const instance = StepFunctionContextService.instance();
296+
it("first test of #1", () => {
297+
const actual = instance["deterministicSha256HashToBigIntString"](
298+
"arn:aws:states:sa-east-1:425362996713:stateMachine:MyStateMachine-b276uka1j#lambda#1",
299+
PARENT_ID,
300+
);
301+
expect(actual).toEqual("3711631873188331089");
302+
});
303+
304+
it("test same hashing number is generated as logs-backend for execution id # state name # entered time", () => {
305+
const actual = instance["deterministicSha256HashToBigIntString"](
306+
"arn:aws:states:sa-east-1:425362996713:stateMachine:MyStateMachine-b276uka1j#lambda#2",
307+
PARENT_ID,
308+
);
309+
expect(actual).toEqual("5759173372325510050");
310+
});
311+
});
303312
});

‎src/trace/step-function-service.ts

Lines changed: 47 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
import { Md5 } from "ts-md5";
21
import { logDebug } from "../utils";
32
import { SampleMode, TraceSource } from "./trace-context-service";
43
import { SpanContextWrapper } from "./span-context-wrapper";
4+
import { Sha256 } from "@aws-crypto/sha256-js";
55

66
export interface StepFunctionContext {
77
"step_function.execution_name": string;
@@ -16,6 +16,10 @@ export interface StepFunctionContext {
1616
"step_function.state_retry_count": number;
1717
}
1818

19+
export const TRACE_ID = "traceId";
20+
export const PARENT_ID = "spanId";
21+
export const DD_P_TID = "_dd.p.tid";
22+
1923
export class StepFunctionContextService {
2024
private static _instance: StepFunctionContextService;
2125
public context?: StepFunctionContext;
@@ -123,51 +127,69 @@ export class StepFunctionContextService {
123127
public get spanContext(): SpanContextWrapper | null {
124128
if (this.context === undefined) return null;
125129

126-
const traceId = this.deterministicMd5HashToBigIntString(this.context["step_function.execution_id"]);
127-
const parentId = this.deterministicMd5HashToBigIntString(
130+
const traceId = this.deterministicSha256HashToBigIntString(this.context["step_function.execution_id"], TRACE_ID);
131+
const parentId = this.deterministicSha256HashToBigIntString(
128132
this.context["step_function.execution_id"] +
129133
"#" +
130134
this.context["step_function.state_name"] +
131135
"#" +
132136
this.context["step_function.state_entered_time"],
137+
PARENT_ID,
133138
);
134139
const sampleMode = SampleMode.AUTO_KEEP;
135140

136-
const spanContext = SpanContextWrapper.fromTraceContext({
137-
traceId,
138-
parentId,
139-
sampleMode,
140-
source: TraceSource.Event,
141-
});
142-
143-
if (spanContext === null) return null;
144-
logDebug(`Extracted trace context from StepFunctionContext`, { traceContext: this.context });
145-
return spanContext;
141+
try {
142+
// Try requiring class from the tracer.
143+
const _DatadogSpanContext = require("dd-trace/packages/dd-trace/src/opentracing/span_context");
144+
const id = require("dd-trace/packages/dd-trace/src/id");
145+
146+
const ddSpanContext = new _DatadogSpanContext({
147+
traceId: id(traceId, 10),
148+
spanId: id(parentId, 10),
149+
sampling: { priority: sampleMode.toString(2) },
150+
});
151+
152+
const ptid = this.deterministicSha256HashToBigIntString(this.context["step_function.execution_id"], DD_P_TID);
153+
ddSpanContext._trace.tags["_dd.p.tid"] = id(ptid, 10).toString(16);
154+
if (ddSpanContext === null) return null;
155+
156+
logDebug(`Extracted trace context from StepFunctionContext`, { traceContext: ddSpanContext });
157+
158+
return new SpanContextWrapper(ddSpanContext, TraceSource.Event);
159+
} catch (error) {
160+
if (error instanceof Error) {
161+
logDebug("Couldn't generate SpanContext with tracer.", error);
162+
}
163+
return null;
164+
}
146165
}
147166

148-
private deterministicMd5HashToBigIntString(s: string): string {
149-
const binaryString = this.deterministicMd5HashInBinary(s);
167+
private deterministicSha256HashToBigIntString(s: string, type: string): string {
168+
const binaryString = this.deterministicSha256Hash(s, type);
150169
return BigInt("0b" + binaryString).toString();
151170
}
152171

153-
private deterministicMd5HashInBinary(s: string): string {
154-
// 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
155-
const hex = Md5.hashStr(s);
172+
private deterministicSha256Hash(s: string, type: string): string {
173+
// returns 128 bits hash unless mostSignificant64Bits options is set to true.
156174

157-
let binary = "";
158-
for (let i = 0; i < hex.length; i++) {
159-
const ch = hex.charAt(i);
160-
binary = binary + this.hexToBinary(ch);
175+
const hash = new Sha256();
176+
hash.update(s);
177+
const uint8Array = hash.digestSync();
178+
// type === SPAN_ID || type === DD_P_TID
179+
let intArray = uint8Array.subarray(0, 8);
180+
if (type === TRACE_ID) {
181+
intArray = uint8Array.subarray(8, 16);
161182
}
183+
const binaryString = intArray.reduce((acc, num) => acc + this.numberToBinaryString(num), "");
162184

163-
const res = "0" + binary.substring(1, 64);
185+
const res = "0" + binaryString.substring(1, 64);
164186
if (res === "0".repeat(64)) {
165187
return "1";
166188
}
167189
return res;
168190
}
169191

170-
private hexToBinary(hex: string) {
171-
return parseInt(hex, 16).toString(2).padStart(4, "0");
192+
private numberToBinaryString(num: number): string {
193+
return num.toString(2).padStart(8, "0");
172194
}
173195
}

0 commit comments

Comments
 (0)