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..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 @@ -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,12 +34,16 @@ 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.TagContext 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 +import java.nio.charset.StandardCharsets import java.util.function.BiConsumer import java.util.function.BiFunction import java.util.function.Function @@ -215,6 +220,44 @@ 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 = 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() >> lambdaRequestContext + getTags() >> lambdaContext.tags + } + + when: + LambdaAppSecHandler.processRequestEnd(span) + + then: + 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, lambdaAppSecContext, []) + 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..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 @@ -121,11 +121,12 @@ static void exit( final AgentSpan span = scope.span(); try { - if (throwable == null) { - AgentTracer.get().notifyAppSecEnd(span, result); - } else { + if (throwable != null) { span.addThrowable(throwable); } + + 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 @@ -141,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/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..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; @@ -59,6 +60,7 @@ abstract class LambdaHandlerInstrumentationTest extends AbstractInstrumentationT Map capturedHeaders; Object capturedBody; boolean appSecEnded; + int appSecEndCount; Integer capturedResponseStatus; Map capturedResponseHeaders; @@ -82,6 +84,7 @@ void setUpAppSec() { capturedHeaders = new HashMap<>(); capturedBody = null; appSecEnded = false; + appSecEndCount = 0; capturedResponseStatus = null; capturedResponseHeaders = new HashMap<>(); capturedResponseBody = null; @@ -121,6 +124,7 @@ void setUpAppSec() { (BiFunction>) (ctx2, spanInfo) -> { appSecEnded = true; + appSecEndCount++; return Flow.ResultFlow.empty(); }); @@ -184,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)); @@ -201,6 +232,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 +519,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 +604,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 +629,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..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 @@ -143,7 +143,9 @@ public static void processRequestEnd(AgentSpan span) { } 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 +159,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()) { @@ -189,6 +190,10 @@ public static void processResponseData(AgentSpan span, Object result) { return; } + RequestContext requestContext = span.getRequestContext(); + boolean hasAppSecContext = + requestContext != null && requestContext.getData(RequestContextSlot.APPSEC) != null; + try { byte[] bytes = ((ByteArrayOutputStream) result).toByteArray(); if (bytes.length == 0 || bytes.length > MAX_EVENT_SIZE) { @@ -214,9 +219,11 @@ 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"); + // 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; } @@ -319,6 +326,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); 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..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 @@ -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,6 +49,7 @@ 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; @@ -1269,6 +1271,7 @@ void processRequestEndInvokesCallbackAndSkipsAsmTagsWhenNotManuallyKept() { 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); @@ -1335,17 +1338,6 @@ void processRequestEndSetsNoUnsupportedEventTypeMetricWhenNoTriggerTypeWasRecord LambdaAppSecHandler.processRequestEnd(span); - verify(span, never()).setMetric(anyString(), anyInt()); - } - - @Test - void processRequestEndSetsNoUnsupportedEventTypeMetricWhenAppSecIsDisabled() { - ActiveSubsystems.APPSEC_ACTIVE = false; - LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.UNKNOWN); - AgentSpan span = mock(AgentSpan.class); - - LambdaAppSecHandler.processRequestEnd(span); - verifyNoInteractions(span); } @@ -1354,6 +1346,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 +1367,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"); @@ -1660,10 +1691,44 @@ private RequestContext captureTemporaryRequestContext(Object appSecContext) { @Test void processResponseDataDoesNothingWhenAppSecIsDisabled() { ActiveSubsystems.APPSEC_ACTIVE = false; + LambdaAppSecHandler.setCurrentTriggerType(LambdaTriggerType.API_GATEWAY_V1_REST); AgentSpan span = mock(AgentSpan.class); - ByteArrayOutputStream result = createOutputStream("{\"statusCode\": 200, \"body\": \"ok\"}"); - LambdaAppSecHandler.processResponseData(span, result); - verify(span, never()).getRequestContext(); + + LambdaAppSecHandler.processResponseData( + span, createOutputStream("{\"statusCode\": 200, \"body\": \"ok\"}")); + + verifyNoInteractions(span); + } + + @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 @@ -1688,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 @@ -2377,6 +2464,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 +2712,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 +2721,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 +2959,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);