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 a73de84d7d0..1ad8c7ccc4e 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 b5f01f0c91b..e610e7267c0 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 @@ -96,6 +96,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; @@ -135,6 +136,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( @@ -175,6 +177,7 @@ public void reset() { initialReqDataSubInfo = null; rawRequestBodySubInfo = null; requestBodySubInfo = null; + responseBodySubInfo = null; pathParamsSubInfo = null; respDataSubInfo = null; grpcServerMethodSubInfo = null; @@ -636,6 +639,40 @@ 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; + } + // TODO: review schema extraction limits + 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/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index b79c3f2ca0b..d4610272a97 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 @@ -450,6 +451,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 } @@ -1327,4 +1329,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..bff52a670ee --- /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,58 @@ +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 io.vertx.core.json.JsonObject; +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) Object source, @ActiveRequestContext RequestContext reqCtx) { + + if (source == null) { + return; + } + + Object object = source; + if (object instanceof JsonObject) { + object = ((JsonObject) object).getMap(); + } + + CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); + 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 afd26bc61bc..0e752208850 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 @@ -135,6 +135,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 +336,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 +355,10 @@ abstract class HttpServerTest extends WithHttpServer { false } + boolean testResponseBodyJson() { + false + } + boolean testBlocking() { false } @@ -1581,6 +1587,40 @@ 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()) + def request = request( + BODY_JSON, 'POST', + RequestBody.create(MediaType.get('application/json'), '{"a": "x"}')) + .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) + + then: + TEST_WRITER.get(0).any { + it.getTag('response.body') == '[a:[x]]' + } + + 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 +2320,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 +2344,7 @@ abstract class HttpServerTest extends WithHttpServer { boolean bodyEndBlock boolean bodyConvertedBlock boolean responseHeadersInTags + boolean responseBodyTag } static final String stringOrEmpty(String string) { @@ -2356,6 +2398,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 +2495,33 @@ abstract class HttpServerTest extends WithHttpServer { } } as BiFunction>) + final BiFunction> responseBodyObjectCb = + ({ RequestContext rqCtxt, Object obj -> + if (obj instanceof Map) { + obj = obj.collectEntries { + [ + it.key, + (it.value instanceof Iterable || it.value instanceof String[]) ? it.value : [it.value] + ] + } + } else if (!(obj instanceof String) && !(obj instanceof List)) { + obj = obj.properties + .findAll { it.key != 'class' } + .collectEntries { [it.key, it.value instanceof Iterable ? it.value : [it.value]] } + } + Context context = rqCtxt.getData(RequestContextSlot.APPSEC) + if (context.responseBodyTag) { + rqCtxt.traceSegment.setTagTop('response.body', obj as String) + } + 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..55916f936a8 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,6 +1,9 @@ import datadog.smoketest.appsec.AbstractAppSecServerSmokeTest import datadog.trace.agent.test.utils.OkHttpUtils +import groovy.json.JsonOutput +import okhttp3.MediaType import okhttp3.Request +import okhttp3.RequestBody import okhttp3.Response import spock.lang.IgnoreIf @@ -70,4 +73,37 @@ 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 = [ + source: 'AppSecVertxSmokeTest', + tests : [ + [ + name : 'API Security samples only one request per endpoint', + status: 'SUCCESS' + ], + [ + name : 'test response schema extraction', + status: 'FAILED' + ] + ] + ] + 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') + } } 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..b57305fc541 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 @@ -312,7 +312,19 @@ public EventType>> shellCmd() { return (EventType>>) SHELL_CMD; } - static final int HTTP_ROUTE_ID = 26; + static final int RESPONSE_BODY_ID = 26; + + @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 HTTP_ROUTE_ID = 27; @SuppressWarnings("rawtypes") private static final EventType HTTP_ROUTE = new ET<>("http.route", HTTP_ROUTE_ID); 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 8f2840d4308..e7dd23076a8 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); @@ -260,6 +263,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);