diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index bed9b4d2178..8e9a746cd68 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -36,6 +36,7 @@ import datadog.trace.bootstrap.instrumentation.api.URIUtils; import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; import datadog.trace.bootstrap.instrumentation.decorator.http.ClientIpAddressResolver; +import java.io.OutputStream; import java.net.InetAddress; import java.util.BitSet; import java.util.LinkedHashMap; @@ -383,6 +384,8 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) { final int status = status(response); onResponseStatus(span, status); + final OutputStream responseBody = responseBody(response); + AgentPropagation.ContextVisitor getter = responseGetter(); if (getter != null) { ResponseHeaderTagClassifier tagger = @@ -393,12 +396,17 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) { } if (!isAppSecOnResponseSeparate()) { - callIGCallbackResponseAndHeaders(span, response, status); + callIGCallbackResponseAndHeaders(span, response, status, responseBody); } } return span; } + // TODO this should be abstract by the end of developments + protected OutputStream responseBody(RESPONSE response) { + return null; + } + // @Override // public Span onError(final Span span, final Throwable throwable) { // assert span != null; @@ -464,7 +472,8 @@ private Flow callIGCallbackRequestHeaders(AgentSpan span, REQUEST_CARRIER IGKeyClassifier.create( requestContext, cbp.getCallback(EVENTS.requestHeader()), - cbp.getCallback(EVENTS.requestHeaderDone())); + cbp.getCallback(EVENTS.requestHeaderDone()), + null); if (null != igKeyClassifier) { getter.forEachKey(carrier, igKeyClassifier); return igKeyClassifier.done(); @@ -493,15 +502,16 @@ private Flow callIGCallbackRequestSessionId(final AgentSpan span, final RE } private Flow callIGCallbackResponseAndHeaders( - AgentSpan span, RESPONSE carrier, int status) { - return callIGCallbackResponseAndHeaders(span, carrier, status, responseGetter()); + AgentSpan span, RESPONSE carrier, int status, OutputStream responseBody) { + return callIGCallbackResponseAndHeaders(span, carrier, status, responseGetter(), responseBody); } public Flow callIGCallbackResponseAndHeaders( AgentSpan span, RESP carrier, int status, - AgentPropagation.ContextVisitor contextVisitor) { + AgentPropagation.ContextVisitor contextVisitor, + OutputStream responseBody) { CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC); RequestContext requestContext = span.getRequestContext(); if (cbp == null || requestContext == null) { @@ -520,7 +530,8 @@ public Flow callIGCallbackResponseAndHeaders( IGKeyClassifier.create( requestContext, cbp.getCallback(EVENTS.responseHeader()), - cbp.getCallback(EVENTS.responseHeaderDone())); + cbp.getCallback(EVENTS.responseHeaderDone()), + cbp.getCallback(EVENTS.responseBodyDone())); if (null != igKeyClassifier) { contextVisitor.forEachKey(carrier, igKeyClassifier); return igKeyClassifier.done(); @@ -604,7 +615,8 @@ protected static final class IGKeyClassifier implements AgentPropagation.KeyClas public static IGKeyClassifier create( RequestContext requestContext, TriConsumer headerCallback, - Function> doneCallback) { + Function> doneCallback, + BiFunction> bodyDoneCallback) { if (null == requestContext || null == headerCallback) { return null; } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java index d88c2fb0311..3b395648db3 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java @@ -48,7 +48,7 @@ public interface KnownAddresses { Address RESPONSE_BODY_OBJECT = new Address<>("server.response.body"); /** First chars of HTTP response body */ - Address RESPONSE_BODY_RAW = new Address<>("server.response.body.raw"); + Address RESPONSE_BODY_RAW = new Address<>("server.response.body.raw"); /** Reponse headers excluding cookies */ Address>> RESPONSE_HEADERS_NO_COOKIES = 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 6314b76ba44..87b8376e162 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 @@ -15,6 +15,7 @@ import datadog.trace.api.internal.TraceSegment; import datadog.trace.util.stacktrace.StackTraceEvent; import java.io.Closeable; +import java.io.OutputStream; import java.util.*; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; @@ -98,6 +99,7 @@ public class AppSecRequestContext implements DataBundle, Closeable { private String inferredClientIp; private volatile StoredBodySupplier storedRequestBodySupplier; + private volatile OutputStream storedResponseBodySupplier; private String dbType; private int responseStatus; @@ -106,6 +108,7 @@ public class AppSecRequestContext implements DataBundle, Closeable { private boolean rawReqBodyPublished; private boolean convertedReqBodyPublished; private boolean respDataPublished; + private boolean rawResBodyPublished; private boolean pathParamsPublished; private volatile Map derivatives; @@ -451,6 +454,10 @@ void setStoredRequestBodySupplier(StoredBodySupplier storedRequestBodySupplier) this.storedRequestBodySupplier = storedRequestBodySupplier; } + void setStoredResponseBodySupplier(OutputStream storedResponseBodySupplier) { + this.storedResponseBodySupplier = storedResponseBodySupplier; + } + public String getDbType() { return dbType; } @@ -503,10 +510,18 @@ public boolean isRespDataPublished() { return respDataPublished; } + public boolean isRawResBodyPublished() { + return rawResBodyPublished; + } + public void setRespDataPublished(boolean respDataPublished) { this.respDataPublished = respDataPublished; } + public void setRawResBodyPublished(boolean rawResBodyPublished) { + this.rawResBodyPublished = rawResBodyPublished; + } + /** * Updates the current used usr.id * @@ -580,6 +595,14 @@ public CharSequence getStoredRequestBody() { return storedRequestBodySupplier.get(); } + /** @return the contents of stream */ + public CharSequence getStoredResponseBody() { + CharSequence storedResponseBody = null; + storedResponseBody = + this.storedResponseBodySupplier != null ? this.storedResponseBodySupplier.toString() : null; + return storedResponseBody; + } + public void reportEvents(Collection appSecEvents) { for (AppSecEvent event : appSecEvents) { StandardizedLogging.attackDetected(log, event); 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 310dc5e4853..f9a5bf27c3a 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 @@ -37,6 +37,7 @@ import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; import datadog.trace.util.stacktrace.StackTraceEvent; import datadog.trace.util.stacktrace.StackUtils; +import java.io.OutputStream; import java.net.URI; import java.net.URISyntaxException; import java.nio.charset.Charset; @@ -94,6 +95,7 @@ public class GatewayBridge { // subscriber cache private volatile DataSubscriberInfo initialReqDataSubInfo; private volatile DataSubscriberInfo rawRequestBodySubInfo; + private volatile DataSubscriberInfo rawResponseBodySubInfo; private volatile DataSubscriberInfo requestBodySubInfo; private volatile DataSubscriberInfo pathParamsSubInfo; private volatile DataSubscriberInfo respDataSubInfo; @@ -141,6 +143,8 @@ public void init() { subscriptionService.registerCallback(EVENTS.responseStarted(), this::onResponseStarted); subscriptionService.registerCallback(EVENTS.responseHeader(), this::onResponseHeader); subscriptionService.registerCallback(EVENTS.responseHeaderDone(), this::onResponseHeaderDone); + subscriptionService.registerCallback(EVENTS.responseBodyStart(), this::onResponseBodyStart); + subscriptionService.registerCallback(EVENTS.responseBodyDone(), this::onResponseBodyDone); subscriptionService.registerCallback(EVENTS.grpcServerMethod(), this::onGrpcServerMethod); subscriptionService.registerCallback( EVENTS.grpcServerRequestMessage(), this::onGrpcServerRequestMessage); @@ -602,7 +606,7 @@ private Flow onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup } CharSequence bodyContent = supplier.get(); - if (bodyContent == null || bodyContent.length() == 0) { + if (bodyContent.length() == 0) { return NoopFlow.INSTANCE; } DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.REQUEST_BODY_RAW, bodyContent); @@ -615,6 +619,38 @@ private Flow onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup } } + private Flow onResponseBodyDone(RequestContext ctx_, OutputStream supplier) { + AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); + if (ctx == null || ctx.isRawResBodyPublished()) { + return NoopFlow.INSTANCE; + } + ctx.setRawResBodyPublished(true); + + while (true) { + DataSubscriberInfo subInfo = rawResponseBodySubInfo; + if (subInfo == null) { + subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_RAW); + rawResponseBodySubInfo = subInfo; + } + if (subInfo == null || subInfo.isEmpty()) { + return NoopFlow.INSTANCE; + } + + CharSequence bodyContent = supplier.toString(); + if (bodyContent.length() == 0) { + return NoopFlow.INSTANCE; + } + DataBundle bundle = + new SingletonDataBundle<>(KnownAddresses.RESPONSE_BODY_OBJECT, bodyContent); + try { + GatewayContext gwCtx = new GatewayContext(false); + return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx); + } catch (ExpiredSubscriberInfoException e) { + rawResponseBodySubInfo = null; + } + } + } + private Flow onRequestPathParams(RequestContext ctx_, Map data) { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null || ctx.isPathParamsPublished()) { @@ -651,6 +687,16 @@ private Void onRequestBodyStart(RequestContext ctx_, StoredBodySupplier supplier return null; } + private Void onResponseBodyStart(RequestContext ctx_, OutputStream supplier) { + AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); + if (ctx == null) { + return null; + } + + ctx.setStoredResponseBodySupplier(supplier); + return null; + } + private Flow onRequestStarted() { if (!AppSecSystem.isActive()) { return RequestContextSupplier.EMPTY; @@ -1014,8 +1060,12 @@ private Flow maybePublishResponseData(AppSecRequestContext ctx) { MapDataBundle bundle = MapDataBundle.of( - KnownAddresses.RESPONSE_STATUS, String.valueOf(ctx.getResponseStatus()), - KnownAddresses.RESPONSE_HEADERS_NO_COOKIES, ctx.getResponseHeaders()); + KnownAddresses.RESPONSE_STATUS, + String.valueOf(ctx.getResponseStatus()), + KnownAddresses.RESPONSE_HEADERS_NO_COOKIES, + ctx.getResponseHeaders(), + KnownAddresses.RESPONSE_BODY_OBJECT, + ctx.getStoredResponseBody()); while (true) { DataSubscriberInfo subInfo = respDataSubInfo; @@ -1110,6 +1160,9 @@ private static class IGAppSecEventDependencies { KnownAddresses.REQUEST_BODY_RAW, l(EVENTS.requestBodyStart(), EVENTS.requestBodyDone())); DATA_DEPENDENCIES.put(KnownAddresses.REQUEST_PATH_PARAMS, l(EVENTS.requestPathParams())); DATA_DEPENDENCIES.put(KnownAddresses.REQUEST_BODY_OBJECT, l(EVENTS.requestBodyProcessed())); + DATA_DEPENDENCIES.put( + KnownAddresses.RESPONSE_BODY_RAW, + l(EVENTS.responseBodyStart(), EVENTS.responseBodyDone())); } private static Collection> l( 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 6839b7061b9..3f61897f296 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 @@ -11,7 +11,6 @@ import com.datadog.appsec.report.AppSecEvent import com.datadog.appsec.report.AppSecEventWrapper import datadog.trace.api.ProductTraceSource import datadog.trace.api.config.GeneralConfig -import static datadog.trace.api.config.IastConfig.IAST_DEDUPLICATION_ENABLED import datadog.trace.api.function.TriConsumer import datadog.trace.api.function.TriFunction import datadog.trace.api.gateway.BlockResponseFunction @@ -114,6 +113,8 @@ class GatewayBridgeSpecification extends DDSpecification { BiFunction> shellCmdCB BiFunction> userCB TriFunction> loginEventCB + BiFunction responseBodyStartCB + BiFunction> responseBodyDoneCB WafMetricCollector wafMetricCollector = Mock(WafMetricCollector) @@ -452,6 +453,8 @@ class GatewayBridgeSpecification extends DDSpecification { 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 } + 1 * ig.registerCallback(EVENTS.responseBodyStart(), _) >> { responseBodyStartCB = it[1]; null } + 1 * ig.registerCallback(EVENTS.responseBodyDone(), _) >> { responseBodyDoneCB = it[1]; null } 1 * ig.registerCallback(EVENTS.grpcServerMethod(), _) >> { grpcServerMethodCB = it[1]; null } 1 * ig.registerCallback(EVENTS.grpcServerRequestMessage(), _) >> { grpcServerRequestMessageCB = it[1]; null } 1 * ig.registerCallback(EVENTS.graphqlServerRequestMessage(), _) >> { graphqlServerRequestMessageCB = it[1]; null } @@ -933,7 +936,7 @@ class GatewayBridgeSpecification extends DDSpecification { } @Override - def T getOrCreateMetaStructTop(String key, Function defaultValue) { + T getOrCreateMetaStructTop(String key, Function defaultValue) { return null } @@ -991,7 +994,7 @@ class GatewayBridgeSpecification extends DDSpecification { getTraceSegment() >> traceSegment } final spanInfo = Mock(AgentSpan) { - getTags() >> ['http.route':'/'] + getTags() >> ['http.route': '/'] } when: diff --git a/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/BlockingResponseHelper.java b/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/BlockingResponseHelper.java index 046d6c96bdb..f78b63ea05b 100644 --- a/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/BlockingResponseHelper.java +++ b/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/BlockingResponseHelper.java @@ -37,7 +37,11 @@ public static HttpResponse handleFinishForWaf(final AgentSpan span, final HttpRe } Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, response, response.status().intValue(), AkkaHttpServerHeaders.responseGetter()); + span, + response, + response.status().intValue(), + AkkaHttpServerHeaders.responseGetter(), + null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java index 4728c725ca0..bd52c6e31ed 100644 --- a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java @@ -84,7 +84,7 @@ static class CommitResponseAdvice { Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER); + span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java index 380bf5a30d1..ce630c0cb77 100644 --- a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java @@ -98,7 +98,7 @@ static class CommitResponseAdvice { Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER); + span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java index fb92469eaa7..cd542069615 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java @@ -119,7 +119,7 @@ static class CommitResponseAdvice { Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER); + span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java index 009965e4337..ba70a7eba0a 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java @@ -69,7 +69,7 @@ public class JettyCommitResponseHelper { Response resp = connection.getResponse(); Flow flow = JettyDecorator.DECORATE.callIGCallbackResponseAndHeaders( - span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER); + span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java index 7fb3bec4fcc..d4a409dd20c 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java @@ -69,7 +69,7 @@ public class JettyCommitResponseHelper { Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER); + span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java index 3e9afaebc7d..b960276775e 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java @@ -68,7 +68,7 @@ public class JettyCommitResponseHelper { Response resp = connection.getResponse(); Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER); + span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java index 2fd93b65141..5850b11460c 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java @@ -71,7 +71,7 @@ public class JettyCommitResponseHelper { Response resp = connection.getResponse(); Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER); + span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java index 7511cff8c56..b5e6ea6db2c 100644 --- a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java @@ -69,7 +69,11 @@ public void writeRequested(ChannelHandlerContext ctx, MessageEvent msg) throws E Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, origResponse, origResponse.getStatus().getCode(), ResponseExtractAdapter.GETTER); + span, + origResponse, + origResponse.getStatus().getCode(), + ResponseExtractAdapter.GETTER, + null); channelTraceContext.setAnalyzedResponse(true); Flow.Action action = flow.getAction(); if (!(action instanceof Flow.Action.RequestBlockingAction)) { diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java index 6e3ec443840..49eaf7f49d0 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java @@ -88,7 +88,11 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, origResponse, origResponse.getStatus().code(), ResponseExtractAdapter.GETTER); + span, + origResponse, + origResponse.getStatus().code(), + ResponseExtractAdapter.GETTER, + null); markAnalyzedResponse(channel); Flow.Action action = flow.getAction(); if (!(action instanceof Flow.Action.RequestBlockingAction)) { diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java index 9139e318703..f562206f6c3 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java @@ -85,7 +85,11 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr Flow flow = DECORATE.callIGCallbackResponseAndHeaders( - span, origResponse, origResponse.getStatus().code(), ResponseExtractAdapter.GETTER); + span, + origResponse, + origResponse.getStatus().code(), + ResponseExtractAdapter.GETTER, + null); markAnalyzedResponse(channel); Flow.Action action = flow.getAction(); if (!(action instanceof Flow.Action.RequestBlockingAction)) { diff --git a/dd-java-agent/instrumentation/tomcat-appsec-5.5/src/main/java/datadog/trace/instrumentation/tomcat55/CommitActionInstrumentation.java b/dd-java-agent/instrumentation/tomcat-appsec-5.5/src/main/java/datadog/trace/instrumentation/tomcat55/CommitActionInstrumentation.java index d15a19f09dd..1a8a68038ac 100644 --- a/dd-java-agent/instrumentation/tomcat-appsec-5.5/src/main/java/datadog/trace/instrumentation/tomcat55/CommitActionInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-appsec-5.5/src/main/java/datadog/trace/instrumentation/tomcat55/CommitActionInstrumentation.java @@ -106,7 +106,8 @@ static class ProcessCommitActionAdvice { agentSpan, coyoteResponse, coyoteResponse.getStatus(), - ExtractAdapter.CoyoteResponse.GETTER); + ExtractAdapter.CoyoteResponse.GETTER, + null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/tomcat-appsec-7/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java b/dd-java-agent/instrumentation/tomcat-appsec-7/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java index 670c877ffad..93ac615c2c6 100644 --- a/dd-java-agent/instrumentation/tomcat-appsec-7/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-appsec-7/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java @@ -106,7 +106,8 @@ static class ProcessCommitActionAdvice { agentSpan, coyoteResponse, coyoteResponse.getStatus(), - ExtractAdapter.CoyoteResponse.GETTER); + ExtractAdapter.CoyoteResponse.GETTER, + null); Flow.Action action = flow.getAction(); if (action instanceof Flow.Action.RequestBlockingAction) { Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action; diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpServerExchangeSenderInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpServerExchangeSenderInstrumentation.java index e3594c2d459..5e5550892e3 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpServerExchangeSenderInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpServerExchangeSenderInstrumentation.java @@ -72,7 +72,7 @@ static class GetResponseChannelAdvice { AgentSpan span = continuation.span(); Flow flow = UndertowDecorator.DECORATE.callIGCallbackResponseAndHeaders( - span, xchg, xchg.getStatusCode(), UndertowExtractAdapter.Response.GETTER); + span, xchg, xchg.getStatusCode(), UndertowExtractAdapter.Response.GETTER, null); Flow.Action action = flow.getAction(); if (!(action instanceof Flow.Action.RequestBlockingAction)) { return false; 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 d840cad01c3..9651cd5e33e 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 @@ -5,6 +5,7 @@ import datadog.trace.api.http.StoredBodySupplier; import datadog.trace.api.telemetry.LoginEvent; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; +import java.io.OutputStream; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; @@ -312,6 +313,28 @@ public EventType>> shellCmd() { return (EventType>>) SHELL_CMD; } + static final int RESPONSE_BODY_START_ID = 26; + + @SuppressWarnings("rawtypes") + private static final EventType RESPONSE_BODY_START = + new ET<>("response.body.started", RESPONSE_BODY_START_ID); + /** The request body has started being read */ + @SuppressWarnings("unchecked") + public EventType> responseBodyStart() { + return (EventType>) RESPONSE_BODY_START; + } + + static final int RESPONSE_BODY_DONE_ID = 27; + + @SuppressWarnings("rawtypes") + private static final EventType RESPONSE_BODY_DONE = + new ET<>("response.body.done", RESPONSE_BODY_DONE_ID); + /** The request body is done being read */ + @SuppressWarnings("unchecked") + public EventType>> responseBodyDone() { + return (EventType>>) RESPONSE_BODY_DONE; + } + 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 cd5219f1dfc..5c18cfc337c 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 @@ -22,6 +22,8 @@ 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_DONE_ID; +import static datadog.trace.api.gateway.Events.RESPONSE_BODY_START_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; @@ -315,6 +317,7 @@ public boolean equals(Object obj) { return callback.equals(obj); } }; + case RESPONSE_BODY_START_ID: case REQUEST_BODY_START_ID: return (C) new BiFunction() { @@ -329,6 +332,7 @@ public Void apply(RequestContext ctx, StoredBodySupplier storedBodySupplier) { } } }; + case RESPONSE_BODY_DONE_ID: case REQUEST_BODY_DONE_ID: return (C) new BiFunction>() {