From c1560d856ddba6d52b52bafaa9952e9248a1e71b Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Mon, 21 Sep 2026 15:17:39 +0200 Subject: [PATCH 1/8] Complete Lambda AppSec telemetry integration --- .../gateway/GatewayBridgeSpecification.groovy | 32 ++ .../lambda/LambdaHandlerInstrumentation.java | 56 ++-- .../src/test/java/HandlerStreamingNested.java | 13 + ...dlerStreamingWritesResponseThenThrows.java | 18 ++ .../LambdaHandlerInstrumentationTest.java | 35 ++- .../java/datadog/trace/core/CoreTracer.java | 7 +- .../trace/lambda/LambdaAppSecHandler.java | 285 ++++++++++-------- .../trace/lambda/LambdaAppSecHandlerTest.java | 202 ++++++++++++- 8 files changed, 492 insertions(+), 156 deletions(-) create mode 100644 dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingNested.java create mode 100644 dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingWritesResponseThenThrows.java diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index f648fabfc83..75fed0d7dac 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -21,6 +21,7 @@ import datadog.trace.api.function.TriFunction import datadog.appsec.api.blocking.BlockingContentType import datadog.trace.bootstrap.blocking.BlockingActionHelper import datadog.trace.api.gateway.BlockResponseFunction +import datadog.trace.api.gateway.CallbackProvider import datadog.trace.api.gateway.Flow import datadog.trace.api.gateway.IGSpanInfo import datadog.trace.api.gateway.RequestContext @@ -33,9 +34,11 @@ import datadog.trace.api.telemetry.LoginEvent import datadog.trace.api.telemetry.RuleType import datadog.trace.api.telemetry.WafMetricCollector import datadog.trace.bootstrap.instrumentation.api.AgentSpan +import datadog.trace.bootstrap.instrumentation.api.AgentTracer import datadog.trace.bootstrap.instrumentation.api.Tags import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase +import datadog.trace.lambda.LambdaAppSecHandler import datadog.trace.test.util.DDSpecification import spock.lang.Shared @@ -215,6 +218,35 @@ class GatewayBridgeSpecification extends DDSpecification { flow.action == Flow.Action.Noop.INSTANCE } + void 'lambda request end reaches shared waf telemetry with its framework'() { + given: + AgentTracer.TracerAPI originalTracer = AgentTracer.get() + CallbackProvider callbackProvider = Stub() { + getCallback(EVENTS.requestEnded()) >> requestEndedCB + } + AgentTracer.TracerAPI tracer = Stub() { + getCallbackProvider(RequestContextSlot.APPSEC) >> callbackProvider + } + AgentSpan span = Mock() { + getRequestContext() >> ctx + getTags() >> TagMap.fromMap([(Tags.COMPONENT): 'aws-lambda']) + } + AgentTracer.forceRegister(tracer) + + when: + LambdaAppSecHandler.processRequestEnd(span) + + then: + 1 * requestSampler.preSampleRequest(arCtx, 'aws-lambda') >> false + 1 * span.setMetric('_dd.appsec.enabled', 1) + 1 * span.setTag('_dd.runtime_family', 'jvm') + 1 * pp.processTraceSegment(traceSegment, arCtx, []) + 1 * wafMetricCollector.wafRequest(false, false, false, false, false, false, false, false) + + cleanup: + AgentTracer.forceRegister(originalTracer) + } + void 'actor ip calculated from headers'() { AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext) mockAppSecCtx.requestHeaders >> [ diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index cb8d369d51e..fda8ecd695d 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -121,31 +121,43 @@ static void exit( final AgentSpan span = scope.span(); try { - if (throwable == null) { - AgentTracer.get().notifyAppSecEnd(span, result); - } else { + if (throwable != null) { span.addThrowable(throwable); } - // Force the resource name back to the literal placeholder marker right - // before finish so that the Datadog Lambda Extension's filter - // (filter_span_from_lambda_library_or_runtime in - // bottlecap/src/traces/trace_processor.rs, which compares - // span.resource == "dd-tracer-serverless-span") drops the placeholder. - // Other instrumentation (HTTP/JAX-RS) may have overwritten it with the - // route ("POST /") during the invocation, in which case the extension - // would fail to dedup, leading to the placeholder leaking to the backend - // with parent_id=0 and detaching the inferred apigateway root from the - // rest of the trace. - // Use TAG_INTERCEPTOR priority because DDSpanContext.setResourceName - // ignores writes whose priority is below the current resource priority, - // and the HTTP/JAX-RS instrumentation will already have written - // HTTP_FRAMEWORK_ROUTE (3) by this point. - span.setResourceName(INVOCATION_SPAN_NAME, ResourceNamePriorities.TAG_INTERCEPTOR); } finally { - scope.close(); - span.finish(); - AgentTracer.get() - .notifyExtensionEnd(span, result, null != throwable, awsContext.getAwsRequestId()); + try { + AgentTracer.get().notifyAppSecEnd(span, throwable == null ? result : null); + } finally { + try { + // Force the resource name back to the literal placeholder marker right + // before finish so that the Datadog Lambda Extension's filter + // (filter_span_from_lambda_library_or_runtime in + // bottlecap/src/traces/trace_processor.rs, which compares + // span.resource == "dd-tracer-serverless-span") drops the placeholder. + // Other instrumentation (HTTP/JAX-RS) may have overwritten it with the + // route ("POST /") during the invocation, in which case the extension + // would fail to dedup, leading to the placeholder leaking to the backend + // with parent_id=0 and detaching the inferred apigateway root from the + // rest of the trace. + // Use TAG_INTERCEPTOR priority because DDSpanContext.setResourceName + // ignores writes whose priority is below the current resource priority, + // and the HTTP/JAX-RS instrumentation will already have written + // HTTP_FRAMEWORK_ROUTE (3) by this point. + span.setResourceName(INVOCATION_SPAN_NAME, ResourceNamePriorities.TAG_INTERCEPTOR); + } finally { + try { + scope.close(); + } finally { + try { + span.finish(); + } finally { + AgentTracer.get() + .notifyExtensionEnd( + span, result, null != throwable, awsContext.getAwsRequestId()); + } + } + } + } } } } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingNested.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingNested.java new file mode 100644 index 00000000000..ce4188a7256 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingNested.java @@ -0,0 +1,13 @@ +import com.amazonaws.services.lambda.runtime.Context; +import com.amazonaws.services.lambda.runtime.RequestStreamHandler; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; + +public class HandlerStreamingNested implements RequestStreamHandler { + @Override + public void handleRequest(InputStream inputStream, OutputStream outputStream, Context context) + throws IOException { + new HandlerStreamingWithApiGwResponse().handleRequest(inputStream, outputStream, context); + } +} diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingWritesResponseThenThrows.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingWritesResponseThenThrows.java new file mode 100644 index 00000000000..4169b8a2d9f --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingWritesResponseThenThrows.java @@ -0,0 +1,18 @@ +import com.amazonaws.services.lambda.runtime.Context; +import com.amazonaws.services.lambda.runtime.RequestStreamHandler; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; + +public class HandlerStreamingWritesResponseThenThrows implements RequestStreamHandler { + @Override + public void handleRequest(InputStream inputStream, OutputStream outputStream, Context context) + throws IOException { + outputStream.write( + ("{\"statusCode\":200,\"headers\":{\"content-type\":\"application/json\"}," + + "\"body\":\"{\\\"discarded\\\":true}\"}") + .getBytes(StandardCharsets.UTF_8)); + throw new Error("Some error"); + } +} diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java index a25e2f8fb4f..3e5e78b8056 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java @@ -59,6 +59,7 @@ abstract class LambdaHandlerInstrumentationTest extends AbstractInstrumentationT Map capturedHeaders; Object capturedBody; boolean appSecEnded; + int appSecEndCount; Integer capturedResponseStatus; Map capturedResponseHeaders; @@ -82,6 +83,7 @@ void setUpAppSec() { capturedHeaders = new HashMap<>(); capturedBody = null; appSecEnded = false; + appSecEndCount = 0; capturedResponseStatus = null; capturedResponseHeaders = new HashMap<>(); capturedResponseBody = null; @@ -121,6 +123,7 @@ void setUpAppSec() { (BiFunction>) (ctx2, spanInfo) -> { appSecEnded = true; + appSecEndCount++; return Flow.ResultFlow.empty(); }); @@ -201,6 +204,7 @@ void testStreamingHandlerWithError() { .tags( defaultTags(), tag("request_id", is(REQUEST_ID)), + tag("_dd.appsec.unsupported_event_type", is(1)), error(Error.class, "Some error")))); } @@ -487,9 +491,26 @@ void invocationSpanCarriesHttpTags() throws IOException { tag(Tags.HTTP_USER_AGENT, is("test-agent")), tag(Tags.HTTP_ROUTE, is("/api/users/{id}")), tag(Tags.HTTP_HOSTNAME, is("api.example.com")), + tag(Tags.COMPONENT, is("aws-lambda")), tag(Tags.HTTP_STATUS, is(200))))); } + @Test + void nestedHandlerFinalizesAppSecRequestOnce() throws IOException { + String eventJson = + "{" + "\"path\": \"/api/nested\"," + "\"requestContext\": {\"httpMethod\": \"GET\"}" + "}"; + ByteArrayInputStream input = + new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8)); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + + new HandlerStreamingNested().handleRequest(input, output, newContext()); + + assertTrue(appSecStarted); + assertTrue(appSecEnded); + assertEquals(1, appSecEndCount); + assertTraces(trace(span().type(DDSpanTypes.SERVERLESS).error(false))); + } + @Test void responseCallbacksFireBeforeRequestEnded() throws IOException { List callOrder = new ArrayList<>(); @@ -555,13 +576,21 @@ void responseCallbacksFireBeforeRequestEnded() throws IOException { @Test void responseCallbacksReceiveNoDataWhenHandlerThrows() { - ByteArrayInputStream input = new ByteArrayInputStream("Hello".getBytes(StandardCharsets.UTF_8)); + String eventJson = + "{" + "\"path\": \"/api/failure\"," + "\"requestContext\": {\"httpMethod\": \"GET\"}" + "}"; + ByteArrayInputStream input = + new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8)); ByteArrayOutputStream output = new ByteArrayOutputStream(); assertThrows( Error.class, - () -> new HandlerStreamingWithError().handleRequest(input, output, newContext())); + () -> + new HandlerStreamingWritesResponseThenThrows() + .handleRequest(input, output, newContext())); + assertTrue(appSecStarted, "request callbacks should run before the handler throws"); + assertTrue(appSecEnded, "requestEnded should run after the handler throws"); + assertEquals(1, appSecEndCount); assertNull(capturedResponseStatus, "response status should not be set when handler throws"); assertNull(capturedResponseBody, "response body should not be set when handler throws"); assertTraces( @@ -572,6 +601,8 @@ void responseCallbacksReceiveNoDataWhenHandlerThrows() { .tags( defaultTags(), tag("request_id", is(REQUEST_ID)), + tag(Tags.HTTP_METHOD, is("GET")), + tag(Tags.COMPONENT, is("aws-lambda")), error(Error.class, "Some error")))); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index e9e130847b9..fa735966e0b 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -1330,8 +1330,11 @@ public void notifyExtensionEnd( @Override public void notifyAppSecEnd(AgentSpan span, Object result) { - LambdaAppSecHandler.processResponseData(span, result); - LambdaAppSecHandler.processRequestEnd(span); + try { + LambdaAppSecHandler.processResponseData(span, result); + } finally { + LambdaAppSecHandler.processRequestEnd(span); + } } @Override diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java index 32a36f9c207..4d308c2e302 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java @@ -36,6 +36,7 @@ import datadog.trace.lambda.LambdaEventParser.LambdaTriggerType; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.io.Closeable; import java.nio.charset.StandardCharsets; import java.util.List; import java.util.Map; @@ -110,11 +111,16 @@ public static AgentSpanContext processRequestStart(Object event) { } LambdaURIDataAdapter uriAdapter = new LambdaURIDataAdapter(fullPath, eventData.headers, eventData.host); - AgentSpanContext context = processAppSecRequestData(eventData, uriAdapter); - if (context instanceof TagContext) { - applyHttpTags((TagContext) context, eventData, uriAdapter); + AgentSpanContext appSecContext = processAppSecRequestData(eventData, uriAdapter); + try { + if (appSecContext instanceof TagContext) { + applyHttpTags((TagContext) appSecContext, eventData, uriAdapter); + } + return appSecContext; + } catch (Exception e) { + closeAppSecData(appSecContext); + throw e; } - return context; } catch (Exception e) { log.debug("Failed to process AppSec request data", e); return null; @@ -131,19 +137,14 @@ public static void processRequestEnd(AgentSpan span) { LambdaTriggerType triggerType = CURRENT_TRIGGER_TYPE.get(); CURRENT_TRIGGER_TYPE.remove(); - if (!ActiveSubsystems.APPSEC_ACTIVE || span == null || triggerType == null) { - return; - } - - // A null trigger type means processRequestStart never ran, so the invocation was not analysed - // at all, which is not the same as an unsupported trigger. - if (!triggerType.isHttp()) { - span.setMetric(UNSUPPORTED_EVENT_TYPE_METRIC, 1); + if (span == null) { return; } RequestContext requestContext = span.getRequestContext(); - if (requestContext != null) { + Object rawAppSecCtx = + requestContext != null ? requestContext.getData(RequestContextSlot.APPSEC) : null; + if (rawAppSecCtx != null) { AgentTracer.TracerAPI tracer = AgentTracer.get(); BiFunction> requestEndedCallback = tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestEnded()); @@ -157,7 +158,6 @@ public static void processRequestEnd(AgentSpan span) { // GatewayBridge propagates ASM_KEEP based on WAF attack events, but not on // isManuallyKept(), which is set by trace-tagging rules that produce no events. // Apply it here so those traces are not silently dropped. - Object rawAppSecCtx = requestContext.getData(RequestContextSlot.APPSEC); AppSecContext appSecCtx = rawAppSecCtx instanceof AppSecContext ? (AppSecContext) rawAppSecCtx : null; if (appSecCtx != null && appSecCtx.isManuallyKept()) { @@ -165,6 +165,13 @@ public static void processRequestEnd(AgentSpan span) { traceSeg.setTagTop(Tags.ASM_KEEP, true); traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); } + return; + } + + // A null trigger type means processRequestStart never ran, so the invocation was not analysed + // at all, which is not the same as an unsupported trigger. + if (triggerType != null && !triggerType.isHttp()) { + span.setMetric(UNSUPPORTED_EVENT_TYPE_METRIC, 1); } } @@ -177,9 +184,7 @@ public static void processRequestEnd(AgentSpan span) { * @param result the Lambda handler result (expected to be a ByteArrayOutputStream) */ public static void processResponseData(AgentSpan span, Object result) { - if (!ActiveSubsystems.APPSEC_ACTIVE - || span == null - || !(result instanceof ByteArrayOutputStream)) { + if (span == null || !(result instanceof ByteArrayOutputStream)) { return; } @@ -189,6 +194,12 @@ public static void processResponseData(AgentSpan span, Object result) { return; } + RequestContext requestContext = span.getRequestContext(); + if (requestContext == null || requestContext.getData(RequestContextSlot.APPSEC) == null) { + log.debug("Span has no AppSec request context, skipping response processing"); + return; + } + try { byte[] bytes = ((ByteArrayOutputStream) result).toByteArray(); if (bytes.length == 0 || bytes.length > MAX_EVENT_SIZE) { @@ -214,12 +225,6 @@ public static void processResponseData(AgentSpan span, Object result) { span.setError(isError, ErrorPriorities.HTTP_SERVER_DECORATOR); } - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { - log.debug("Span has no RequestContext, skipping response processing"); - return; - } - AgentTracer.TracerAPI tracer = AgentTracer.get(); CallbackProvider cbp = tracer.getCallbackProvider(RequestContextSlot.APPSEC); @@ -291,24 +296,33 @@ public static AgentSpanContext mergeContexts( if (extensionContext instanceof TagContext) { TagContext merged = (TagContext) extensionContext; - if (appSecData != null) { - merged.withRequestContextDataAppSec(appSecData); - } - // The extension context is the one that survives, so the HTTP tags applied to the AppSec - // context have to be carried over: CoreTracer copies them onto the span at creation. - // The AppSec-derived values win on a key collision. No collision is reachable today: the - // extension context only carries tags for headers mapped through - // DD_TRACE_REQUEST_HEADER_TAGS - // (ContextInterpreter.handleTags), and those would have to be mapped onto an http.* key. - for (TagMap.EntryReader tag : extracted.getTags()) { - merged.putTag(tag.tag(), tag.stringValue()); + boolean transferred = false; + try { + // The extension context is the one that survives, so the HTTP tags applied to the AppSec + // context have to be carried over: CoreTracer copies them onto the span at creation. + // The AppSec-derived values win on a key collision. No collision is reachable today: the + // extension context only carries tags for headers mapped through + // DD_TRACE_REQUEST_HEADER_TAGS + // (ContextInterpreter.handleTags), and those would have to be mapped onto an http.* key. + for (TagMap.EntryReader tag : extracted.getTags()) { + merged.putTag(tag.tag(), tag.stringValue()); + } + if (appSecData != null) { + merged.withRequestContextDataAppSec(appSecData); + } + transferred = true; + return merged; + } finally { + if (!transferred) { + closeAppSecData(appSecData); + } } - return merged; } rlLog.warn( "Cannot merge AppSec data: extension context is not a TagContext: {}", extensionContext.getClass()); + closeAppSecData(appSecData); } return extensionContext; } @@ -319,6 +333,8 @@ public static AgentSpanContext mergeContexts( * tags, {@code span.kind} and {@code http.fragment}. */ static void applyHttpTags(TagContext ctx, LambdaRequestData req, LambdaURIDataAdapter url) { + ctx.putTag(Tags.COMPONENT, "aws-lambda"); + // The synthetic "WEBSOCKET" method stays inside the AppSec path; none is fabricated here. if (req.method != null && req.triggerType != LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET) { ctx.putTag(Tags.HTTP_METHOD, req.method); @@ -370,111 +386,134 @@ private static AgentSpanContext processAppSecRequestData( } TagContext tagContext = new TagContext(); - Object appSecRequestContext; - - // Call requestStarted - appSecRequestContext = requestStartedCallback.get().getResult(); - tagContext.withRequestContextDataAppSec(appSecRequestContext); - - if (appSecRequestContext != null) { - TemporaryRequestContext requestContext = new TemporaryRequestContext(appSecRequestContext); - - // Call requestMethodUriRaw - if (eventData.method != null && eventData.path != null) { - datadog.trace.api.function.TriFunction> - methodUriCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestMethodUriRaw()); - if (methodUriCallback != null) { - methodUriCallback.apply(requestContext, eventData.method, uriAdapter); - } else { - log.debug("requestMethodUriRaw callback is null"); + Object appSecRequestContext = null; + boolean transferred = false; + try { + // Call requestStarted + appSecRequestContext = requestStartedCallback.get().getResult(); + tagContext.withRequestContextDataAppSec(appSecRequestContext); + + if (appSecRequestContext != null) { + TemporaryRequestContext requestContext = new TemporaryRequestContext(appSecRequestContext); + + // Call requestMethodUriRaw + if (eventData.method != null && eventData.path != null) { + datadog.trace.api.function.TriFunction> + methodUriCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestMethodUriRaw()); + if (methodUriCallback != null) { + methodUriCallback.apply(requestContext, eventData.method, uriAdapter); + } else { + log.debug("requestMethodUriRaw callback is null"); + } } - } - // Call requestHeader for each header - if (eventData.headers != null && !eventData.headers.isEmpty()) { - TriConsumer headerCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestHeader()); - if (headerCallback != null) { - for (Map.Entry header : eventData.headers.entrySet()) { - headerCallback.accept(requestContext, header.getKey(), header.getValue()); + // Call requestHeader for each header + if (eventData.headers != null && !eventData.headers.isEmpty()) { + TriConsumer headerCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestHeader()); + if (headerCallback != null) { + for (Map.Entry header : eventData.headers.entrySet()) { + headerCallback.accept(requestContext, header.getKey(), header.getValue()); + } + } else { + log.debug("requestHeader callback is null"); } - } else { - log.debug("requestHeader callback is null"); } - } - // Call requestClientSocketAddress - if (eventData.sourceIp != null) { - datadog.trace.api.function.TriFunction> - socketAddrCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestClientSocketAddress()); - if (socketAddrCallback != null) { - Integer port = eventData.sourcePort != null ? eventData.sourcePort : 0; - socketAddrCallback.apply(requestContext, eventData.sourceIp, port); - } else { - log.debug("requestClientSocketAddress callback is null"); + // Call requestClientSocketAddress + if (eventData.sourceIp != null) { + datadog.trace.api.function.TriFunction> + socketAddrCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestClientSocketAddress()); + if (socketAddrCallback != null) { + Integer port = eventData.sourcePort != null ? eventData.sourcePort : 0; + socketAddrCallback.apply(requestContext, eventData.sourceIp, port); + } else { + log.debug("requestClientSocketAddress callback is null"); + } } - } - // Call requestHeaderDone - Function> headerDoneCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestHeaderDone()); - if (headerDoneCallback != null) { - headerDoneCallback.apply(requestContext); - } else { - log.debug("requestHeaderDone callback is null"); - } - - // Call requestPathParams - if (eventData.pathParameters != null && !eventData.pathParameters.isEmpty()) { - BiFunction, Flow> pathParamsCallback = + // Call requestHeaderDone + Function> headerDoneCallback = tracer .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestPathParams()); - if (pathParamsCallback != null) { - pathParamsCallback.apply(requestContext, eventData.pathParameters); + .getCallback(EVENTS.requestHeaderDone()); + if (headerDoneCallback != null) { + headerDoneCallback.apply(requestContext); } else { - log.debug("requestPathParams callback is null"); + log.debug("requestHeaderDone callback is null"); } - } - // Call requestBodyProcessed - if (eventData.body != null) { - BiFunction> bodyCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestBodyProcessed()); - if (bodyCallback != null) { - bodyCallback.apply(requestContext, eventData.body); - } else { - log.debug("requestBodyProcessed callback is null"); + // Call requestPathParams + if (eventData.pathParameters != null && !eventData.pathParameters.isEmpty()) { + BiFunction, Flow> pathParamsCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestPathParams()); + if (pathParamsCallback != null) { + pathParamsCallback.apply(requestContext, eventData.pathParameters); + } else { + log.debug("requestPathParams callback is null"); + } } - } - // Call requestFilesFilenames. Only the names are reported: the file content shares the - // body's UTF-8 decode, so for anything that is not text it is already lossy. - if (!eventData.filenames.isEmpty()) { - BiFunction, Flow> filenamesCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestFilesFilenames()); - if (filenamesCallback != null) { - filenamesCallback.apply(requestContext, eventData.filenames); - } else { - log.debug("requestFilesFilenames callback is null"); + // Call requestBodyProcessed + if (eventData.body != null) { + BiFunction> bodyCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestBodyProcessed()); + if (bodyCallback != null) { + bodyCallback.apply(requestContext, eventData.body); + } else { + log.debug("requestBodyProcessed callback is null"); + } } + + // Call requestFilesFilenames. Only the names are reported: the file content shares the + // body's UTF-8 decode, so for anything that is not text it is already lossy. + if (!eventData.filenames.isEmpty()) { + BiFunction, Flow> filenamesCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestFilesFilenames()); + if (filenamesCallback != null) { + filenamesCallback.apply(requestContext, eventData.filenames); + } else { + log.debug("requestFilesFilenames callback is null"); + } + } + } + transferred = true; + return tagContext; + } finally { + if (!transferred) { + closeAppSecData(appSecRequestContext); + } + } + } + + private static void closeAppSecData(AgentSpanContext context) { + if (context instanceof TagContext) { + closeAppSecData(((TagContext) context).getRequestContextDataAppSec()); + } + } + + private static void closeAppSecData(Object appSecData) { + if (appSecData instanceof Closeable) { + try { + ((Closeable) appSecData).close(); + } catch (Exception e) { + log.debug("Failed to close abandoned AppSec request context", e); } } - return tagContext; } /** Sets the current trigger type thread-local. Package-private for use in tests only. */ diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java index e3104df0db5..8ba85a83b51 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java @@ -12,6 +12,7 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; @@ -48,11 +49,13 @@ import datadog.trace.bootstrap.instrumentation.api.TagContext; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; +import datadog.trace.core.CoreTracer; import datadog.trace.core.DDCoreJavaSpecification; import datadog.trace.lambda.LambdaEventParser.LambdaResponseData; import datadog.trace.lambda.LambdaEventParser.LambdaTriggerType; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.io.Closeable; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; @@ -1216,13 +1219,14 @@ void processRequestEndDoesNothingWhenSpanIsNull() { } @Test - void processRequestEndDoesNothingWhenAppSecIsDisabled() { + void processRequestEndDoesNothingWithoutAttachedContextWhenAppSecIsDisabled() { ActiveSubsystems.APPSEC_ACTIVE = false; AgentSpan span = mock(AgentSpan.class); LambdaAppSecHandler.processRequestEnd(span); - verifyNoInteractions(span); + verify(span).getRequestContext(); + verifyNoMoreInteractions(span); } @Test @@ -1265,10 +1269,58 @@ void processRequestEndInvokesCallbackAndSkipsAsmTagsWhenNotManuallyKept() { verify(mockTraceSegment, never()).setTagTop(any(), any()); } + @Test + @SuppressWarnings("unchecked") + void processRequestEndFinalizesAttachedContextWhenAppSecBecomesInactive() { + ActiveSubsystems.APPSEC_ACTIVE = false; + Object appSecContext = new Object(); + RequestContext requestContext = mock(RequestContext.class); + when(requestContext.getData(RequestContextSlot.APPSEC)).thenReturn(appSecContext); + AgentSpan span = mock(AgentSpan.class); + when(span.getRequestContext()).thenReturn(requestContext); + + BiFunction> requestEndedCallback = + mock(BiFunction.class); + when(requestEndedCallback.apply(any(), any())).thenReturn(Flow.ResultFlow.empty()); + CallbackProvider callbackProvider = mock(CallbackProvider.class); + when(callbackProvider.getCallback(EVENTS.requestEnded())).thenReturn(requestEndedCallback); + AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); + when(tracer.getCallbackProvider(RequestContextSlot.APPSEC)).thenReturn(callbackProvider); + AgentTracer.forceRegister(tracer); + + LambdaAppSecHandler.processRequestEnd(span); + + verify(requestEndedCallback).apply(requestContext, span); + } + + @Test + @SuppressWarnings("unchecked") + void processRequestEndFinalizesAttachedContextWithoutRecordedTrigger() { + Object appSecContext = new Object(); + RequestContext requestContext = mock(RequestContext.class); + when(requestContext.getData(RequestContextSlot.APPSEC)).thenReturn(appSecContext); + AgentSpan span = mock(AgentSpan.class); + when(span.getRequestContext()).thenReturn(requestContext); + + BiFunction> requestEndedCallback = + mock(BiFunction.class); + when(requestEndedCallback.apply(any(), any())).thenReturn(Flow.ResultFlow.empty()); + CallbackProvider callbackProvider = mock(CallbackProvider.class); + when(callbackProvider.getCallback(EVENTS.requestEnded())).thenReturn(requestEndedCallback); + AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); + when(tracer.getCallbackProvider(RequestContextSlot.APPSEC)).thenReturn(callbackProvider); + AgentTracer.forceRegister(tracer); + + LambdaAppSecHandler.processRequestEnd(span); + + verify(requestEndedCallback).apply(requestContext, span); + } + @Test void processRequestEndHandlesNullRequestEndedCallbackGracefully() { LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); RequestContext mockRequestContext = mock(RequestContext.class); + when(mockRequestContext.getData(RequestContextSlot.APPSEC)).thenReturn(new Object()); AgentSpan span = mock(AgentSpan.class); when(span.getRequestContext()).thenReturn(mockRequestContext); @@ -1324,6 +1376,7 @@ void processRequestEndSetsUnsupportedEventTypeMetricForNonHttpTrigger() { LambdaAppSecHandler.processRequestEnd(span); + verify(span).getRequestContext(); verify(span).setMetric("_dd.appsec.unsupported_event_type", 1); verifyNoMoreInteractions(span); } @@ -1339,14 +1392,16 @@ void processRequestEndSetsNoUnsupportedEventTypeMetricWhenNoTriggerTypeWasRecord } @Test - void processRequestEndSetsNoUnsupportedEventTypeMetricWhenAppSecIsDisabled() { + void processRequestEndPreservesUnsupportedMetricWhenAppSecBecomesDisabled() { ActiveSubsystems.APPSEC_ACTIVE = false; LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.UNKNOWN); AgentSpan span = mock(AgentSpan.class); LambdaAppSecHandler.processRequestEnd(span); - verifyNoInteractions(span); + verify(span).getRequestContext(); + verify(span).setMetric("_dd.appsec.unsupported_event_type", 1); + verifyNoMoreInteractions(span); } @Test @@ -1354,6 +1409,7 @@ void processRequestEndSetsNoUnsupportedEventTypeMetricWhenAppSecIsDisabled() { void processRequestEndSetsNoUnsupportedEventTypeMetricForHttpTrigger() { LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); RequestContext mockRequestContext = mock(RequestContext.class); + when(mockRequestContext.getData(RequestContextSlot.APPSEC)).thenReturn(new Object()); when(mockRequestContext.getTraceSegment()).thenReturn(mock(TraceSegment.class)); AgentSpan span = mock(AgentSpan.class); when(span.getRequestContext()).thenReturn(mockRequestContext); @@ -1374,6 +1430,44 @@ void processRequestEndSetsNoUnsupportedEventTypeMetricForHttpTrigger() { verify(span, never()).setMetric(anyString(), anyInt()); } + @Test + void processRequestEndDoesNotReportHttpTriggerWithoutAttachedAppSecContextAsUnsupported() { + LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); + RequestContext requestContext = mock(RequestContext.class); + AgentSpan span = mock(AgentSpan.class); + when(span.getRequestContext()).thenReturn(requestContext); + + LambdaAppSecHandler.processRequestEnd(span); + + verify(span, never()).setMetric(anyString(), anyInt()); + } + + @Test + @SuppressWarnings("unchecked") + void processRequestEndClearsTriggerAfterSuccessfulFinalization() { + LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); + RequestContext requestContext = mock(RequestContext.class); + when(requestContext.getData(RequestContextSlot.APPSEC)).thenReturn(new Object()); + AgentSpan span = mock(AgentSpan.class); + when(span.getRequestContext()).thenReturn(requestContext); + + BiFunction> requestEndedCallback = + mock(BiFunction.class); + when(requestEndedCallback.apply(any(), any())).thenReturn(Flow.ResultFlow.empty()); + CallbackProvider callbackProvider = mock(CallbackProvider.class); + when(callbackProvider.getCallback(EVENTS.requestEnded())).thenReturn(requestEndedCallback); + AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); + when(tracer.getCallbackProvider(RequestContextSlot.APPSEC)).thenReturn(callbackProvider); + AgentTracer.forceRegister(tracer); + + LambdaAppSecHandler.processRequestEnd(span); + + AgentSpan nextSpan = mock(AgentSpan.class); + LambdaAppSecHandler.processRequestEnd(nextSpan); + verify(requestEndedCallback).apply(requestContext, span); + verify(nextSpan, never()).setMetric(anyString(), anyInt()); + } + @Test void processRequestEndSetsUnsupportedEventTypeMetricAfterAnUnparseablePayload() { ByteArrayInputStream event = createInputStream("{invalid json"); @@ -1382,6 +1476,7 @@ void processRequestEndSetsUnsupportedEventTypeMetricAfterAnUnparseablePayload() AgentSpan span = mock(AgentSpan.class); LambdaAppSecHandler.processRequestEnd(span); + verify(span).getRequestContext(); verify(span).setMetric("_dd.appsec.unsupported_event_type", 1); verifyNoMoreInteractions(span); } @@ -1436,6 +1531,33 @@ void mergeContextsReturnsExtensionContextWhenItIsNotTagContext() { extensionContext, LambdaAppSecHandler.mergeContexts(extensionContext, appSecContext)); } + @Test + void mergeContextsClosesAppSecDataWhenExtensionContextIsIncompatible() throws IOException { + Closeable appSecData = mock(Closeable.class); + TagContext appSecContext = new TagContext(); + appSecContext.withRequestContextDataAppSec(appSecData); + AgentSpanContext extensionContext = mock(AgentSpanContext.class); + + assertSame( + extensionContext, LambdaAppSecHandler.mergeContexts(extensionContext, appSecContext)); + + verify(appSecData).close(); + } + + @Test + void mergeContextsDoesNotCloseTransferredAppSecData() throws IOException { + Closeable appSecData = mock(Closeable.class); + TagContext appSecContext = new TagContext(); + appSecContext.withRequestContextDataAppSec(appSecData); + TagContext extensionContext = new TagContext(); + + assertSame( + extensionContext, LambdaAppSecHandler.mergeContexts(extensionContext, appSecContext)); + + assertSame(appSecData, extensionContext.getRequestContextDataAppSec()); + verify(appSecData, never()).close(); + } + // ============================================================================ // Error Handling and Null Callback Tests // ============================================================================ @@ -1457,6 +1579,29 @@ void processRequestStartHandlesNullRequestStartedCallbackGracefully() { assertNull(LambdaAppSecHandler.processRequestStart(event)); } + @Test + @SuppressWarnings("unchecked") + void processRequestStartClosesAppSecDataWhenRequestProcessingFails() throws IOException { + String eventJson = "{\"path\": \"/test\", \"requestContext\": {\"httpMethod\": \"GET\"}}"; + Closeable appSecData = mock(Closeable.class); + Supplier> requestStartedCallback = mock(Supplier.class); + when(requestStartedCallback.get()).thenReturn(new Flow.ResultFlow<>(appSecData)); + TriFunction> methodUriCallback = + mock(TriFunction.class); + when(methodUriCallback.apply(any(), anyString(), any(URIDataAdapter.class))) + .thenThrow(new IllegalStateException("request processing failed")); + CallbackProvider callbackProvider = mock(CallbackProvider.class); + when(callbackProvider.getCallback(EVENTS.requestStarted())).thenReturn(requestStartedCallback); + when(callbackProvider.getCallback(EVENTS.requestMethodUriRaw())).thenReturn(methodUriCallback); + AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); + when(tracer.getCallbackProvider(RequestContextSlot.APPSEC)).thenReturn(callbackProvider); + AgentTracer.forceRegister(tracer); + + assertNull(LambdaAppSecHandler.processRequestStart(createInputStream(eventJson))); + + verify(appSecData).close(); + } + @Test @SuppressWarnings("unchecked") void processRequestStartHandlesNullMethodUriCallbackGracefully() { @@ -1658,12 +1803,48 @@ private RequestContext captureTemporaryRequestContext(Object appSecContext) { // ============================================================================ @Test - void processResponseDataDoesNothingWhenAppSecIsDisabled() { + void processResponseDataProcessesAttachedContextWhenAppSecBecomesDisabled() { ActiveSubsystems.APPSEC_ACTIVE = false; - AgentSpan span = mock(AgentSpan.class); + LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); + Integer[] capturedStatus = {null}; + AgentSpan span = + setupMockResponseCallbacks(status -> capturedStatus[0] = status, null, null, null); ByteArrayOutputStream result = createOutputStream("{\"statusCode\": 200, \"body\": \"ok\"}"); + LambdaAppSecHandler.processResponseData(span, result); - verify(span, never()).getRequestContext(); + + assertEquals(200, capturedStatus[0]); + } + + @Test + @SuppressWarnings("unchecked") + void notifyAppSecEndFinalizesRequestWhenResponseProcessingThrows() { + LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); + RequestContext requestContext = mock(RequestContext.class); + when(requestContext.getData(RequestContextSlot.APPSEC)).thenReturn(new Object()); + AgentSpan span = mock(AgentSpan.class); + when(span.getRequestContext()).thenReturn(requestContext); + + BiFunction> requestEndedCallback = + mock(BiFunction.class); + when(requestEndedCallback.apply(any(), any())).thenReturn(Flow.ResultFlow.empty()); + CallbackProvider callbackProvider = mock(CallbackProvider.class); + when(callbackProvider.getCallback(EVENTS.requestEnded())).thenReturn(requestEndedCallback); + AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); + when(tracer.getCallbackProvider(RequestContextSlot.APPSEC)).thenReturn(callbackProvider); + AgentTracer.forceRegister(tracer); + + ByteArrayOutputStream result = + new ByteArrayOutputStream() { + @Override + public synchronized byte[] toByteArray() { + throw new AssertionError("response processing failed"); + } + }; + + CoreTracer coreTracer = tracerBuilder().build(); + assertThrows(AssertionError.class, () -> coreTracer.notifyAppSecEnd(span, result)); + verify(requestEndedCallback).apply(requestContext, span); } @Test @@ -2377,6 +2558,7 @@ void appliesHttpTagsForRestApiEvent() { + " \"api.example.com\"}}")); Map tags = tagsOf(context); + assertEquals("aws-lambda", tags.get(Tags.COMPONENT)); assertEquals("GET", tags.get(Tags.HTTP_METHOD)); assertEquals("https://api.example.com/users/42", tags.get(Tags.HTTP_URL)); assertEquals("q=hello", tags.get(DDTags.HTTP_QUERY)); @@ -2624,6 +2806,7 @@ void omitsRouteTagForAlbEvent() { @Test void mergeContextsCopiesHttpTagsIntoExtensionContext() { TagContext appSecContext = new TagContext(); + appSecContext.putTag(Tags.COMPONENT, "aws-lambda"); appSecContext.putTag(Tags.HTTP_URL, "https://api.example.com/users/42"); appSecContext.putTag(Tags.HTTP_ROUTE, "/users/{id}"); TagContext extensionContext = new TagContext(); @@ -2632,8 +2815,12 @@ void mergeContextsCopiesHttpTagsIntoExtensionContext() { assertSame(extensionContext, merged); Map tags = tagsOf(merged); + assertEquals("aws-lambda", tags.get(Tags.COMPONENT)); assertEquals("https://api.example.com/users/42", tags.get(Tags.HTTP_URL)); assertEquals("/users/{id}", tags.get(Tags.HTTP_ROUTE)); + + ((TagContext) merged).putTag(Tags.COMPONENT, "spring-web-controller"); + assertEquals("spring-web-controller", tagsOf(merged).get(Tags.COMPONENT)); } @Test @@ -2866,6 +3053,7 @@ private AgentSpan setupMockResponseCallbacks( Consumer onResponseBody) { RequestContext mockRequestContext = mock(RequestContext.class); + when(mockRequestContext.getData(RequestContextSlot.APPSEC)).thenReturn(new Object()); AgentSpan mockSpan = mock(AgentSpan.class); when(mockSpan.getRequestContext()).thenReturn(mockRequestContext); From df852ad96a2892ac67c3ead7a5345ec153dbc447 Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Mon, 21 Sep 2026 16:43:52 +0200 Subject: [PATCH 2/8] Fix AppSec test CodeNarc violations --- .../appsec/gateway/GatewayBridgeSpecification.groovy | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index 75fed0d7dac..7073ef3ecf9 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -221,13 +221,13 @@ class GatewayBridgeSpecification extends DDSpecification { void 'lambda request end reaches shared waf telemetry with its framework'() { given: AgentTracer.TracerAPI originalTracer = AgentTracer.get() - CallbackProvider callbackProvider = Stub() { + CallbackProvider callbackProvider = Stub { getCallback(EVENTS.requestEnded()) >> requestEndedCB } - AgentTracer.TracerAPI tracer = Stub() { + AgentTracer.TracerAPI tracer = Stub { getCallbackProvider(RequestContextSlot.APPSEC) >> callbackProvider } - AgentSpan span = Mock() { + AgentSpan span = Mock { getRequestContext() >> ctx getTags() >> TagMap.fromMap([(Tags.COMPONENT): 'aws-lambda']) } From 60c20e83eeaa61740205cc48465e3dc5ac58899c Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Mon, 21 Sep 2026 17:10:55 +0200 Subject: [PATCH 3/8] Simplify Lambda AppSec lifecycle handling --- .../lambda/LambdaHandlerInstrumentation.java | 55 ++-- .../trace/lambda/LambdaAppSecHandler.java | 246 ++++++++---------- .../trace/lambda/LambdaAppSecHandlerTest.java | 51 ---- 3 files changed, 126 insertions(+), 226 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index fda8ecd695d..73beb4a7314 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -124,40 +124,29 @@ static void exit( if (throwable != null) { span.addThrowable(throwable); } + + AgentTracer.get().notifyAppSecEnd(span, throwable == null ? result : null); + + // Force the resource name back to the literal placeholder marker right + // before finish so that the Datadog Lambda Extension's filter + // (filter_span_from_lambda_library_or_runtime in + // bottlecap/src/traces/trace_processor.rs, which compares + // span.resource == "dd-tracer-serverless-span") drops the placeholder. + // Other instrumentation (HTTP/JAX-RS) may have overwritten it with the + // route ("POST /") during the invocation, in which case the extension + // would fail to dedup, leading to the placeholder leaking to the backend + // with parent_id=0 and detaching the inferred apigateway root from the + // rest of the trace. + // Use TAG_INTERCEPTOR priority because DDSpanContext.setResourceName + // ignores writes whose priority is below the current resource priority, + // and the HTTP/JAX-RS instrumentation will already have written + // HTTP_FRAMEWORK_ROUTE (3) by this point. + span.setResourceName(INVOCATION_SPAN_NAME, ResourceNamePriorities.TAG_INTERCEPTOR); } finally { - try { - AgentTracer.get().notifyAppSecEnd(span, throwable == null ? result : null); - } finally { - try { - // Force the resource name back to the literal placeholder marker right - // before finish so that the Datadog Lambda Extension's filter - // (filter_span_from_lambda_library_or_runtime in - // bottlecap/src/traces/trace_processor.rs, which compares - // span.resource == "dd-tracer-serverless-span") drops the placeholder. - // Other instrumentation (HTTP/JAX-RS) may have overwritten it with the - // route ("POST /") during the invocation, in which case the extension - // would fail to dedup, leading to the placeholder leaking to the backend - // with parent_id=0 and detaching the inferred apigateway root from the - // rest of the trace. - // Use TAG_INTERCEPTOR priority because DDSpanContext.setResourceName - // ignores writes whose priority is below the current resource priority, - // and the HTTP/JAX-RS instrumentation will already have written - // HTTP_FRAMEWORK_ROUTE (3) by this point. - span.setResourceName(INVOCATION_SPAN_NAME, ResourceNamePriorities.TAG_INTERCEPTOR); - } finally { - try { - scope.close(); - } finally { - try { - span.finish(); - } finally { - AgentTracer.get() - .notifyExtensionEnd( - span, result, null != throwable, awsContext.getAwsRequestId()); - } - } - } - } + scope.close(); + span.finish(); + AgentTracer.get() + .notifyExtensionEnd(span, result, null != throwable, awsContext.getAwsRequestId()); } } } diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java index 4d308c2e302..fd7bd59bbe6 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java @@ -36,7 +36,6 @@ import datadog.trace.lambda.LambdaEventParser.LambdaTriggerType; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.io.Closeable; import java.nio.charset.StandardCharsets; import java.util.List; import java.util.Map; @@ -111,16 +110,11 @@ public static AgentSpanContext processRequestStart(Object event) { } LambdaURIDataAdapter uriAdapter = new LambdaURIDataAdapter(fullPath, eventData.headers, eventData.host); - AgentSpanContext appSecContext = processAppSecRequestData(eventData, uriAdapter); - try { - if (appSecContext instanceof TagContext) { - applyHttpTags((TagContext) appSecContext, eventData, uriAdapter); - } - return appSecContext; - } catch (Exception e) { - closeAppSecData(appSecContext); - throw e; + AgentSpanContext context = processAppSecRequestData(eventData, uriAdapter); + if (context instanceof TagContext) { + applyHttpTags((TagContext) context, eventData, uriAdapter); } + return context; } catch (Exception e) { log.debug("Failed to process AppSec request data", e); return null; @@ -296,33 +290,24 @@ public static AgentSpanContext mergeContexts( if (extensionContext instanceof TagContext) { TagContext merged = (TagContext) extensionContext; - boolean transferred = false; - try { - // The extension context is the one that survives, so the HTTP tags applied to the AppSec - // context have to be carried over: CoreTracer copies them onto the span at creation. - // The AppSec-derived values win on a key collision. No collision is reachable today: the - // extension context only carries tags for headers mapped through - // DD_TRACE_REQUEST_HEADER_TAGS - // (ContextInterpreter.handleTags), and those would have to be mapped onto an http.* key. - for (TagMap.EntryReader tag : extracted.getTags()) { - merged.putTag(tag.tag(), tag.stringValue()); - } - if (appSecData != null) { - merged.withRequestContextDataAppSec(appSecData); - } - transferred = true; - return merged; - } finally { - if (!transferred) { - closeAppSecData(appSecData); - } + if (appSecData != null) { + merged.withRequestContextDataAppSec(appSecData); + } + // The extension context is the one that survives, so the HTTP tags applied to the AppSec + // context have to be carried over: CoreTracer copies them onto the span at creation. + // The AppSec-derived values win on a key collision. No collision is reachable today: the + // extension context only carries tags for headers mapped through + // DD_TRACE_REQUEST_HEADER_TAGS + // (ContextInterpreter.handleTags), and those would have to be mapped onto an http.* key. + for (TagMap.EntryReader tag : extracted.getTags()) { + merged.putTag(tag.tag(), tag.stringValue()); } + return merged; } rlLog.warn( "Cannot merge AppSec data: extension context is not a TagContext: {}", extensionContext.getClass()); - closeAppSecData(appSecData); } return extensionContext; } @@ -386,134 +371,111 @@ private static AgentSpanContext processAppSecRequestData( } TagContext tagContext = new TagContext(); - Object appSecRequestContext = null; - boolean transferred = false; - try { - // Call requestStarted - appSecRequestContext = requestStartedCallback.get().getResult(); - tagContext.withRequestContextDataAppSec(appSecRequestContext); - - if (appSecRequestContext != null) { - TemporaryRequestContext requestContext = new TemporaryRequestContext(appSecRequestContext); - - // Call requestMethodUriRaw - if (eventData.method != null && eventData.path != null) { - datadog.trace.api.function.TriFunction> - methodUriCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestMethodUriRaw()); - if (methodUriCallback != null) { - methodUriCallback.apply(requestContext, eventData.method, uriAdapter); - } else { - log.debug("requestMethodUriRaw callback is null"); - } + Object appSecRequestContext; + + // Call requestStarted + appSecRequestContext = requestStartedCallback.get().getResult(); + tagContext.withRequestContextDataAppSec(appSecRequestContext); + + if (appSecRequestContext != null) { + TemporaryRequestContext requestContext = new TemporaryRequestContext(appSecRequestContext); + + // Call requestMethodUriRaw + if (eventData.method != null && eventData.path != null) { + datadog.trace.api.function.TriFunction> + methodUriCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestMethodUriRaw()); + if (methodUriCallback != null) { + methodUriCallback.apply(requestContext, eventData.method, uriAdapter); + } else { + log.debug("requestMethodUriRaw callback is null"); } + } - // Call requestHeader for each header - if (eventData.headers != null && !eventData.headers.isEmpty()) { - TriConsumer headerCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestHeader()); - if (headerCallback != null) { - for (Map.Entry header : eventData.headers.entrySet()) { - headerCallback.accept(requestContext, header.getKey(), header.getValue()); - } - } else { - log.debug("requestHeader callback is null"); + // Call requestHeader for each header + if (eventData.headers != null && !eventData.headers.isEmpty()) { + TriConsumer headerCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestHeader()); + if (headerCallback != null) { + for (Map.Entry header : eventData.headers.entrySet()) { + headerCallback.accept(requestContext, header.getKey(), header.getValue()); } + } else { + log.debug("requestHeader callback is null"); } + } - // Call requestClientSocketAddress - if (eventData.sourceIp != null) { - datadog.trace.api.function.TriFunction> - socketAddrCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestClientSocketAddress()); - if (socketAddrCallback != null) { - Integer port = eventData.sourcePort != null ? eventData.sourcePort : 0; - socketAddrCallback.apply(requestContext, eventData.sourceIp, port); - } else { - log.debug("requestClientSocketAddress callback is null"); - } + // Call requestClientSocketAddress + if (eventData.sourceIp != null) { + datadog.trace.api.function.TriFunction> + socketAddrCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestClientSocketAddress()); + if (socketAddrCallback != null) { + Integer port = eventData.sourcePort != null ? eventData.sourcePort : 0; + socketAddrCallback.apply(requestContext, eventData.sourceIp, port); + } else { + log.debug("requestClientSocketAddress callback is null"); } + } - // Call requestHeaderDone - Function> headerDoneCallback = + // Call requestHeaderDone + Function> headerDoneCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestHeaderDone()); + if (headerDoneCallback != null) { + headerDoneCallback.apply(requestContext); + } else { + log.debug("requestHeaderDone callback is null"); + } + + // Call requestPathParams + if (eventData.pathParameters != null && !eventData.pathParameters.isEmpty()) { + BiFunction, Flow> pathParamsCallback = tracer .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestHeaderDone()); - if (headerDoneCallback != null) { - headerDoneCallback.apply(requestContext); + .getCallback(EVENTS.requestPathParams()); + if (pathParamsCallback != null) { + pathParamsCallback.apply(requestContext, eventData.pathParameters); } else { - log.debug("requestHeaderDone callback is null"); - } - - // Call requestPathParams - if (eventData.pathParameters != null && !eventData.pathParameters.isEmpty()) { - BiFunction, Flow> pathParamsCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestPathParams()); - if (pathParamsCallback != null) { - pathParamsCallback.apply(requestContext, eventData.pathParameters); - } else { - log.debug("requestPathParams callback is null"); - } - } - - // Call requestBodyProcessed - if (eventData.body != null) { - BiFunction> bodyCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestBodyProcessed()); - if (bodyCallback != null) { - bodyCallback.apply(requestContext, eventData.body); - } else { - log.debug("requestBodyProcessed callback is null"); - } + log.debug("requestPathParams callback is null"); } + } - // Call requestFilesFilenames. Only the names are reported: the file content shares the - // body's UTF-8 decode, so for anything that is not text it is already lossy. - if (!eventData.filenames.isEmpty()) { - BiFunction, Flow> filenamesCallback = - tracer - .getCallbackProvider(RequestContextSlot.APPSEC) - .getCallback(EVENTS.requestFilesFilenames()); - if (filenamesCallback != null) { - filenamesCallback.apply(requestContext, eventData.filenames); - } else { - log.debug("requestFilesFilenames callback is null"); - } + // Call requestBodyProcessed + if (eventData.body != null) { + BiFunction> bodyCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestBodyProcessed()); + if (bodyCallback != null) { + bodyCallback.apply(requestContext, eventData.body); + } else { + log.debug("requestBodyProcessed callback is null"); } } - transferred = true; - return tagContext; - } finally { - if (!transferred) { - closeAppSecData(appSecRequestContext); - } - } - } - - private static void closeAppSecData(AgentSpanContext context) { - if (context instanceof TagContext) { - closeAppSecData(((TagContext) context).getRequestContextDataAppSec()); - } - } - private static void closeAppSecData(Object appSecData) { - if (appSecData instanceof Closeable) { - try { - ((Closeable) appSecData).close(); - } catch (Exception e) { - log.debug("Failed to close abandoned AppSec request context", e); + // Call requestFilesFilenames. Only the names are reported: the file content shares the + // body's UTF-8 decode, so for anything that is not text it is already lossy. + if (!eventData.filenames.isEmpty()) { + BiFunction, Flow> filenamesCallback = + tracer + .getCallbackProvider(RequestContextSlot.APPSEC) + .getCallback(EVENTS.requestFilesFilenames()); + if (filenamesCallback != null) { + filenamesCallback.apply(requestContext, eventData.filenames); + } else { + log.debug("requestFilesFilenames callback is null"); + } } } + return tagContext; } /** Sets the current trigger type thread-local. Package-private for use in tests only. */ diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java index 8ba85a83b51..c6bc868cbbf 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java @@ -55,7 +55,6 @@ import datadog.trace.lambda.LambdaEventParser.LambdaTriggerType; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.io.Closeable; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; @@ -1531,33 +1530,6 @@ void mergeContextsReturnsExtensionContextWhenItIsNotTagContext() { extensionContext, LambdaAppSecHandler.mergeContexts(extensionContext, appSecContext)); } - @Test - void mergeContextsClosesAppSecDataWhenExtensionContextIsIncompatible() throws IOException { - Closeable appSecData = mock(Closeable.class); - TagContext appSecContext = new TagContext(); - appSecContext.withRequestContextDataAppSec(appSecData); - AgentSpanContext extensionContext = mock(AgentSpanContext.class); - - assertSame( - extensionContext, LambdaAppSecHandler.mergeContexts(extensionContext, appSecContext)); - - verify(appSecData).close(); - } - - @Test - void mergeContextsDoesNotCloseTransferredAppSecData() throws IOException { - Closeable appSecData = mock(Closeable.class); - TagContext appSecContext = new TagContext(); - appSecContext.withRequestContextDataAppSec(appSecData); - TagContext extensionContext = new TagContext(); - - assertSame( - extensionContext, LambdaAppSecHandler.mergeContexts(extensionContext, appSecContext)); - - assertSame(appSecData, extensionContext.getRequestContextDataAppSec()); - verify(appSecData, never()).close(); - } - // ============================================================================ // Error Handling and Null Callback Tests // ============================================================================ @@ -1579,29 +1551,6 @@ void processRequestStartHandlesNullRequestStartedCallbackGracefully() { assertNull(LambdaAppSecHandler.processRequestStart(event)); } - @Test - @SuppressWarnings("unchecked") - void processRequestStartClosesAppSecDataWhenRequestProcessingFails() throws IOException { - String eventJson = "{\"path\": \"/test\", \"requestContext\": {\"httpMethod\": \"GET\"}}"; - Closeable appSecData = mock(Closeable.class); - Supplier> requestStartedCallback = mock(Supplier.class); - when(requestStartedCallback.get()).thenReturn(new Flow.ResultFlow<>(appSecData)); - TriFunction> methodUriCallback = - mock(TriFunction.class); - when(methodUriCallback.apply(any(), anyString(), any(URIDataAdapter.class))) - .thenThrow(new IllegalStateException("request processing failed")); - CallbackProvider callbackProvider = mock(CallbackProvider.class); - when(callbackProvider.getCallback(EVENTS.requestStarted())).thenReturn(requestStartedCallback); - when(callbackProvider.getCallback(EVENTS.requestMethodUriRaw())).thenReturn(methodUriCallback); - AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); - when(tracer.getCallbackProvider(RequestContextSlot.APPSEC)).thenReturn(callbackProvider); - AgentTracer.forceRegister(tracer); - - assertNull(LambdaAppSecHandler.processRequestStart(createInputStream(eventJson))); - - verify(appSecData).close(); - } - @Test @SuppressWarnings("unchecked") void processRequestStartHandlesNullMethodUriCallbackGracefully() { From b2223df740eea46a928dc43644c00f5577dca52a Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Mon, 21 Sep 2026 17:33:05 +0200 Subject: [PATCH 4/8] Restore Lambda AppSec activation gates --- .../trace/lambda/LambdaAppSecHandler.java | 6 ++- .../trace/lambda/LambdaAppSecHandlerTest.java | 54 +++---------------- 2 files changed, 11 insertions(+), 49 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java index fd7bd59bbe6..6e4b14b65a6 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java @@ -131,7 +131,7 @@ public static void processRequestEnd(AgentSpan span) { LambdaTriggerType triggerType = CURRENT_TRIGGER_TYPE.get(); CURRENT_TRIGGER_TYPE.remove(); - if (span == null) { + if (!ActiveSubsystems.APPSEC_ACTIVE || span == null) { return; } @@ -178,7 +178,9 @@ public static void processRequestEnd(AgentSpan span) { * @param result the Lambda handler result (expected to be a ByteArrayOutputStream) */ public static void processResponseData(AgentSpan span, Object result) { - if (span == null || !(result instanceof ByteArrayOutputStream)) { + if (!ActiveSubsystems.APPSEC_ACTIVE + || span == null + || !(result instanceof ByteArrayOutputStream)) { return; } diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java index c6bc868cbbf..12781193e1a 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java @@ -1218,14 +1218,13 @@ void processRequestEndDoesNothingWhenSpanIsNull() { } @Test - void processRequestEndDoesNothingWithoutAttachedContextWhenAppSecIsDisabled() { + void processRequestEndDoesNothingWhenAppSecIsDisabled() { ActiveSubsystems.APPSEC_ACTIVE = false; AgentSpan span = mock(AgentSpan.class); LambdaAppSecHandler.processRequestEnd(span); - verify(span).getRequestContext(); - verifyNoMoreInteractions(span); + verifyNoInteractions(span); } @Test @@ -1268,30 +1267,6 @@ void processRequestEndInvokesCallbackAndSkipsAsmTagsWhenNotManuallyKept() { verify(mockTraceSegment, never()).setTagTop(any(), any()); } - @Test - @SuppressWarnings("unchecked") - void processRequestEndFinalizesAttachedContextWhenAppSecBecomesInactive() { - ActiveSubsystems.APPSEC_ACTIVE = false; - Object appSecContext = new Object(); - RequestContext requestContext = mock(RequestContext.class); - when(requestContext.getData(RequestContextSlot.APPSEC)).thenReturn(appSecContext); - AgentSpan span = mock(AgentSpan.class); - when(span.getRequestContext()).thenReturn(requestContext); - - BiFunction> requestEndedCallback = - mock(BiFunction.class); - when(requestEndedCallback.apply(any(), any())).thenReturn(Flow.ResultFlow.empty()); - CallbackProvider callbackProvider = mock(CallbackProvider.class); - when(callbackProvider.getCallback(EVENTS.requestEnded())).thenReturn(requestEndedCallback); - AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); - when(tracer.getCallbackProvider(RequestContextSlot.APPSEC)).thenReturn(callbackProvider); - AgentTracer.forceRegister(tracer); - - LambdaAppSecHandler.processRequestEnd(span); - - verify(requestEndedCallback).apply(requestContext, span); - } - @Test @SuppressWarnings("unchecked") void processRequestEndFinalizesAttachedContextWithoutRecordedTrigger() { @@ -1390,19 +1365,6 @@ void processRequestEndSetsNoUnsupportedEventTypeMetricWhenNoTriggerTypeWasRecord verify(span, never()).setMetric(anyString(), anyInt()); } - @Test - void processRequestEndPreservesUnsupportedMetricWhenAppSecBecomesDisabled() { - ActiveSubsystems.APPSEC_ACTIVE = false; - LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.UNKNOWN); - AgentSpan span = mock(AgentSpan.class); - - LambdaAppSecHandler.processRequestEnd(span); - - verify(span).getRequestContext(); - verify(span).setMetric("_dd.appsec.unsupported_event_type", 1); - verifyNoMoreInteractions(span); - } - @Test @SuppressWarnings("unchecked") void processRequestEndSetsNoUnsupportedEventTypeMetricForHttpTrigger() { @@ -1752,17 +1714,15 @@ private RequestContext captureTemporaryRequestContext(Object appSecContext) { // ============================================================================ @Test - void processResponseDataProcessesAttachedContextWhenAppSecBecomesDisabled() { + void processResponseDataDoesNothingWhenAppSecIsDisabled() { ActiveSubsystems.APPSEC_ACTIVE = false; LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); - Integer[] capturedStatus = {null}; - AgentSpan span = - setupMockResponseCallbacks(status -> capturedStatus[0] = status, null, null, null); - ByteArrayOutputStream result = createOutputStream("{\"statusCode\": 200, \"body\": \"ok\"}"); + AgentSpan span = mock(AgentSpan.class); - LambdaAppSecHandler.processResponseData(span, result); + LambdaAppSecHandler.processResponseData( + span, createOutputStream("{\"statusCode\": 200, \"body\": \"ok\"}")); - assertEquals(200, capturedStatus[0]); + verifyNoInteractions(span); } @Test From ae5d81f7073e71481789c620cac0ff60bbee3529 Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Tue, 22 Sep 2026 14:28:52 +0200 Subject: [PATCH 5/8] Restore Lambda AppSec trigger gates --- .../trace/lambda/LambdaAppSecHandler.java | 14 +++++----- .../trace/lambda/LambdaAppSecHandlerTest.java | 27 +------------------ 2 files changed, 7 insertions(+), 34 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java index 6e4b14b65a6..5885a942e47 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java @@ -131,7 +131,12 @@ public static void processRequestEnd(AgentSpan span) { LambdaTriggerType triggerType = CURRENT_TRIGGER_TYPE.get(); CURRENT_TRIGGER_TYPE.remove(); - if (!ActiveSubsystems.APPSEC_ACTIVE || span == null) { + if (!ActiveSubsystems.APPSEC_ACTIVE || span == null || triggerType == null) { + return; + } + + if (!triggerType.isHttp()) { + span.setMetric(UNSUPPORTED_EVENT_TYPE_METRIC, 1); return; } @@ -159,13 +164,6 @@ public static void processRequestEnd(AgentSpan span) { traceSeg.setTagTop(Tags.ASM_KEEP, true); traceSeg.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); } - return; - } - - // A null trigger type means processRequestStart never ran, so the invocation was not analysed - // at all, which is not the same as an unsupported trigger. - if (triggerType != null && !triggerType.isHttp()) { - span.setMetric(UNSUPPORTED_EVENT_TYPE_METRIC, 1); } } diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java index 12781193e1a..ba292eff21a 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java @@ -1267,29 +1267,6 @@ void processRequestEndInvokesCallbackAndSkipsAsmTagsWhenNotManuallyKept() { verify(mockTraceSegment, never()).setTagTop(any(), any()); } - @Test - @SuppressWarnings("unchecked") - void processRequestEndFinalizesAttachedContextWithoutRecordedTrigger() { - Object appSecContext = new Object(); - RequestContext requestContext = mock(RequestContext.class); - when(requestContext.getData(RequestContextSlot.APPSEC)).thenReturn(appSecContext); - AgentSpan span = mock(AgentSpan.class); - when(span.getRequestContext()).thenReturn(requestContext); - - BiFunction> requestEndedCallback = - mock(BiFunction.class); - when(requestEndedCallback.apply(any(), any())).thenReturn(Flow.ResultFlow.empty()); - CallbackProvider callbackProvider = mock(CallbackProvider.class); - when(callbackProvider.getCallback(EVENTS.requestEnded())).thenReturn(requestEndedCallback); - AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); - when(tracer.getCallbackProvider(RequestContextSlot.APPSEC)).thenReturn(callbackProvider); - AgentTracer.forceRegister(tracer); - - LambdaAppSecHandler.processRequestEnd(span); - - verify(requestEndedCallback).apply(requestContext, span); - } - @Test void processRequestEndHandlesNullRequestEndedCallbackGracefully() { LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); @@ -1350,7 +1327,6 @@ void processRequestEndSetsUnsupportedEventTypeMetricForNonHttpTrigger() { LambdaAppSecHandler.processRequestEnd(span); - verify(span).getRequestContext(); verify(span).setMetric("_dd.appsec.unsupported_event_type", 1); verifyNoMoreInteractions(span); } @@ -1362,7 +1338,7 @@ void processRequestEndSetsNoUnsupportedEventTypeMetricWhenNoTriggerTypeWasRecord LambdaAppSecHandler.processRequestEnd(span); - verify(span, never()).setMetric(anyString(), anyInt()); + verifyNoInteractions(span); } @Test @@ -1437,7 +1413,6 @@ void processRequestEndSetsUnsupportedEventTypeMetricAfterAnUnparseablePayload() AgentSpan span = mock(AgentSpan.class); LambdaAppSecHandler.processRequestEnd(span); - verify(span).getRequestContext(); verify(span).setMetric("_dd.appsec.unsupported_event_type", 1); verifyNoMoreInteractions(span); } From 5ee22c258c44e8c1980f80244bf21feab48d395b Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Tue, 22 Sep 2026 15:53:00 +0200 Subject: [PATCH 6/8] fix tests --- .../gateway/GatewayBridgeSpecification.groovy | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index 7073ef3ecf9..33bc2c2e29c 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -35,6 +35,7 @@ import datadog.trace.api.telemetry.RuleType import datadog.trace.api.telemetry.WafMetricCollector import datadog.trace.bootstrap.instrumentation.api.AgentSpan import datadog.trace.bootstrap.instrumentation.api.AgentTracer +import datadog.trace.bootstrap.instrumentation.api.TagContext import datadog.trace.bootstrap.instrumentation.api.Tags import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase @@ -42,6 +43,7 @@ import datadog.trace.lambda.LambdaAppSecHandler import datadog.trace.test.util.DDSpecification import spock.lang.Shared +import java.nio.charset.StandardCharsets import java.util.function.BiConsumer import java.util.function.BiFunction import java.util.function.Function @@ -221,26 +223,35 @@ class GatewayBridgeSpecification extends DDSpecification { void 'lambda request end reaches shared waf telemetry with its framework'() { given: AgentTracer.TracerAPI originalTracer = AgentTracer.get() - CallbackProvider callbackProvider = Stub { + CallbackProvider callbackProvider = Mock { + getCallback(EVENTS.requestStarted()) >> requestStartedCB getCallback(EVENTS.requestEnded()) >> requestEndedCB } AgentTracer.TracerAPI tracer = Stub { getCallbackProvider(RequestContextSlot.APPSEC) >> callbackProvider } + AgentTracer.forceRegister(tracer) + + byte[] event = '{"path":"/","requestContext":{"httpMethod":"GET"}}'.getBytes(StandardCharsets.UTF_8) + TagContext lambdaContext = LambdaAppSecHandler.processRequestStart(new ByteArrayInputStream(event)) as TagContext + AppSecRequestContext lambdaAppSecContext = lambdaContext.requestContextDataAppSec as AppSecRequestContext + RequestContext lambdaRequestContext = Stub { + getData(RequestContextSlot.APPSEC) >> lambdaAppSecContext + getTraceSegment() >> traceSegment + } AgentSpan span = Mock { - getRequestContext() >> ctx - getTags() >> TagMap.fromMap([(Tags.COMPONENT): 'aws-lambda']) + getRequestContext() >> lambdaRequestContext + getTags() >> lambdaContext.tags } - AgentTracer.forceRegister(tracer) when: LambdaAppSecHandler.processRequestEnd(span) then: - 1 * requestSampler.preSampleRequest(arCtx, 'aws-lambda') >> false + 1 * requestSampler.preSampleRequest(lambdaAppSecContext, 'aws-lambda') >> false 1 * span.setMetric('_dd.appsec.enabled', 1) 1 * span.setTag('_dd.runtime_family', 'jvm') - 1 * pp.processTraceSegment(traceSegment, arCtx, []) + 1 * pp.processTraceSegment(traceSegment, lambdaAppSecContext, []) 1 * wafMetricCollector.wafRequest(false, false, false, false, false, false, false, false) cleanup: From 03452084b9f4bc6fdb548f1b368767a4a0bffe2b Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Wed, 23 Sep 2026 17:37:30 +0200 Subject: [PATCH 7/8] Preserve Lambda resource name on AppSec failure --- .../lambda/LambdaHandlerInstrumentation.java | 3 +- .../LambdaHandlerInstrumentationTest.java | 28 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index 73beb4a7314..c44190fbd13 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -126,7 +126,7 @@ static void exit( } AgentTracer.get().notifyAppSecEnd(span, throwable == null ? result : null); - + } finally { // Force the resource name back to the literal placeholder marker right // before finish so that the Datadog Lambda Extension's filter // (filter_span_from_lambda_library_or_runtime in @@ -142,7 +142,6 @@ static void exit( // and the HTTP/JAX-RS instrumentation will already have written // HTTP_FRAMEWORK_ROUTE (3) by this point. span.setResourceName(INVOCATION_SPAN_NAME, ResourceNamePriorities.TAG_INTERCEPTOR); - } finally { scope.close(); span.finish(); AgentTracer.get() diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java index 3e5e78b8056..0aff2114e51 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java @@ -26,6 +26,7 @@ import datadog.trace.api.gateway.RequestContextSlot; import datadog.trace.api.gateway.SubscriptionService; import datadog.trace.bootstrap.ActiveSubsystems; +import datadog.trace.bootstrap.InstrumentationErrors; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; @@ -187,6 +188,33 @@ void serverlessInvocationSpanResourceResetAfterHttpFrameworkOverwrite() throws I .error(false))); } + @Test + void serverlessInvocationSpanResourceResetWhenAppSecEndThrows() throws IOException { + String eventJson = + "{" + "\"path\": \"/\"," + "\"requestContext\": {\"httpMethod\": \"GET\"}" + "}"; + ByteArrayInputStream input = + new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8)); + ByteArrayOutputStream output = + new ByteArrayOutputStream() { + @Override + public synchronized byte[] toByteArray() { + throw new AssertionError("response processing failed"); + } + }; + + new HandlerStreamingSimulatesHttpFrameworkResource().handleRequest(input, output, newContext()); + + assertFalse(InstrumentationErrors.noErrors()); + InstrumentationErrors.resetErrors(); + assertTrue(appSecEnded); + assertTraces( + trace( + span() + .resourceName(name -> operation().equals(name.toString())) + .type(DDSpanTypes.SERVERLESS) + .error(false))); + } + @Test void testStreamingHandlerWithError() { ByteArrayInputStream input = new ByteArrayInputStream("Hello".getBytes(StandardCharsets.UTF_8)); From ae8ced69e4fd34beaed5d0d5b53f0b549adcabe8 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Thu, 24 Sep 2026 10:44:18 -0400 Subject: [PATCH 8/8] Keep http.status_code when the Lambda AppSec context is missing (#12623) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Keep http.status_code when the Lambda AppSec context is missing processResponseData moved its RequestContext guard above the parse and tightened it to also require APPSEC slot data. Because DDSpanContext.getRequestContext() returns `this` and is never null for a real span, the live half of that guard is the new APPSEC check — so an exception inside processRequestStart (caught and swallowed there, while CURRENT_TRIGGER_TYPE is already set to the HTTP trigger) now costs the invocation span its http.status_code and error flag as well. Split the guard instead: compute the AppSec-context flag up front, keep publishing the status code, and gate only the WAF response callbacks on it. This is the previous ordering with the stricter condition retained. Also renames processResponseDataDoesNothingWhenSpanHasNoRequestContext, which asserted nothing and whose name described behaviour the span never had, and adds coverage for the non-null-RequestContext-without-APPSEC case that the regression actually goes through. Co-Authored-By: Claude Opus 5 (1M context) * Restore comment explaining the null trigger type check The comment distinguishing a null trigger type (processRequestStart never ran) from a non-HTTP one (ran, trigger unsupported) was dropped while the requestContext guard below it was reworked. Both branches are still there and still behave differently, so the explanation still applies. Co-Authored-By: Claude Opus 5 (1M context) --------- Co-authored-by: Joey Zhao <234797088+joey-zhao_ddog@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- .../trace/lambda/LambdaAppSecHandler.java | 16 +++++++--- .../trace/lambda/LambdaAppSecHandlerTest.java | 32 ++++++++++++++++--- 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java index 5885a942e47..8ed1eb4cb6a 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java @@ -135,6 +135,8 @@ public static void processRequestEnd(AgentSpan span) { return; } + // A null trigger type means processRequestStart never ran, so the invocation was not analysed + // at all, which is not the same as an unsupported trigger. if (!triggerType.isHttp()) { span.setMetric(UNSUPPORTED_EVENT_TYPE_METRIC, 1); return; @@ -189,10 +191,8 @@ public static void processResponseData(AgentSpan span, Object result) { } RequestContext requestContext = span.getRequestContext(); - if (requestContext == null || requestContext.getData(RequestContextSlot.APPSEC) == null) { - log.debug("Span has no AppSec request context, skipping response processing"); - return; - } + boolean hasAppSecContext = + requestContext != null && requestContext.getData(RequestContextSlot.APPSEC) != null; try { byte[] bytes = ((ByteArrayOutputStream) result).toByteArray(); @@ -219,6 +219,14 @@ public static void processResponseData(AgentSpan span, Object result) { span.setError(isError, ErrorPriorities.HTTP_SERVER_DECORATOR); } + // http.status_code is a tracing tag and is published above whether or not AppSec ran: a + // failure inside processRequestStart must not also cost the span its status. The WAF + // callbacks below, in contrast, have nowhere to deliver without an AppSec request context. + if (!hasAppSecContext) { + log.debug("Span has no AppSec request context, skipping response WAF callbacks"); + return; + } + AgentTracer.TracerAPI tracer = AgentTracer.get(); CallbackProvider cbp = tracer.getCallbackProvider(RequestContextSlot.APPSEC); diff --git a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java index ba292eff21a..f51ec9e4e78 100644 --- a/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/lambda/LambdaAppSecHandlerTest.java @@ -1753,13 +1753,35 @@ void processResponseDataDoesNothingForNullResult() { } @Test - void processResponseDataDoesNothingWhenSpanHasNoRequestContext() { + void processResponseDataPublishesStatusButNoWafEventsWhenSpanHasNoRequestContext() { + LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); AgentSpan span = mock(AgentSpan.class); when(span.getRequestContext()).thenReturn(null); - ByteArrayOutputStream result = createOutputStream("{\"statusCode\": 200}"); - setupMockResponseCallbacks(null, null, null, null); - LambdaAppSecHandler.processResponseData(span, result); - // no exception expected + AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); + AgentTracer.forceRegister(tracer); + + LambdaAppSecHandler.processResponseData(span, createOutputStream("{\"statusCode\": 200}")); + + verify(span).setHttpStatusCode(200); + verify(tracer, never()).getCallbackProvider(RequestContextSlot.APPSEC); + } + + @Test + void processResponseDataStillPublishesStatusWhenSpanHasNoAppSecContext() { + // An exception inside processRequestStart leaves an HTTP trigger type recorded but no AppSec + // context on the span. http.status_code is a tracing tag and must survive that. + LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); + RequestContext requestContext = mock(RequestContext.class); + AgentSpan span = mock(AgentSpan.class); + when(span.getRequestContext()).thenReturn(requestContext); + AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); + AgentTracer.forceRegister(tracer); + + LambdaAppSecHandler.processResponseData( + span, createOutputStream("{\"statusCode\": 503, \"body\": \"boom\"}")); + + verify(span).setHttpStatusCode(503); + verify(tracer, never()).getCallbackProvider(RequestContextSlot.APPSEC); } @Test