Complete AWS Lambda AppSec telemetry integration

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 >> [

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()

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);
+  }
+}

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");
+  }
+}

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<String, String> capturedHeaders;
   Object capturedBody;
   boolean appSecEnded;
+  int appSecEndCount;
 
   Integer capturedResponseStatus;
   Map<String, String> 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<RequestContext, IGSpanInfo, Flow<Void>>)
             (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<String> 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"))));
   }
 

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

dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java

@@ -135,15 +135,15 @@ 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;
     }
 
     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<RequestContext, IGSpanInfo, Flow<Void>> requestEndedCallback =
           tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestEnded());
@@ -157,7 +157,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 +188,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 +219,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);
 
@@ -319,6 +318,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);