From bd96ea3c17dcd580556ca97702684739fc38d39f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Tue, 24 Jun 2025 10:16:20 +0200 Subject: [PATCH] Extract Vert.x json body response schemas --- .../event/data/ObjectIntrospection.java | 45 ++++++++--- .../appsec/gateway/AppSecRequestContext.java | 9 +++ .../datadog/appsec/gateway/GatewayBridge.java | 36 +++++++++ .../ObjectIntrospectionSpecification.groovy | 28 ++++++- .../gateway/GatewayBridgeSpecification.groovy | 15 ++++ .../server/RoutingContextInstrumentation.java | 41 ++++++++++ .../RoutingContextJsonResponseAdvice.java | 55 ++++++++++++++ .../server/VertxHttpServerForkedTest.groovy | 5 ++ .../src/test/java/server/VertxTestServer.java | 3 +- .../server/VertxHttpServerForkedTest.groovy | 5 ++ .../src/test/java/server/VertxTestServer.java | 3 +- .../agent/test/base/HttpServerTest.groovy | 76 +++++++++++++++++++ .../java/datadog/vertx_4_2/MainVerticle.java | 10 +++ .../test/groovy/AppSecVertxSmokeTest.groovy | 38 ++++++++++ .../datadog/trace/api/gateway/Events.java | 12 +++ .../api/gateway/InstrumentationGateway.java | 2 + .../gateway/InstrumentationGatewayTest.java | 6 ++ 17 files changed, 373 insertions(+), 16 deletions(-) create mode 100644 dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextInstrumentation.java create mode 100644 dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextJsonResponseAdvice.java diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/ObjectIntrospection.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/ObjectIntrospection.java index 241e7e2965f..a024081860d 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/ObjectIntrospection.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/ObjectIntrospection.java @@ -11,6 +11,7 @@ import java.lang.reflect.Method; import java.lang.reflect.Modifier; import java.util.ArrayList; +import java.util.Collection; import java.util.HashMap; import java.util.Iterator; import java.util.List; @@ -211,20 +212,42 @@ private static Object doConversion(Object obj, int depth, State state) { // iterables if (obj instanceof Iterable) { - List newList; - if (obj instanceof List) { - newList = new ArrayList<>(((List) obj).size()); + final Iterator it = ((Iterable) obj).iterator(); + final boolean isMap = it.hasNext() && it.next() instanceof Map.Entry; + // some json libraries implement objects as Iterable + if (isMap) { + Map newMap; + if (obj instanceof Collection) { + newMap = new HashMap<>(((Collection) obj).size()); + } else { + newMap = new HashMap<>(); + } + for (Map.Entry e : ((Iterable>) obj)) { + Object key = e.getKey(); + Object newKey = keyConversion(e.getKey(), state); + if (newKey == null && key != null) { + // probably we're out of elements anyway + continue; + } + newMap.put(newKey, guardedConversion(e.getValue(), depth + 1, state)); + } + return newMap; } else { - newList = new ArrayList<>(); - } - for (Object o : ((Iterable) obj)) { - if (state.elemsLeft <= 0) { - state.listMapTooLarge = true; - break; + List newList; + if (obj instanceof Collection) { + newList = new ArrayList<>(((Collection) obj).size()); + } else { + newList = new ArrayList<>(); } - newList.add(guardedConversion(o, depth + 1, state)); + for (Object o : ((Iterable) obj)) { + if (state.elemsLeft <= 0) { + state.listMapTooLarge = true; + break; + } + newList.add(guardedConversion(o, depth + 1, state)); + } + return newList; } - return newList; } // arrays diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index 741f03c17ac..376b0448591 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java @@ -105,6 +105,7 @@ public class AppSecRequestContext implements DataBundle, Closeable { private boolean reqDataPublished; private boolean rawReqBodyPublished; private boolean convertedReqBodyPublished; + private boolean responseBodyPublished; private boolean respDataPublished; private boolean pathParamsPublished; private volatile Map derivatives; @@ -502,6 +503,14 @@ public void setConvertedReqBodyPublished(boolean convertedReqBodyPublished) { this.convertedReqBodyPublished = convertedReqBodyPublished; } + public boolean isResponseBodyPublished() { + return responseBodyPublished; + } + + public void setResponseBodyPublished(final boolean responseBodyPublished) { + this.responseBodyPublished = responseBodyPublished; + } + public boolean isRespDataPublished() { return respDataPublished; } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index 743c93c0728..c055f709a74 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -98,6 +98,7 @@ public class GatewayBridge { private volatile DataSubscriberInfo initialReqDataSubInfo; private volatile DataSubscriberInfo rawRequestBodySubInfo; private volatile DataSubscriberInfo requestBodySubInfo; + private volatile DataSubscriberInfo responseBodySubInfo; private volatile DataSubscriberInfo pathParamsSubInfo; private volatile DataSubscriberInfo respDataSubInfo; private volatile DataSubscriberInfo grpcServerMethodSubInfo; @@ -137,6 +138,7 @@ public void init() { subscriptionService.registerCallback(EVENTS.requestMethodUriRaw(), this::onRequestMethodUriRaw); subscriptionService.registerCallback(EVENTS.requestBodyStart(), this::onRequestBodyStart); subscriptionService.registerCallback(EVENTS.requestBodyDone(), this::onRequestBodyDone); + subscriptionService.registerCallback(EVENTS.responseBody(), this::onResponseBody); subscriptionService.registerCallback( EVENTS.requestClientSocketAddress(), this::onRequestClientSocketAddress); subscriptionService.registerCallback( @@ -177,6 +179,7 @@ public void reset() { initialReqDataSubInfo = null; rawRequestBodySubInfo = null; requestBodySubInfo = null; + responseBodySubInfo = null; pathParamsSubInfo = null; respDataSubInfo = null; grpcServerMethodSubInfo = null; @@ -638,6 +641,39 @@ private Flow onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup } } + private Flow onResponseBody(RequestContext ctx_, Object obj) { + AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); + if (ctx == null) { + return NoopFlow.INSTANCE; + } + + if (ctx.isResponseBodyPublished()) { + log.debug( + "Response body already published; will ignore new value of type {}", obj.getClass()); + return NoopFlow.INSTANCE; + } + ctx.setResponseBodyPublished(true); + + while (true) { + DataSubscriberInfo subInfo = responseBodySubInfo; + if (subInfo == null) { + subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT); + responseBodySubInfo = subInfo; + } + if (subInfo == null || subInfo.isEmpty()) { + return NoopFlow.INSTANCE; + } + Object converted = ObjectIntrospection.convert(obj, ctx); + DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.RESPONSE_BODY_OBJECT, converted); + try { + GatewayContext gwCtx = new GatewayContext(false); + return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx); + } catch (ExpiredSubscriberInfoException e) { + responseBodySubInfo = null; + } + } + } + private Flow onRequestPathParams(RequestContext ctx_, Map data) { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null || ctx.isPathParamsPublished()) { diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/ObjectIntrospectionSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/ObjectIntrospectionSpecification.groovy index e2e08ecb498..a5f4553eb5d 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/ObjectIntrospectionSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/ObjectIntrospectionSpecification.groovy @@ -2,13 +2,10 @@ package com.datadog.appsec.event.data import com.datadog.appsec.gateway.AppSecRequestContext import com.fasterxml.jackson.databind.ObjectMapper -import com.fasterxml.jackson.databind.node.ArrayNode -import com.fasterxml.jackson.databind.node.ObjectNode import datadog.trace.api.telemetry.WafMetricCollector import datadog.trace.test.util.DDSpecification import groovy.json.JsonBuilder import groovy.json.JsonOutput -import groovy.json.JsonSlurper import spock.lang.Shared import java.nio.CharBuffer @@ -465,6 +462,17 @@ class ObjectIntrospectionSpecification extends DDSpecification { MAPPER.readTree('"unicode: \\u0041"') || 'unicode: A' } + void 'iterable json objects'() { + setup: + final map = [name: 'This is just a test', list: [1, 2, 3, 4, 5]] + + when: + final result = convert(new IterableJsonObject(map), ctx) + + then: + result == map + } + private static int countNesting(final Mapobject, final int levels) { if (object.isEmpty()) { return levels @@ -475,4 +483,18 @@ class ObjectIntrospectionSpecification extends DDSpecification { } return countNesting(object.values().first() as Map, levels + 1) } + + private static class IterableJsonObject implements Iterable> { + + private final Map map + + IterableJsonObject(Map map) { + this.map = map + } + + @Override + Iterator> iterator() { + return map.entrySet().iterator() + } + } } 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 18c34db6d37..38eaf9f1208 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 @@ -99,6 +99,7 @@ class GatewayBridgeSpecification extends DDSpecification { BiFunction requestBodyStartCB BiFunction> requestBodyDoneCB BiFunction> requestBodyProcessedCB + BiFunction> responseBodyCB BiFunction> responseStartedCB TriConsumer respHeaderCB Function> respHeadersDoneCB @@ -463,6 +464,7 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * ig.registerCallback(EVENTS.requestBodyStart(), _) >> { requestBodyStartCB = it[1]; null } 1 * ig.registerCallback(EVENTS.requestBodyDone(), _) >> { requestBodyDoneCB = it[1]; null } 1 * ig.registerCallback(EVENTS.requestBodyProcessed(), _) >> { requestBodyProcessedCB = it[1]; null } + 1 * ig.registerCallback(EVENTS.responseBody(), _) >> { responseBodyCB = it[1]; null } 1 * ig.registerCallback(EVENTS.responseStarted(), _) >> { responseStartedCB = it[1]; null } 1 * ig.registerCallback(EVENTS.responseHeader(), _) >> { respHeaderCB = it[1]; null } 1 * ig.registerCallback(EVENTS.responseHeaderDone(), _) >> { respHeadersDoneCB = it[1]; null } @@ -1340,4 +1342,17 @@ class GatewayBridgeSpecification extends DDSpecification { arCtx.getRoute() == route } + void 'test on response body callback'() { + when: + responseBodyCB.apply(ctx, [test: 'this is a test']) + + then: + 1 * eventDispatcher.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT) >> nonEmptyDsInfo + 1 * eventDispatcher.publishDataEvent(_, _, _, _) >> { + final bundle = it[2] as DataBundle + final body = bundle.get(KnownAddresses.RESPONSE_BODY_OBJECT) + assert body['test'] == 'this is a test' + } + } + } diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextInstrumentation.java b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextInstrumentation.java new file mode 100644 index 00000000000..86908831f53 --- /dev/null +++ b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextInstrumentation.java @@ -0,0 +1,41 @@ +package datadog.trace.instrumentation.vertx_4_0.server; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.muzzle.Reference; +import io.vertx.ext.web.impl.RoutingContextImpl; + +/** + * @see RoutingContextImpl#getBodyAsJson(int) + * @see RoutingContextImpl#getBodyAsJsonArray(int) + */ +@AutoService(InstrumenterModule.class) +public class RoutingContextInstrumentation extends InstrumenterModule.AppSec + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + + public RoutingContextInstrumentation() { + super("vertx", "vertx-4.0"); + } + + @Override + public Reference[] additionalMuzzleReferences() { + return new Reference[] {VertxVersionMatcher.HTTP_1X_SERVER_RESPONSE}; + } + + @Override + public String instrumentedType() { + return "io.vertx.ext.web.RoutingContext"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("json").and(takesArguments(1)).and(takesArgument(0, Object.class)), + packageName + ".RoutingContextJsonResponseAdvice"); + } +} diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextJsonResponseAdvice.java b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextJsonResponseAdvice.java new file mode 100644 index 00000000000..0fbb7cd8990 --- /dev/null +++ b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextJsonResponseAdvice.java @@ -0,0 +1,55 @@ +package datadog.trace.instrumentation.vertx_4_0.server; + +import static datadog.trace.api.gateway.Events.EVENTS; + +import datadog.appsec.api.blocking.BlockingException; +import datadog.trace.advice.ActiveRequestContext; +import datadog.trace.advice.RequiresRequestContext; +import datadog.trace.api.gateway.BlockResponseFunction; +import datadog.trace.api.gateway.CallbackProvider; +import datadog.trace.api.gateway.Flow; +import datadog.trace.api.gateway.RequestContext; +import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import java.util.function.BiFunction; +import net.bytebuddy.asm.Advice; + +@RequiresRequestContext(RequestContextSlot.APPSEC) +class RoutingContextJsonResponseAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + static void before( + @Advice.Argument(0) final Object object, @ActiveRequestContext final RequestContext reqCtx) { + + if (object == null) { + return; + } + + CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); + if (cbp == null) { + return; + } + BiFunction> callback = + cbp.getCallback(EVENTS.responseBody()); + if (callback == null) { + return; + } + + Flow flow = callback.apply(reqCtx, object); + Flow.Action action = flow.getAction(); + if (action instanceof Flow.Action.RequestBlockingAction) { + BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction(); + if (blockResponseFunction == null) { + return; + } + Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; + blockResponseFunction.tryCommitBlockingResponse( + reqCtx.getTraceSegment(), + rba.getStatusCode(), + rba.getBlockingContentType(), + rba.getExtraHeaders()); + + throw new BlockingException("Blocked request (for RoutingContext/json)"); + } + } +} diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy b/dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy index db268592aef..e18172999ff 100644 --- a/dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy +++ b/dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy @@ -83,6 +83,11 @@ class VertxHttpServerForkedTest extends HttpServerTest { true } + @Override + boolean testResponseBodyJson() { + true + } + @Override boolean testBlocking() { true diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/test/java/server/VertxTestServer.java b/dd-java-agent/instrumentation/vertx-web-4.0/src/test/java/server/VertxTestServer.java index 350f945b625..d4dc482cb2a 100644 --- a/dd-java-agent/instrumentation/vertx-web-4.0/src/test/java/server/VertxTestServer.java +++ b/dd-java-agent/instrumentation/vertx-web-4.0/src/test/java/server/VertxTestServer.java @@ -127,7 +127,8 @@ public void start(final Promise startPromise) { BODY_JSON, () -> { JsonObject json = ctx.getBodyAsJson(); - ctx.response().setStatusCode(BODY_JSON.getStatus()).end(json.toString()); + ctx.response().setStatusCode(BODY_JSON.getStatus()); + ctx.json(json); })); router .route(QUERY_ENCODED_BOTH.getRawPath()) diff --git a/dd-java-agent/instrumentation/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy b/dd-java-agent/instrumentation/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy index c1111f61ecc..261430c0023 100644 --- a/dd-java-agent/instrumentation/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy +++ b/dd-java-agent/instrumentation/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy @@ -67,6 +67,11 @@ class VertxHttpServerForkedTest extends HttpServerTest { true } + @Override + boolean testResponseBodyJson() { + true + } + @Override boolean testBodyUrlencoded() { true diff --git a/dd-java-agent/instrumentation/vertx-web-5.0/src/test/java/server/VertxTestServer.java b/dd-java-agent/instrumentation/vertx-web-5.0/src/test/java/server/VertxTestServer.java index f711678b3bd..17feffcc79f 100644 --- a/dd-java-agent/instrumentation/vertx-web-5.0/src/test/java/server/VertxTestServer.java +++ b/dd-java-agent/instrumentation/vertx-web-5.0/src/test/java/server/VertxTestServer.java @@ -118,7 +118,8 @@ public void start(final Promise startPromise) { BODY_JSON, () -> { JsonObject json = ctx.body().asJsonObject(); - ctx.response().setStatusCode(BODY_JSON.getStatus()).end(json.toString()); + ctx.response().setStatusCode(BODY_JSON.getStatus()); + ctx.json(json); })); router .route(QUERY_ENCODED_BOTH.getRawPath()) diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy index c9278a71520..0bb81fd1102 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy @@ -40,6 +40,8 @@ import datadog.trace.bootstrap.instrumentation.api.URIUtils import datadog.trace.core.DDSpan import datadog.trace.core.datastreams.StatsGroup import datadog.trace.test.util.Flaky +import groovy.json.JsonOutput +import groovy.json.JsonSlurper import groovy.transform.Canonical import groovy.transform.CompileStatic import net.bytebuddy.utility.RandomString @@ -135,6 +137,7 @@ abstract class HttpServerTest extends WithHttpServer { ss.registerCallback(events.requestBodyStart(), callbacks.requestBodyStartCb) ss.registerCallback(events.requestBodyDone(), callbacks.requestBodyEndCb) ss.registerCallback(events.requestBodyProcessed(), callbacks.requestBodyObjectCb) + ss.registerCallback(events.responseBody(), callbacks.responseBodyObjectCb) ss.registerCallback(events.responseStarted(), callbacks.responseStartedCb) ss.registerCallback(events.responseHeader(), callbacks.responseHeaderCb) ss.registerCallback(events.responseHeaderDone(), callbacks.responseHeaderDoneCb) @@ -335,6 +338,7 @@ abstract class HttpServerTest extends WithHttpServer { false } + boolean isRequestBodyNoStreaming() { // if true, plain text request body tests expect the requestBodyProcessed // callback to tbe called, not requestBodyStart/requestBodyDone @@ -353,6 +357,10 @@ abstract class HttpServerTest extends WithHttpServer { false } + boolean testResponseBodyJson() { + false + } + boolean testBlocking() { false } @@ -1581,6 +1589,44 @@ abstract class HttpServerTest extends WithHttpServer { true | 'text/html;q=0.8, application/json;q=0.9' } + void 'test instrumentation gateway json response body'() { + setup: + assumeTrue(testResponseBodyJson()) + final body = [a: 'x'] + def request = request( + BODY_JSON, 'POST', + RequestBody.create(MediaType.get('application/json'), JsonOutput.toJson(body))) + .header(IG_RESPONSE_BODY_TAG, 'true') + .build() + def response = client.newCall(request).execute() + if (isDataStreamsEnabled()) { + TEST_DATA_STREAMS_WRITER.waitForGroups(1) + } + + expect: + response.body().charStream().text == BODY_JSON.body + + when: + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + + then: + !trace.isEmpty() + def rootSpan = trace.find { it.parentId == 0 } + assert rootSpan != null + final responseBody = rootSpan.getTag('response.body') as String + new JsonSlurper().parseText(responseBody) == body + + and: + if (isDataStreamsEnabled()) { + StatsGroup first = TEST_DATA_STREAMS_WRITER.groups.find { it.parentHash == 0 } + verifyAll(first) { + edgeTags.containsAll(DSM_EDGE_TAGS) + edgeTags.size() == DSM_EDGE_TAGS.size() + } + } + } + @Flaky(value = "https://github.com/DataDog/dd-trace-java/issues/4681", suites = ["GrizzlyAsyncTest", "GrizzlyTest"]) def 'test blocking of request with json response'() { setup: @@ -2280,6 +2326,7 @@ abstract class HttpServerTest extends WithHttpServer { static final String IG_BODY_END_BLOCK_HEADER = "x-block-body-end" static final String IG_BODY_CONVERTED_HEADER = "x-block-body-converted" static final String IG_ASK_FOR_RESPONSE_HEADER_TAGS_HEADER = "x-include-response-headers-in-tags" + static final String IG_RESPONSE_BODY_TAG = "x-include-response-body-in-tags" static final String IG_PEER_ADDRESS = "ig-peer-address" static final String IG_PEER_PORT = "ig-peer-port" static final String IG_RESPONSE_STATUS = "ig-response-status" @@ -2303,6 +2350,7 @@ abstract class HttpServerTest extends WithHttpServer { boolean bodyEndBlock boolean bodyConvertedBlock boolean responseHeadersInTags + boolean responseBodyTag } static final String stringOrEmpty(String string) { @@ -2356,6 +2404,9 @@ abstract class HttpServerTest extends WithHttpServer { if (IG_ASK_FOR_RESPONSE_HEADER_TAGS_HEADER.equalsIgnoreCase(key)) { context.responseHeadersInTags = true } + if (IG_RESPONSE_BODY_TAG.equalsIgnoreCase(key)) { + context.responseBodyTag = true + } } as TriConsumer final Function> requestHeaderDoneCb = @@ -2450,6 +2501,31 @@ abstract class HttpServerTest extends WithHttpServer { } } as BiFunction>) + final BiFunction> responseBodyObjectCb = + ({ RequestContext rqCtxt, Object obj -> + String body + // we need to extract a JSON representation of the response object, some frameworks classes might need updating + // as they might not work with a simple toString() call + if (obj instanceof String) { + body = obj as String + } else if (obj instanceof Map | obj instanceof List) { + body = JsonOutput.toJson(obj) + } else { + body = obj.toString() + } + Context context = rqCtxt.getData(RequestContextSlot.APPSEC) + if (context.responseBodyTag) { + rqCtxt.traceSegment.setTagTop('response.body', body) + } + if (context.responseBlock) { + new RbaFlow( + new Flow.Action.RequestBlockingAction(413, BlockingContentType.JSON) + ) + } else { + Flow.ResultFlow.empty() + } + } as BiFunction>) + final BiFunction> responseStartedCb = ({ RequestContext rqCtxt, Integer resultCode -> Context context = rqCtxt.getData(RequestContextSlot.APPSEC) diff --git a/dd-smoke-tests/vertx-4.2/application/src/main/java/datadog/vertx_4_2/MainVerticle.java b/dd-smoke-tests/vertx-4.2/application/src/main/java/datadog/vertx_4_2/MainVerticle.java index 64f795433bf..391a743f08f 100644 --- a/dd-smoke-tests/vertx-4.2/application/src/main/java/datadog/vertx_4_2/MainVerticle.java +++ b/dd-smoke-tests/vertx-4.2/application/src/main/java/datadog/vertx_4_2/MainVerticle.java @@ -6,6 +6,7 @@ import io.vertx.core.VertxOptions; import io.vertx.core.http.HttpServerOptions; import io.vertx.ext.web.Router; +import io.vertx.ext.web.handler.BodyHandler; import java.math.BigInteger; import java.util.concurrent.ThreadLocalRandom; @@ -51,6 +52,15 @@ public void start(Promise startPromise) throws Exception { .setStatusCode(Integer.parseInt(ctx.request().getParam("status_code"))) .end("EXECUTED")); + router.route("/api_security/response").handler(BodyHandler.create()); + router + .route("/api_security/response") + .handler( + ctx -> { + ctx.response().setStatusCode(200); + ctx.json(ctx.getBodyAsJson()); + }); + vertx .createHttpServer(new HttpServerOptions().setHandle100ContinueAutomatically(true)) .requestHandler( diff --git a/dd-smoke-tests/vertx-4.2/src/test/groovy/AppSecVertxSmokeTest.groovy b/dd-smoke-tests/vertx-4.2/src/test/groovy/AppSecVertxSmokeTest.groovy index 1ccf20d6fd1..6132e884c45 100644 --- a/dd-smoke-tests/vertx-4.2/src/test/groovy/AppSecVertxSmokeTest.groovy +++ b/dd-smoke-tests/vertx-4.2/src/test/groovy/AppSecVertxSmokeTest.groovy @@ -1,9 +1,15 @@ import datadog.smoketest.appsec.AbstractAppSecServerSmokeTest import datadog.trace.agent.test.utils.OkHttpUtils +import groovy.json.JsonOutput +import groovy.json.JsonSlurper +import okhttp3.MediaType import okhttp3.Request +import okhttp3.RequestBody import okhttp3.Response import spock.lang.IgnoreIf +import java.util.zip.GZIPInputStream + @IgnoreIf({ // TODO https://github.com/eclipse-vertx/vert.x/issues/2172 new BigDecimal(System.getProperty("java.specification.version")).isAtLeast(17.0) }) @@ -70,4 +76,36 @@ class AppSecVertxSmokeTest extends AbstractAppSecServerSmokeTest { span.meta.containsKey('_dd.appsec.s.req.params') span.meta.containsKey('_dd.appsec.s.req.headers') } + + void 'test response schema extraction'() { + given: + def url = "http://localhost:${httpPort}/api_security/response" + def client = OkHttpUtils.clientBuilder().build() + def body = [ + "main" : [["key": "id001", "value": 1345.67], ["value": 1567.89, "key": "id002"]], + "nullable": null, + ] + def request = new Request.Builder() + .url(url) + .post(RequestBody.create(MediaType.get('application/json'), JsonOutput.toJson(body))) + .build() + + when: + final response = client.newCall(request).execute() + waitForTraceCount(1) + + then: + response.code() == 200 + def span = rootSpans.first() + span.meta.containsKey('_dd.appsec.s.res.headers') + span.meta.containsKey('_dd.appsec.s.res.body') + final schema = new JsonSlurper().parse(unzip(span.meta.get('_dd.appsec.s.res.body'))) + assert schema == [["main": [[[["key": [8], "value": [16]]]], ["len": 2]], "nullable": [1]]] + } + + + private static byte[] unzip(final String text) { + final inflaterStream = new GZIPInputStream(new ByteArrayInputStream(text.decodeBase64())) + return inflaterStream.getBytes() + } } diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java index 11a19eedcb7..41d32658d4b 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/Events.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/Events.java @@ -322,6 +322,18 @@ public EventType> httpRoute() { return (EventType>) HTTP_ROUTE; } + static final int RESPONSE_BODY_ID = 27; + + @SuppressWarnings("rawtypes") + private static final EventType RESPONSE_BODY = new ET<>("response.body", RESPONSE_BODY_ID); + /** + * The original response body object used by the framework before being serialized to the response + */ + @SuppressWarnings("unchecked") + public EventType>> responseBody() { + return (EventType>>) RESPONSE_BODY; + } + static final int MAX_EVENTS = nextId.get(); private static final class ET extends EventType { diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java b/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java index ac20fc5997f..d8ba93910a3 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java @@ -23,6 +23,7 @@ import static datadog.trace.api.gateway.Events.REQUEST_PATH_PARAMS_ID; import static datadog.trace.api.gateway.Events.REQUEST_SESSION_ID; import static datadog.trace.api.gateway.Events.REQUEST_STARTED_ID; +import static datadog.trace.api.gateway.Events.RESPONSE_BODY_ID; import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_DONE_ID; import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_ID; import static datadog.trace.api.gateway.Events.RESPONSE_STARTED_ID; @@ -347,6 +348,7 @@ public Flow apply(RequestContext ctx, StoredBodySupplier storedBodySupplie case GRPC_SERVER_REQUEST_MESSAGE_ID: case GRAPHQL_SERVER_REQUEST_MESSAGE_ID: case REQUEST_BODY_CONVERTED_ID: + case RESPONSE_BODY_ID: return (C) new BiFunction>() { @Override diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java index 5cbede406d1..93841161245 100644 --- a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java @@ -190,6 +190,9 @@ public void testNormalCalls() { ss.registerCallback(events.requestBodyProcessed(), callback); assertThat(cbp.getCallback(events.requestBodyProcessed()).apply(null, null).getAction()) .isEqualTo(Flow.Action.Noop.INSTANCE); + ss.registerCallback(events.responseBody(), callback); + assertThat(cbp.getCallback(events.responseBody()).apply(null, null).getAction()) + .isEqualTo(Flow.Action.Noop.INSTANCE); ss.registerCallback(events.grpcServerMethod(), callback); assertThat(cbp.getCallback(events.grpcServerMethod()).apply(null, null).getAction()) .isEqualTo(Flow.Action.Noop.INSTANCE); @@ -262,6 +265,9 @@ public void testThrowableBlocking() { ss.registerCallback(events.requestBodyProcessed(), throwback); assertThat(cbp.getCallback(events.requestBodyProcessed()).apply(null, null).getAction()) .isEqualTo(Flow.Action.Noop.INSTANCE); + ss.registerCallback(events.responseBody(), throwback); + assertThat(cbp.getCallback(events.responseBody()).apply(null, null).getAction()) + .isEqualTo(Flow.Action.Noop.INSTANCE); ss.registerCallback(events.grpcServerMethod(), throwback); assertThat(cbp.getCallback(events.grpcServerMethod()).apply(null, null).getAction()) .isEqualTo(Flow.Action.Noop.INSTANCE);