diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/HttpServerResponseEndHandlerInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/HttpServerResponseEndHandlerInstrumentation.java index 5cfdf5767b2..da2d589e475 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/HttpServerResponseEndHandlerInstrumentation.java +++ b/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 }; } diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapper.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapper.java index c34fb32cdfb..70aaa687dde 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapper.java +++ b/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; } } diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy index 8a9372dd9e7..2be39db6f4a 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy +++ b/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 { @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) + } } diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-5.0/src/test/java/server/HttpServerResponseEndHandlerInstrumentationTest.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-5.0/src/test/java/server/HttpServerResponseEndHandlerInstrumentationTest.java new file mode 100644 index 00000000000..25201f3e872 --- /dev/null +++ b/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 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 ? "" : 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"); + } +}