feat(aap): update tags on inferred spans for API Gateway

src/trace/listener.spec.ts

@@ -16,6 +16,7 @@ let mockExtract: jest.Mock<any, any>;
 let mockSpanContextWrapper: any;
 let mockSpanContext: any;
 let mockTraceSource: TraceSource | undefined = undefined;
+let mockCurrentSpan: any = null;
 
 jest.mock("./tracer-wrapper", () => {
   mockWrap = jest.fn().mockImplementation((name, options, func) => func);
@@ -25,6 +26,10 @@ jest.mock("./tracer-wrapper", () => {
 
     constructor() {}
 
+    get currentSpan(): any {
+      return mockCurrentSpan;
+    }
+
     wrap(name: any, options: any, fn: any): any {
       return mockWrap(name, options, fn);
     }
@@ -560,4 +565,97 @@ describe("TraceListener", () => {
       );
     });
   });
+
+  describe("AppSec tag propagation", () => {
+    beforeEach(() => {
+      mockWrap.mockClear();
+      mockExtract.mockClear();
+      mockSpanContext = undefined;
+      mockSpanContextWrapper = undefined;
+      mockTraceSource = undefined;
+      mockCurrentSpan = null;
+      process.env = { ...oldEnv };
+    });
+
+    afterEach(() => {
+      process.env = oldEnv;
+      mockCurrentSpan = null;
+    });
+
+    it("copies _dd.appsec.json from lambda span to inferred span when present", async () => {
+      const mockSetTag = jest.fn();
+      const appsecJsonValue = '{"triggers":[{"rule":{"id":"rule-1"}}]}';
+
+      mockCurrentSpan = {
+        _tags: {
+          "_dd.appsec.json": appsecJsonValue,
+        },
+      };
+
+      const listener = new TraceListener(defaultConfig);
+      await listener.onStartInvocation({}, context as any);
+
+      const mockInferredSpan = {
+        isAsync: () => false,
+        setTag: mockSetTag,
+        finish: jest.fn(),
+      };
+      (listener as any).inferredSpan = mockInferredSpan;
+
+      const unwrappedFunc = () => {};
+      const wrappedFunc = listener.onWrap(unwrappedFunc);
+      wrappedFunc();
+      await listener.onCompleteInvocation();
+
+      expect(mockSetTag).toHaveBeenCalledWith("_dd.appsec.json", appsecJsonValue);
+    });
+
+    it("does not set _dd.appsec.json on inferred span when not present on lambda span", async () => {
+      const mockSetTag = jest.fn();
+
+      mockCurrentSpan = {
+        _tags: {},
+      };
+
+      const listener = new TraceListener(defaultConfig);
+      await listener.onStartInvocation({}, context as any);
+
+      const mockInferredSpan = {
+        isAsync: () => false,
+        setTag: mockSetTag,
+        finish: jest.fn(),
+      };
+      (listener as any).inferredSpan = mockInferredSpan;
+
+      const unwrappedFunc = () => {};
+      const wrappedFunc = listener.onWrap(unwrappedFunc);
+      wrappedFunc();
+      await listener.onCompleteInvocation();
+
+      expect(mockSetTag).not.toHaveBeenCalledWith("_dd.appsec.json", expect.anything());
+    });
+
+    it("does not set _dd.appsec.json when lambda span is not available", async () => {
+      const mockSetTag = jest.fn();
+
+      mockCurrentSpan = null;
+
+      const listener = new TraceListener(defaultConfig);
+      await listener.onStartInvocation({}, context as any);
+
+      const mockInferredSpan = {
+        isAsync: () => false,
+        setTag: mockSetTag,
+        finish: jest.fn(),
+      };
+      (listener as any).inferredSpan = mockInferredSpan;
+
+      const unwrappedFunc = () => {};
+      const wrappedFunc = listener.onWrap(unwrappedFunc);
+      wrappedFunc();
+      await listener.onCompleteInvocation();
+
+      expect(mockSetTag).not.toHaveBeenCalledWith("_dd.appsec.json", expect.anything());
+    });
+  });
 });

src/trace/listener.ts

@@ -267,6 +267,15 @@ export class TraceListener {
         logDebug("Setting error tag to inferred span");
         this.inferredSpan.setTag("error", error);
       }
+
+      const lambdaSpan = this.tracerWrapper.currentSpan;
+      if (lambdaSpan) {
+        const appsecJson = lambdaSpan._tags?.["_dd.appsec.json"];
+        if (appsecJson) {
+          this.inferredSpan.setTag("_dd.appsec.json", appsecJson);
+        }
+      }
+
       if (this.inferredSpan.isAsync()) {
         finishTime = this.wrappedCurrentSpan?.startTime() || Date.now();
       } else {