Fix missing instrumentation for vertx 5.1 responses

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/HttpServerResponseEndHandlerInstrumentation.java

@@ -34,7 +34,10 @@ public String[] helperClassNames() {
   @Override
   public String[] knownMatchingTypes() {
     return new String[] {
-      "io.vertx.core.http.impl.Http1xServerResponse", "io.vertx.core.http.impl.Http2ServerResponse "
+      "io.vertx.core.http.impl.Http1xServerResponse",
+      "io.vertx.core.http.impl.Http2ServerResponse",
+      "io.vertx.core.http.impl.http1.Http1ServerResponse", // HTTP/1 response when v >= 5.1
+      "io.vertx.core.http.impl.HttpServerResponseImpl" // HTTP/2 response when v >= 5.1
     };
   }
 

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapper.java

@@ -8,6 +8,7 @@
 import static datadog.trace.instrumentation.vertx_4_0.server.VertxDecorator.DECORATE;
 import static datadog.trace.instrumentation.vertx_4_0.server.VertxDecorator.INSTRUMENTATION_NAME;
 
+import datadog.appsec.api.blocking.BlockingException;
 import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
@@ -65,6 +66,12 @@ public void handle(final RoutingContext routingContext) {
         actual.handle(routingContext);
       } catch (final Throwable t) {
         DECORATE.onError(span, t);
+        if (t instanceof BlockingException) {
+          // AppSec uses BlockingException as control flow after committing a blocking response
+          // from advice such as WafPublishingBodyHandler and RoutingContextJsonAdvice. Finish
+          // immediately because that abort path may bypass Vert.x response/routing end callbacks.
+          finishHandlerSpan(routingContext);
+        }
         throw t;
       }
     }

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy

@@ -1,5 +1,6 @@
 package server
 
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_JSON
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.LOGIN
@@ -16,6 +17,8 @@ import datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator
 import datadog.trace.instrumentation.vertx_4_0.server.VertxDecorator
 import io.vertx.core.AbstractVerticle
 import io.vertx.core.Vertx
+import okhttp3.MediaType
+import okhttp3.RequestBody
 
 class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
   @Override
@@ -181,4 +184,28 @@ class VertxHttpServerWorkerForkedTest extends VertxHttpServerForkedTest {
   HttpServer server() {
     return new VertxServer(verticle(), routerBasePath(), true)
   }
+
+  def 'test blocking of JSON request body finishes route handler span'() {
+    setup:
+    // VertxTestServer handles BODY_JSON by calling ctx.body().asJsonObject().
+    // The IG_BODY_CONVERTED_HEADER is consumed by HttpServerTest's AppSec test callback, which
+    // returns a RequestBlockingAction from requestBodyProcessed() when that JSON body is converted.
+    def request = request(
+      BODY_JSON, 'POST',
+      RequestBody.create(MediaType.get('application/json'), '{"a": "x"}'))
+      .header(IG_BODY_CONVERTED_HEADER, 'true')
+      .build()
+
+    when:
+    def response = client.newCall(request).execute()
+
+    then:
+    response.code() == 413
+    response.body().charStream().text.contains('"title":"You\'ve been blocked"')
+    !handlerRan
+    // The client receiving a 413 only proves the blocking response was committed.
+    // We want to make sure that a BlockingException does now abort the worker route handler
+    // before the vertx.route-handler span has been finished (which would leave it dangling)
+    TEST_WRITER.waitForTraces(1)
+  }
 }

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-5.0/src/test/java/server/HttpServerResponseEndHandlerInstrumentationTest.java

