From 07b8cb9e711c9d5c3ef10c6efae5c02a7afbae82 Mon Sep 17 00:00:00 2001 From: Carles Capell Date: Thu, 3 Sep 2026 17:02:31 +0200 Subject: [PATCH 1/4] feat(appsec): publish response body on end invocation --- src/appsec/index.spec.ts | 106 +++++++++++++++++++++++++++++++++++++++ src/appsec/index.ts | 23 +++++++++ 2 files changed, 129 insertions(+) diff --git a/src/appsec/index.spec.ts b/src/appsec/index.spec.ts index 694f1bb5..e9c375f4 100644 --- a/src/appsec/index.spec.ts +++ b/src/appsec/index.spec.ts @@ -87,6 +87,8 @@ describe("AppSec orchestrator", () => { span, statusCode: "200", responseHeaders: { "content-type": "application/json" }, + responseBody: undefined, + isBase64Encoded: false, }); }); @@ -99,6 +101,8 @@ describe("AppSec orchestrator", () => { span, statusCode: undefined, responseHeaders: undefined, + responseBody: {}, + isBase64Encoded: false, }); }); @@ -111,6 +115,8 @@ describe("AppSec orchestrator", () => { span, statusCode: "502", responseHeaders: undefined, + responseBody: undefined, + isBase64Encoded: false, }); }); @@ -123,6 +129,8 @@ describe("AppSec orchestrator", () => { span, statusCode: "200", responseHeaders: { "content-type": "application/json" }, + responseBody: undefined, + isBase64Encoded: false, }); }); @@ -135,6 +143,8 @@ describe("AppSec orchestrator", () => { span, statusCode: "200", responseHeaders: { "x-option": "test_value" }, + responseBody: undefined, + isBase64Encoded: false, }); }); @@ -155,6 +165,8 @@ describe("AppSec orchestrator", () => { span, statusCode: "200", responseHeaders: { "content-type": "application/json", "x-option": "a, b" }, + responseBody: undefined, + isBase64Encoded: false, }); }); @@ -167,6 +179,8 @@ describe("AppSec orchestrator", () => { span, statusCode: "200", responseHeaders: {}, + responseBody: undefined, + isBase64Encoded: false, }); }); @@ -179,6 +193,8 @@ describe("AppSec orchestrator", () => { span, statusCode: "200", responseHeaders: { "content-length": "42", "x-flag": "true" }, + responseBody: undefined, + isBase64Encoded: false, }); }); @@ -191,6 +207,8 @@ describe("AppSec orchestrator", () => { span, statusCode: "200", responseHeaders: {}, + responseBody: undefined, + isBase64Encoded: false, }); }); @@ -203,6 +221,94 @@ describe("AppSec orchestrator", () => { span, statusCode: undefined, responseHeaders: undefined, + responseBody: undefined, + isBase64Encoded: false, + }); + }); + + it("should publish the body of a proxy integration response", () => { + const span = { setTag: jest.fn() }; + const body = JSON.stringify({ payload: { key: "value" } }); + + processAppsecResponse(span, { statusCode: 200, headers: { "Content-Type": "application/json" }, body }, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { "content-type": "application/json" }, + responseBody: body, + isBase64Encoded: false, + }); + }); + + it("should publish the body raw, without parsing or decoding it", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, { statusCode: 200, body: "eyJhIjoiYiJ9", isBase64Encoded: true }, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: undefined, + responseBody: "eyJhIjoiYiJ9", + isBase64Encoded: true, + }); + }); + + it("should publish the whole result as the body when it is not a proxy integration response", () => { + const span = { setTag: jest.fn() }; + const result = { message: "ok", items: [1, 2] }; + + processAppsecResponse(span, result, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: undefined, + responseBody: result, + isBase64Encoded: false, + }); + }); + + it("should publish a non object result as the body", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, "plain text", "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: undefined, + responseBody: "plain text", + isBase64Encoded: false, + }); + }); + + it("should publish no body when the handler returned nothing", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, undefined, "502"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "502", + responseHeaders: undefined, + responseBody: undefined, + isBase64Encoded: false, + }); + }); + + it("should publish no body when a proxy integration response carries none", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, { statusCode: 204, body: null }, "204"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "204", + responseHeaders: undefined, + responseBody: undefined, + isBase64Encoded: false, }); }); }); diff --git a/src/appsec/index.ts b/src/appsec/index.ts index c28a0951..a0ca379e 100644 --- a/src/appsec/index.ts +++ b/src/appsec/index.ts @@ -43,6 +43,8 @@ export function processAppsecResponse(span: any, result: any, statusCode?: strin span, statusCode, responseHeaders: normalizeResponseHeaders(result), + responseBody: extractResponseBody(result), + isBase64Encoded: !!result?.isBase64Encoded, }); } @@ -57,3 +59,24 @@ function normalizeResponseHeaders(result: any): Record | undefin return normalizeHeaders(headers, multiValueHeaders); } + +/** + * Keys that mark a result as a proxy integration response rather than a payload. This is the same + * rule API Gateway itself applies to decide whether the handler answered with an envelope or with + * the body directly. + */ +const PROXY_RESPONSE_KEYS = ["statusCode", "body", "headers", "multiValueHeaders"]; + +/** + * The body is published raw, exactly as the handler wrote it. Base64 decoding, content type gating + * and size limits belong to the tracer, which is the side that knows what the WAF accepts. + */ +function extractResponseBody(result: any): unknown { + if (result === undefined || result === null) return undefined; + + if (typeof result !== "object") return result; + + if (PROXY_RESPONSE_KEYS.some((key) => key in result)) return result.body ?? undefined; + + return result; +} From a7d5691c4e77dd275c326722e5e41ce8345a2b77 Mon Sep 17 00:00:00 2001 From: Carles Capell Date: Sat, 5 Sep 2026 22:51:03 +0200 Subject: [PATCH 2/4] fix(appsec): use the statusCode discriminator and default the response content type --- src/appsec/index.spec.ts | 67 ++++++++++++++++++++++++++++++++++------ src/appsec/index.ts | 21 ++++++------- 2 files changed, 67 insertions(+), 21 deletions(-) diff --git a/src/appsec/index.spec.ts b/src/appsec/index.spec.ts index e9c375f4..1773cbc5 100644 --- a/src/appsec/index.spec.ts +++ b/src/appsec/index.spec.ts @@ -100,7 +100,7 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: undefined, - responseHeaders: undefined, + responseHeaders: { "content-type": "application/json" }, responseBody: {}, isBase64Encoded: false, }); @@ -114,7 +114,7 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "502", - responseHeaders: undefined, + responseHeaders: { "content-type": "application/json" }, responseBody: undefined, isBase64Encoded: false, }); @@ -123,13 +123,15 @@ describe("AppSec orchestrator", () => { it("should publish the normalized status code when the result carries none", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { headers: { "content-type": "application/json" } }, "200"); + const result = { headers: { "content-type": "application/json" } }; + + processAppsecResponse(span, result, "200"); expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "200", responseHeaders: { "content-type": "application/json" }, - responseBody: undefined, + responseBody: result, isBase64Encoded: false, }); }); @@ -220,7 +222,7 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: undefined, - responseHeaders: undefined, + responseHeaders: { "content-type": "application/json" }, responseBody: undefined, isBase64Encoded: false, }); @@ -241,6 +243,21 @@ describe("AppSec orchestrator", () => { }); }); + it("should default the response content type to json when the result carries no headers", () => { + const span = { setTag: jest.fn() }; + const body = JSON.stringify({ payload: 1 }); + + processAppsecResponse(span, { statusCode: 200, body }, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { "content-type": "application/json" }, + responseBody: body, + isBase64Encoded: false, + }); + }); + it("should publish the body raw, without parsing or decoding it", () => { const span = { setTag: jest.fn() }; @@ -249,7 +266,7 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "200", - responseHeaders: undefined, + responseHeaders: { "content-type": "application/json" }, responseBody: "eyJhIjoiYiJ9", isBase64Encoded: true, }); @@ -264,7 +281,37 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "200", - responseHeaders: undefined, + responseHeaders: { "content-type": "application/json" }, + responseBody: result, + isBase64Encoded: false, + }); + }); + + it("should preserve an inferred payload that carries a body key", () => { + const span = { setTag: jest.fn() }; + const result = { body: { value: 1 } }; + + processAppsecResponse(span, result, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { "content-type": "application/json" }, + responseBody: result, + isBase64Encoded: false, + }); + }); + + it("should preserve an inferred payload that carries a multiValueHeaders key", () => { + const span = { setTag: jest.fn() }; + const result = { multiValueHeaders: { count: [2] } }; + + processAppsecResponse(span, result, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { count: "2" }, responseBody: result, isBase64Encoded: false, }); @@ -278,7 +325,7 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "200", - responseHeaders: undefined, + responseHeaders: { "content-type": "application/json" }, responseBody: "plain text", isBase64Encoded: false, }); @@ -292,7 +339,7 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "502", - responseHeaders: undefined, + responseHeaders: { "content-type": "application/json" }, responseBody: undefined, isBase64Encoded: false, }); @@ -306,7 +353,7 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "204", - responseHeaders: undefined, + responseHeaders: { "content-type": "application/json" }, responseBody: undefined, isBase64Encoded: false, }); diff --git a/src/appsec/index.ts b/src/appsec/index.ts index a0ca379e..a47e4e68 100644 --- a/src/appsec/index.ts +++ b/src/appsec/index.ts @@ -49,34 +49,33 @@ export function processAppsecResponse(span: any, result: any, statusCode?: strin } /** - * Response headers reach the tracer in the same shape as the request ones + * Response headers reach the tracer in the same shape as the request ones. A result that carries + * no headers at all is served by API Gateway and by Function URLs as `application/json`, so that + * is the default, which is also what makes a raw string body eligible for schema extraction. */ -function normalizeResponseHeaders(result: any): Record | undefined { +function normalizeResponseHeaders(result: any): Record { const headers = result?.headers as Record | undefined; const multiValueHeaders = result?.multiValueHeaders as Record | undefined; - if (!headers && !multiValueHeaders) return undefined; + if (!headers && !multiValueHeaders) return { "content-type": "application/json" }; return normalizeHeaders(headers, multiValueHeaders); } -/** - * Keys that mark a result as a proxy integration response rather than a payload. This is the same - * rule API Gateway itself applies to decide whether the handler answered with an envelope or with - * the body directly. - */ -const PROXY_RESPONSE_KEYS = ["statusCode", "body", "headers", "multiValueHeaders"]; - /** * The body is published raw, exactly as the handler wrote it. Base64 decoding, content type gating * and size limits belong to the tracer, which is the side that knows what the WAF accepts. + * + * `statusCode` is the discriminator API Gateway itself uses: without it, payload format 2.0 and + * Function URLs treat the whole result as the body, so keys like `body` or `headers` are payload + * data rather than an envelope. */ function extractResponseBody(result: any): unknown { if (result === undefined || result === null) return undefined; if (typeof result !== "object") return result; - if (PROXY_RESPONSE_KEYS.some((key) => key in result)) return result.body ?? undefined; + if ("statusCode" in result) return result.body ?? undefined; return result; } From 9a38b6d3d403fbf4ff4da399401ac735849600a1 Mon Sep 17 00:00:00 2001 From: Carles Capell Date: Sat, 5 Sep 2026 23:04:45 +0200 Subject: [PATCH 3/4] fix(appsec): key the response envelope on a serializable statusCode --- src/appsec/index.spec.ts | 49 ++++++++++++++++++++++++++++++++++++++-- src/appsec/index.ts | 43 ++++++++++++----------------------- 2 files changed, 62 insertions(+), 30 deletions(-) diff --git a/src/appsec/index.spec.ts b/src/appsec/index.spec.ts index 1773cbc5..cfbb25d4 100644 --- a/src/appsec/index.spec.ts +++ b/src/appsec/index.spec.ts @@ -302,7 +302,22 @@ describe("AppSec orchestrator", () => { }); }); - it("should preserve an inferred payload that carries a multiValueHeaders key", () => { + it("should not read the headers of an inferred payload as response headers", () => { + const span = { setTag: jest.fn() }; + const result = { headers: { "content-type": "text/plain" }, payload: 1 }; + + processAppsecResponse(span, result, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { "content-type": "application/json" }, + responseBody: result, + isBase64Encoded: false, + }); + }); + + it("should not read the multi value headers of an inferred payload as response headers", () => { const span = { setTag: jest.fn() }; const result = { multiValueHeaders: { count: [2] } }; @@ -311,7 +326,37 @@ describe("AppSec orchestrator", () => { expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "200", - responseHeaders: { count: "2" }, + responseHeaders: { "content-type": "application/json" }, + responseBody: result, + isBase64Encoded: false, + }); + }); + + it("should not read the base64 flag of an inferred payload", () => { + const span = { setTag: jest.fn() }; + const result = { isBase64Encoded: true, payload: 1 }; + + processAppsecResponse(span, result, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { "content-type": "application/json" }, + responseBody: result, + isBase64Encoded: false, + }); + }); + + it("should treat a result whose status code does not survive serialization as an inferred payload", () => { + const span = { setTag: jest.fn() }; + const result = { statusCode: undefined, body: { value: 1 } }; + + processAppsecResponse(span, result, "200"); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { "content-type": "application/json" }, responseBody: result, isBase64Encoded: false, }); diff --git a/src/appsec/index.ts b/src/appsec/index.ts index a47e4e68..c060ecc8 100644 --- a/src/appsec/index.ts +++ b/src/appsec/index.ts @@ -39,43 +39,30 @@ export function processAppsecRequest(event: any, span: any): void { export function processAppsecResponse(span: any, result: any, statusCode?: string): void { if (!span || !endInvocationChannel.hasSubscribers) return; + const structured = isStructuredResponse(result); + endInvocationChannel.publish({ span, statusCode, - responseHeaders: normalizeResponseHeaders(result), - responseBody: extractResponseBody(result), - isBase64Encoded: !!result?.isBase64Encoded, + responseHeaders: structured ? normalizeResponseHeaders(result) : inferredResponseHeaders(), + responseBody: structured ? result.body ?? undefined : result ?? undefined, + isBase64Encoded: structured && !!result.isBase64Encoded, }); } -/** - * Response headers reach the tracer in the same shape as the request ones. A result that carries - * no headers at all is served by API Gateway and by Function URLs as `application/json`, so that - * is the default, which is also what makes a raw string body eligible for schema extraction. - */ -function normalizeResponseHeaders(result: any): Record { - const headers = result?.headers as Record | undefined; - const multiValueHeaders = result?.multiValueHeaders as Record | undefined; - - if (!headers && !multiValueHeaders) return { "content-type": "application/json" }; - - return normalizeHeaders(headers, multiValueHeaders); +function isStructuredResponse(result: any): boolean { + return typeof result === "object" && result !== null && result.statusCode !== undefined; } -/** - * The body is published raw, exactly as the handler wrote it. Base64 decoding, content type gating - * and size limits belong to the tracer, which is the side that knows what the WAF accepts. - * - * `statusCode` is the discriminator API Gateway itself uses: without it, payload format 2.0 and - * Function URLs treat the whole result as the body, so keys like `body` or `headers` are payload - * data rather than an envelope. - */ -function extractResponseBody(result: any): unknown { - if (result === undefined || result === null) return undefined; +function inferredResponseHeaders(): Record { + return { "content-type": "application/json" }; +} - if (typeof result !== "object") return result; +function normalizeResponseHeaders(result: any): Record { + const headers = result.headers as Record | undefined; + const multiValueHeaders = result.multiValueHeaders as Record | undefined; - if ("statusCode" in result) return result.body ?? undefined; + if (!headers && !multiValueHeaders) return inferredResponseHeaders(); - return result; + return normalizeHeaders(headers, multiValueHeaders); } From b03f9e8583001873b45fad99a4244b9900b1f64a Mon Sep 17 00:00:00 2001 From: Carles Capell Date: Tue, 15 Sep 2026 12:30:24 +0200 Subject: [PATCH 4/4] fix(appsec): scope response data to what the trigger actually serves --- src/appsec/index.spec.ts | 161 ++++++++++++++++++++++++++++++------- src/appsec/index.ts | 57 ++++++++++--- src/trace/listener.spec.ts | 65 +++++++++++++-- src/trace/listener.ts | 14 +++- src/trace/trigger.spec.ts | 51 +++++++++++- src/trace/trigger.ts | 14 ++++ 6 files changed, 313 insertions(+), 49 deletions(-) diff --git a/src/appsec/index.spec.ts b/src/appsec/index.spec.ts index cfbb25d4..04c6b580 100644 --- a/src/appsec/index.spec.ts +++ b/src/appsec/index.spec.ts @@ -74,14 +74,17 @@ describe("AppSec orchestrator", () => { describe("processAppSecResponse", () => { it("should not publish when span is falsy", () => { - processAppsecResponse(null, { statusCode: 200 }); + processAppsecResponse(null, { statusCode: 200 }, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).not.toHaveBeenCalled(); }); it("should publish the normalized status code and the response headers", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 200, headers: { "content-type": "application/json" } }, "200"); + processAppsecResponse(span, { statusCode: 200, headers: { "content-type": "application/json" } }, "200", { + kind: "buffered", + supportsInference: true, + }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -95,7 +98,7 @@ describe("AppSec orchestrator", () => { it("should publish with undefined statusCode and headers when result has none", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, {}); + processAppsecResponse(span, {}, undefined, { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -109,12 +112,12 @@ describe("AppSec orchestrator", () => { it("should ignore the status code carried by the result", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 200 }, "502"); + processAppsecResponse(span, { statusCode: 200 }, "502", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "502", - responseHeaders: { "content-type": "application/json" }, + responseHeaders: undefined, responseBody: undefined, isBase64Encoded: false, }); @@ -125,7 +128,7 @@ describe("AppSec orchestrator", () => { const result = { headers: { "content-type": "application/json" } }; - processAppsecResponse(span, result, "200"); + processAppsecResponse(span, result, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -139,7 +142,10 @@ describe("AppSec orchestrator", () => { it("should lowercase the response header names", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 200, headers: { "X-Option": "test_value" } }, "200"); + processAppsecResponse(span, { statusCode: 200, headers: { "X-Option": "test_value" } }, "200", { + kind: "buffered", + supportsInference: true, + }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -161,6 +167,7 @@ describe("AppSec orchestrator", () => { multiValueHeaders: { "X-Option": ["a", "b"] }, }, "200", + { kind: "buffered", supportsInference: true }, ); expect(mockPublish).toHaveBeenCalledWith({ @@ -175,7 +182,10 @@ describe("AppSec orchestrator", () => { it("should ignore multi value response headers that are not arrays", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 200, multiValueHeaders: { "Set-Cookie": "a=b" } }, "200"); + processAppsecResponse(span, { statusCode: 200, multiValueHeaders: { "Set-Cookie": "a=b" } }, "200", { + kind: "buffered", + supportsInference: true, + }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -189,7 +199,10 @@ describe("AppSec orchestrator", () => { it("should stringify non string response header values", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 200, headers: { "Content-Length": 42, "X-Flag": true } }, "200"); + processAppsecResponse(span, { statusCode: 200, headers: { "Content-Length": 42, "X-Flag": true } }, "200", { + kind: "buffered", + supportsInference: true, + }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -203,7 +216,10 @@ describe("AppSec orchestrator", () => { it("should skip response headers with a null value", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 200, headers: { "X-Option": null } }, "200"); + processAppsecResponse(span, { statusCode: 200, headers: { "X-Option": null } }, "200", { + kind: "buffered", + supportsInference: true, + }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -217,12 +233,12 @@ describe("AppSec orchestrator", () => { it("should publish no status code when none is normalized, even if the result carries one", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 204 }, undefined); + processAppsecResponse(span, { statusCode: 204 }, undefined, { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: undefined, - responseHeaders: { "content-type": "application/json" }, + responseHeaders: undefined, responseBody: undefined, isBase64Encoded: false, }); @@ -232,7 +248,10 @@ describe("AppSec orchestrator", () => { const span = { setTag: jest.fn() }; const body = JSON.stringify({ payload: { key: "value" } }); - processAppsecResponse(span, { statusCode: 200, headers: { "Content-Type": "application/json" }, body }, "200"); + processAppsecResponse(span, { statusCode: 200, headers: { "Content-Type": "application/json" }, body }, "200", { + kind: "buffered", + supportsInference: true, + }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -243,16 +262,16 @@ describe("AppSec orchestrator", () => { }); }); - it("should default the response content type to json when the result carries no headers", () => { + it("should publish no headers when a structured response carries none", () => { const span = { setTag: jest.fn() }; const body = JSON.stringify({ payload: 1 }); - processAppsecResponse(span, { statusCode: 200, body }, "200"); + processAppsecResponse(span, { statusCode: 200, body }, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "200", - responseHeaders: { "content-type": "application/json" }, + responseHeaders: undefined, responseBody: body, isBase64Encoded: false, }); @@ -261,12 +280,15 @@ describe("AppSec orchestrator", () => { it("should publish the body raw, without parsing or decoding it", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 200, body: "eyJhIjoiYiJ9", isBase64Encoded: true }, "200"); + processAppsecResponse(span, { statusCode: 200, body: "eyJhIjoiYiJ9", isBase64Encoded: true }, "200", { + kind: "buffered", + supportsInference: true, + }); expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "200", - responseHeaders: { "content-type": "application/json" }, + responseHeaders: undefined, responseBody: "eyJhIjoiYiJ9", isBase64Encoded: true, }); @@ -276,7 +298,7 @@ describe("AppSec orchestrator", () => { const span = { setTag: jest.fn() }; const result = { message: "ok", items: [1, 2] }; - processAppsecResponse(span, result, "200"); + processAppsecResponse(span, result, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -291,7 +313,7 @@ describe("AppSec orchestrator", () => { const span = { setTag: jest.fn() }; const result = { body: { value: 1 } }; - processAppsecResponse(span, result, "200"); + processAppsecResponse(span, result, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -306,7 +328,7 @@ describe("AppSec orchestrator", () => { const span = { setTag: jest.fn() }; const result = { headers: { "content-type": "text/plain" }, payload: 1 }; - processAppsecResponse(span, result, "200"); + processAppsecResponse(span, result, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -321,7 +343,7 @@ describe("AppSec orchestrator", () => { const span = { setTag: jest.fn() }; const result = { multiValueHeaders: { count: [2] } }; - processAppsecResponse(span, result, "200"); + processAppsecResponse(span, result, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -336,7 +358,7 @@ describe("AppSec orchestrator", () => { const span = { setTag: jest.fn() }; const result = { isBase64Encoded: true, payload: 1 }; - processAppsecResponse(span, result, "200"); + processAppsecResponse(span, result, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -351,7 +373,7 @@ describe("AppSec orchestrator", () => { const span = { setTag: jest.fn() }; const result = { statusCode: undefined, body: { value: 1 } }; - processAppsecResponse(span, result, "200"); + processAppsecResponse(span, result, "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -362,10 +384,88 @@ describe("AppSec orchestrator", () => { }); }); + it("should publish neither body nor headers when the trigger does not infer responses", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, { payload: 1 }, "200", { kind: "buffered", supportsInference: false }); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: undefined, + responseBody: undefined, + isBase64Encoded: false, + }); + }); + + it("should still publish the headers of a structured response when the trigger does not infer", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, { statusCode: 200, headers: { "X-Option": "a" } }, "200", { + kind: "buffered", + supportsInference: false, + }); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: { "x-option": "a" }, + responseBody: undefined, + isBase64Encoded: false, + }); + }); + + it("should publish no response data for a streaming function that returned nothing", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, undefined, "200", { kind: "streaming" }); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: undefined, + responseBody: undefined, + isBase64Encoded: false, + }); + }); + + it("should publish no response data for a streaming function that returned a payload", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse(span, { payload: 1 }, "200", { kind: "streaming" }); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "200", + responseHeaders: undefined, + responseBody: undefined, + isBase64Encoded: false, + }); + }); + + it("should publish no response data for a streaming function that returned a structured response", () => { + const span = { setTag: jest.fn() }; + + processAppsecResponse( + span, + { statusCode: 201, headers: { "Content-Type": "text/plain" }, body: "streamed", isBase64Encoded: true }, + "201", + { kind: "streaming" }, + ); + + expect(mockPublish).toHaveBeenCalledWith({ + span, + statusCode: "201", + responseHeaders: undefined, + responseBody: undefined, + isBase64Encoded: false, + }); + }); + it("should publish a non object result as the body", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, "plain text", "200"); + processAppsecResponse(span, "plain text", "200", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, @@ -379,12 +479,12 @@ describe("AppSec orchestrator", () => { it("should publish no body when the handler returned nothing", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, undefined, "502"); + processAppsecResponse(span, undefined, "502", { kind: "buffered", supportsInference: true }); expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "502", - responseHeaders: { "content-type": "application/json" }, + responseHeaders: undefined, responseBody: undefined, isBase64Encoded: false, }); @@ -393,12 +493,15 @@ describe("AppSec orchestrator", () => { it("should publish no body when a proxy integration response carries none", () => { const span = { setTag: jest.fn() }; - processAppsecResponse(span, { statusCode: 204, body: null }, "204"); + processAppsecResponse(span, { statusCode: 204, body: null }, "204", { + kind: "buffered", + supportsInference: true, + }); expect(mockPublish).toHaveBeenCalledWith({ span, statusCode: "204", - responseHeaders: { "content-type": "application/json" }, + responseHeaders: undefined, responseBody: undefined, isBase64Encoded: false, }); diff --git a/src/appsec/index.ts b/src/appsec/index.ts index c060ecc8..212e1104 100644 --- a/src/appsec/index.ts +++ b/src/appsec/index.ts @@ -31,38 +31,71 @@ export function processAppsecRequest(event: any, span: any): void { }); } +export type ResponseMode = { kind: "streaming" } | { kind: "buffered"; supportsInference: boolean }; + /** * @param span * @param result * @param statusCode Status code already normalized by the trigger layer. + * @param mode `streaming` when the function writes its response to responseStream, in which case + * the returned value is not what the client received. `buffered` otherwise, with + * `supportsInference` telling whether the trigger serves a result without a status + * code as the body, resolved by the trigger layer before the handler ran. */ -export function processAppsecResponse(span: any, result: any, statusCode?: string): void { +export function processAppsecResponse( + span: any, + result: any, + statusCode: string | undefined, + mode: ResponseMode, +): void { if (!span || !endInvocationChannel.hasSubscribers) return; - const structured = isStructuredResponse(result); - endInvocationChannel.publish({ span, statusCode, - responseHeaders: structured ? normalizeResponseHeaders(result) : inferredResponseHeaders(), - responseBody: structured ? result.body ?? undefined : result ?? undefined, - isBase64Encoded: structured && !!result.isBase64Encoded, + ...extractResponseData(result, mode), }); } -function isStructuredResponse(result: any): boolean { - return typeof result === "object" && result !== null && result.statusCode !== undefined; +function extractResponseData(result: any, mode: ResponseMode) { + // Streaming functions write the real status, headers and body to responseStream, so nothing + // about the response can be derived from the returned value. + if (mode.kind === "streaming") return noResponseData(); + + if (isStructuredResponse(result)) { + return { + responseHeaders: normalizeResponseHeaders(result), + responseBody: result.body ?? undefined, + isBase64Encoded: !!result.isBase64Encoded, + }; + } + + // The client is answered by the integration error AWS builds, which is not this result. + if (!mode.supportsInference) return noResponseData(); + + // Nothing was returned, so there is no result for the trigger to serve as a JSON body. + if (result === undefined || result === null) return noResponseData(); + + return { + responseHeaders: { "content-type": "application/json" }, + responseBody: result, + isBase64Encoded: false, + }; +} + +function noResponseData() { + return { responseHeaders: undefined, responseBody: undefined, isBase64Encoded: false }; } -function inferredResponseHeaders(): Record { - return { "content-type": "application/json" }; +function isStructuredResponse(result: any): boolean { + return typeof result === "object" && result !== null && result.statusCode !== undefined; } -function normalizeResponseHeaders(result: any): Record { +function normalizeResponseHeaders(result: any): Record | undefined { const headers = result.headers as Record | undefined; const multiValueHeaders = result.multiValueHeaders as Record | undefined; - if (!headers && !multiValueHeaders) return inferredResponseHeaders(); + if (!headers && !multiValueHeaders) return undefined; return normalizeHeaders(headers, multiValueHeaders); } diff --git a/src/trace/listener.spec.ts b/src/trace/listener.spec.ts index 9cede1f8..93ec3f36 100644 --- a/src/trace/listener.spec.ts +++ b/src/trace/listener.spec.ts @@ -700,7 +700,35 @@ describe("TraceListener", () => { expect(mockProcessAppsecResponse).toHaveBeenCalledTimes(1); // Non-HTTP trigger: there is no normalized status code to hand over. - expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, result, undefined); + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, result, undefined, { + kind: "buffered", + supportsInference: false, + }); + } finally { + currentSpanSpy.mockRestore(); + } + }); + + it("resolves the inferred response mode before the handler runs, ignoring later event mutations", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const event = JSON.parse(readFileSync("./event_samples/api-gateway-v2.json", "utf8")); + const result = { payload: 1 }; + await listener.onStartInvocation(event, context as any); + listener.onRequestStart(event); + + delete event.version; + delete event.requestContext; + listener.onEndingInvocation(event, result, false); + + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, result, "200", { + kind: "buffered", + supportsInference: true, + }); } finally { currentSpanSpy.mockRestore(); } @@ -720,7 +748,10 @@ describe("TraceListener", () => { await listener.onStartInvocation(event, context as any); listener.onEndingInvocation(event, result, false); - expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, result, "200"); + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, result, "200", { + kind: "buffered", + supportsInference: false, + }); } finally { currentSpanSpy.mockRestore(); } @@ -737,7 +768,10 @@ describe("TraceListener", () => { await listener.onStartInvocation(event, context as any); listener.onEndingInvocation(event, undefined, false); - expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, undefined, "502"); + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, undefined, "502", { + kind: "buffered", + supportsInference: false, + }); } finally { currentSpanSpy.mockRestore(); } @@ -754,7 +788,25 @@ describe("TraceListener", () => { await listener.onStartInvocation(event, context as any); listener.onEndingInvocation(event, undefined, true); - expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, undefined, "200"); + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, undefined, "200", { kind: "streaming" }); + } finally { + currentSpanSpy.mockRestore(); + } + }); + + it("flags a streaming function so no response data is derived from the returned value", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const event = JSON.parse(readFileSync("./event_samples/lambda-function-urls.json", "utf8")); + await listener.onStartInvocation(event, context as any); + listener.onRequestStart(event); + listener.onEndingInvocation(event, undefined, true); + + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, undefined, "200", { kind: "streaming" }); } finally { currentSpanSpy.mockRestore(); } @@ -794,7 +846,10 @@ describe("TraceListener", () => { const responseIs5xxError = listener.onEndingInvocation(event, { statusCode: 500 }, false); expect(responseIs5xxError).toBe(true); - expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, { statusCode: 500 }, "500"); + expect(mockProcessAppsecResponse).toHaveBeenCalledWith(mockSpan, { statusCode: 500 }, "500", { + kind: "buffered", + supportsInference: false, + }); } finally { currentSpanSpy.mockRestore(); } diff --git a/src/trace/listener.ts b/src/trace/listener.ts index 9766a867..aff0647f 100644 --- a/src/trace/listener.ts +++ b/src/trace/listener.ts @@ -2,7 +2,7 @@ import { Context } from "aws-lambda"; import { patchHttp, unpatchHttp } from "./patch-http"; -import { extractTriggerTags, extractHTTPStatusCodeTag, parseEventSource } from "./trigger"; +import { extractTriggerTags, extractHTTPStatusCodeTag, parseEventSource, supportsInferredResponse } from "./trigger"; import { ColdStartTracerConfig, ColdStartTracer } from "./cold-start-tracer"; import { logDebug, tagObject } from "../utils"; import { @@ -108,6 +108,7 @@ export class TraceListener { private inferredSpan?: SpanWrapper; private wrappedCurrentSpan?: SpanWrapper; private triggerTags?: { [key: string]: string }; + private inferredResponseSupported = false; private lambdaSpanParentContext?: SpanContext; private spanPointerAttributesList: SpanPointerAttributes[] | undefined; @@ -179,6 +180,8 @@ export class TraceListener { */ public onRequestStart(event: any): void { if (!this.config.appsecEnabled) return; + // Resolved here because the user function receives this very object and may mutate it. + this.inferredResponseSupported = supportsInferredResponse(event); processAppsecRequest(event, this.tracerWrapper.currentSpan); } @@ -243,7 +246,14 @@ export class TraceListener { this.inferredSpan?.setTag("http.status_code", statusCode); } if (this.config.appsecEnabled) { - processAppsecResponse(this.tracerWrapper.currentSpan, result, statusCode); + processAppsecResponse( + this.tracerWrapper.currentSpan, + result, + statusCode, + isResponseStreamFunction + ? { kind: "streaming" } + : { kind: "buffered", supportsInference: this.inferredResponseSupported }, + ); } // Kept behind AppSec so 5xx responses still reach the WAF, and still nested on inferredSpan // so the early return only happens when there is an inferred span, as before. diff --git a/src/trace/trigger.spec.ts b/src/trace/trigger.spec.ts index 01a35288..c48432b8 100644 --- a/src/trace/trigger.spec.ts +++ b/src/trace/trigger.spec.ts @@ -1,4 +1,10 @@ -import { parseEventSource, parseEventSourceARN, extractTriggerTags, extractHTTPStatusCodeTag } from "./trigger"; +import { + parseEventSource, + parseEventSourceARN, + extractTriggerTags, + extractHTTPStatusCodeTag, + supportsInferredResponse, +} from "./trigger"; import { readFileSync } from "fs"; import { Context } from "aws-lambda"; @@ -237,3 +243,46 @@ describe("parseEventSource", () => { } }); }); + +describe("supportsInferredResponse", () => { + it("should support an HTTP API payload format 2.0 trigger", () => { + expect( + supportsInferredResponse({ + version: "2.0", + rawQueryString: "", + requestContext: { domainName: "abc.execute-api.eu-west-1.amazonaws.com" }, + }), + ).toBe(true); + }); + + it("should support a function url trigger", () => { + expect( + supportsInferredResponse({ + version: "2.0", + rawQueryString: "", + requestContext: { domainName: "abc.lambda-url.eu-west-1.on.aws" }, + }), + ).toBe(true); + }); + + it("should not support a REST API payload format 1.0 trigger", () => { + expect(supportsInferredResponse({ requestContext: { stage: "dev" }, httpMethod: "GET", resource: "/" })).toBe( + false, + ); + }); + + it("should not support an ALB trigger", () => { + expect( + supportsInferredResponse({ requestContext: { elb: { targetGroupArn: "arn" } }, httpMethod: "GET", path: "/" }), + ).toBe(false); + }); + + it("should not support a non HTTP trigger", () => { + expect(supportsInferredResponse({ Records: [] })).toBe(false); + }); + + it("should not throw on a missing event", () => { + expect(supportsInferredResponse(undefined)).toBe(false); + expect(supportsInferredResponse("string event")).toBe(false); + }); +}); diff --git a/src/trace/trigger.ts b/src/trace/trigger.ts index dc95d3a3..cd7885d7 100644 --- a/src/trace/trigger.ts +++ b/src/trace/trigger.ts @@ -373,6 +373,20 @@ export function extractTriggerTags(event: any, context: Context, eventSource: ev return triggerTags; } +/** + * Only HTTP API payload format 2.0 and Function URLs serve a result without a status code as the + * response body. REST API payload format 1.0 and ALB require the structured shape, and answer the + * client with an integration error instead, so such a result never reaches it as a body. + * + * The event is the same mutable object handed to the user function, so this has to be resolved + * before the handler runs. + */ +export function supportsInferredResponse(event: any): boolean { + if (!event || typeof event !== "object") return false; + + return eventType.isLambdaUrlEvent(event) || eventType.isAPIGatewayEventV2(event); +} + /** * extractHTTPStatusCode extracts a status code from the response if the Lambda was triggered * by API Gateway, ALB, or Lambda Function URL