@@ -0,0 +1,179 @@
+package server;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
+
+import datadog.trace.agent.test.AbstractInstrumentationTest;
+import io.vertx.core.Handler;
+import io.vertx.core.Vertx;
+import io.vertx.core.buffer.Buffer;
+import io.vertx.core.http.HttpClient;
+import io.vertx.core.http.HttpClientOptions;
+import io.vertx.core.http.HttpClientResponse;
+import io.vertx.core.http.HttpMethod;
+import io.vertx.core.http.HttpServer;
+import io.vertx.core.http.HttpServerOptions;
+import io.vertx.core.http.HttpServerResponse;
+import io.vertx.core.http.HttpVersion;
+import io.vertx.ext.web.Router;
+import java.io.InputStream;
+import java.lang.reflect.Field;
+import java.net.HttpURLConnection;
+import java.net.ServerSocket;
+import java.net.URL;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+class HttpServerResponseEndHandlerInstrumentationTest extends AbstractInstrumentationTest {
+  private static final String HTTP1_RESPONSE = "io.vertx.core.http.impl.http1.Http1ServerResponse";
+  private static final String HTTP2_RESPONSE = "io.vertx.core.http.impl.HttpServerResponseImpl";
+  private static final String END_HANDLER_WRAPPER =
+      "datadog.trace.instrumentation.vertx_4_0.server.EndHandlerWrapper";
+
+  private static Vertx vertx;
+  private static HttpServer server;
+  private static int port;
+
+  @BeforeAll
+  static void startServer() throws Exception {
+    try (ServerSocket socket = new ServerSocket(0)) {
+      port = socket.getLocalPort();
+    }
+
+    vertx = Vertx.vertx();
+    Router router = Router.router(vertx);
+    router
+        .route("/end-handler-wrapper")
+        .handler(HttpServerResponseEndHandlerInstrumentationTest::handle);
+
+    CountDownLatch ready = new CountDownLatch(1);
+    vertx
+        .createHttpServer(new HttpServerOptions().setHttp2ClearTextEnabled(true))
+        .requestHandler(router)
+        .listen(port)
+        .andThen(
+            result -> {
+              if (result.failed()) {
+                throw new RuntimeException("Failed to start Vert.x server", result.cause());
+              }
+              server = result.result();
+              ready.countDown();
+            });
+    if (!ready.await(10, TimeUnit.SECONDS)) {
+      throw new IllegalStateException("Vert.x server did not start in time");
+    }
+  }
+
+  @AfterAll
+  static void stopServer() throws Exception {
+    if (server != null) {
+      CountDownLatch closed = new CountDownLatch(1);
+      server.close().andThen(ar -> closed.countDown());
+      closed.await(10, TimeUnit.SECONDS);
+    }
+    if (vertx != null) {
+      CountDownLatch closed = new CountDownLatch(1);
+      vertx.close().andThen(ar -> closed.countDown());
+      closed.await(10, TimeUnit.SECONDS);
+    }
+  }
+
+  @Test
+  void wrapsApplicationEndHandlerOnVertx51Http1Response() throws Exception {
+    assumeTrue(hasVertx51Http1Response(), "Http1ServerResponse exists only in Vert.x 5.1+");
+
+    HttpURLConnection conn =
+        (HttpURLConnection)
+            new URL("http://localhost:" + port + "/end-handler-wrapper").openConnection();
+    conn.setRequestMethod("GET");
+    conn.setConnectTimeout(5000);
+    conn.setReadTimeout(5000);
+
+    int responseCode = conn.getResponseCode();
+    InputStream responseStream =
+        responseCode >= 400 ? conn.getErrorStream() : conn.getInputStream();
+    String body = new String(responseStream.readAllBytes(), StandardCharsets.UTF_8);
+    conn.disconnect();
+
+    assertWrapped(responseCode, body, HTTP1_RESPONSE);
+  }
+
+  @Test
+  void wrapsApplicationEndHandlerOnVertx51Http2Response() throws Exception {
+    assumeTrue(hasVertx51Http2Response(), "HttpServerResponseImpl exists only in Vert.x 5.1+");
+
+    HttpClient client =
+        vertx.createHttpClient(
+            new HttpClientOptions()
+                .setProtocolVersion(HttpVersion.HTTP_2)
+                .setHttp2ClearTextUpgrade(false));
+    try {
+      HttpClientResponse response =
+          client
+              .request(HttpMethod.GET, port, "localhost", "/end-handler-wrapper")
+              .compose(request -> request.send())
+              .toCompletionStage()
+              .toCompletableFuture()
+              .get(10, TimeUnit.SECONDS);
+      Buffer body =
+          response.body().toCompletionStage().toCompletableFuture().get(10, TimeUnit.SECONDS);
+
+      assertWrapped(response.statusCode(), body.toString(StandardCharsets.UTF_8), HTTP2_RESPONSE);
+    } finally {
+      client.close().toCompletionStage().toCompletableFuture().get(10, TimeUnit.SECONDS);
+    }
+  }
+
+  private static boolean hasVertx51Http1Response() {
+    return hasClass(HTTP1_RESPONSE);
+  }
+
+  private static boolean hasVertx51Http2Response() {
+    return hasClass(HTTP2_RESPONSE);
+  }
+
+  private static boolean hasClass(String className) {
+    try {
+      Class.forName(className);
+      return true;
+    } catch (ClassNotFoundException ignored) {
+      return false;
+    }
+  }
+
+  private static void assertWrapped(int responseCode, String body, String expectedResponseClass) {
+    assertEquals(200, responseCode, body);
+    assertTrue(body.contains("response=" + expectedResponseClass), body);
+    assertTrue(body.contains("endHandler=" + END_HANDLER_WRAPPER), body);
+  }
+
+  private static void handle(io.vertx.ext.web.RoutingContext ctx) {
+    HttpServerResponse response = ctx.response();
+    Handler<Void> applicationEndHandler = ignored -> {};
+    response.endHandler(applicationEndHandler);
+
+    String responseClassName = response.getClass().getName();
+    Object installedEndHandler;
+    try {
+      Field endHandler = response.getClass().getDeclaredField("endHandler");
+      endHandler.setAccessible(true);
+      installedEndHandler = endHandler.get(response);
+    } catch (ReflectiveOperationException e) {
+      throw new AssertionError("Could not inspect Vert.x response endHandler field", e);
+    }
+
+    String installedClassName =
+        installedEndHandler == null ? "<null>" : installedEndHandler.getClass().getName();
+    boolean wrapped =
+        (HTTP1_RESPONSE.equals(responseClassName) || HTTP2_RESPONSE.equals(responseClassName))
+            && END_HANDLER_WRAPPER.equals(installedClassName);
+    response
+        .setStatusCode(wrapped ? 200 : 500)
+        .end("response=" + responseClassName + "\nendHandler=" + installedClassName + "\n");
+  }
+}