From 5863cc154a84909b4756d6c909cb2230841d1070 Mon Sep 17 00:00:00 2001 From: "sezen.leblay" Date: Thu, 15 May 2025 14:25:15 +0200 Subject: [PATCH 1/5] HTTP response schema collection and data classification --- .../decorator/HttpServerDecorator.java | 10 +- .../appsec/event/data/KnownAddresses.java | 2 +- .../appsec/gateway/AppSecRequestContext.java | 23 ++++ .../datadog/appsec/gateway/GatewayBridge.java | 58 +++++++++- .../gateway/GatewayBridgeSpecification.groovy | 100 +++++++++++++++++- .../appsec/SpringBootSmokeTest.groovy | 2 + .../datadog/trace/api/gateway/Events.java | 23 ++++ .../api/gateway/InstrumentationGateway.java | 4 + 8 files changed, 212 insertions(+), 10 deletions(-) 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 29e33a3dd8c..30532bc2ef5 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 @@ -18,6 +18,7 @@ import datadog.trace.api.gateway.IGSpanInfo; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.http.StoredBodySupplier; import datadog.trace.api.naming.SpanNaming; import datadog.trace.bootstrap.ActiveSubsystems; import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; @@ -435,7 +436,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(); @@ -491,7 +493,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(); @@ -575,7 +578,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..013038cb24c 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 @@ -98,6 +98,7 @@ public class AppSecRequestContext implements DataBundle, Closeable { private String inferredClientIp; private volatile StoredBodySupplier storedRequestBodySupplier; + private volatile StoredBodySupplier storedResponseBodySupplier; private String dbType; private int responseStatus; @@ -106,6 +107,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 +453,10 @@ void setStoredRequestBodySupplier(StoredBodySupplier storedRequestBodySupplier) this.storedRequestBodySupplier = storedRequestBodySupplier; } + void setStoredResponseBodySupplier(StoredBodySupplier storedResponseBodySupplier) { + this.storedResponseBodySupplier = storedResponseBodySupplier; + } + public String getDbType() { return dbType; } @@ -503,10 +509,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 +594,15 @@ public CharSequence getStoredRequestBody() { return storedRequestBodySupplier.get(); } + /** @return the portion of the body read so far, if any */ + public CharSequence getStoredResponseBody() { + StoredBodySupplier storedResponseBodySupplier = this.storedResponseBodySupplier; + if (storedResponseBodySupplier == null) { + return null; + } + return storedResponseBodySupplier.get(); + } + 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 9ba2bdb07a3..56a33d9ab38 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 @@ -94,6 +94,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 +142,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 +605,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 +618,38 @@ private Flow onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup } } + private Flow onResponseBodyDone(RequestContext ctx_, StoredBodySupplier 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.get(); + 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 +686,16 @@ private Void onRequestBodyStart(RequestContext ctx_, StoredBodySupplier supplier return null; } + private Void onResponseBodyStart(RequestContext ctx_, StoredBodySupplier 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; @@ -962,8 +1007,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; @@ -1058,6 +1107,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 4eb14da17dc..8f35bf838ad 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: @@ -1222,4 +1225,95 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) } + void 'forwards response body start events and stores the supplier'() { + StoredBodySupplier supplier = Stub() + + setup: + supplier.get() >> 'response content' + + expect: + ctx.data.storedResponseBody == null + + when: + responseBodyStartCB.apply(ctx, supplier) + + then: + ctx.data.storedResponseBody == 'response content' + } + + void 'forwards response body done events and distributes the body contents'() { + DataBundle bundle + GatewayContext gatewayContext + StoredBodySupplier supplier = Stub() + + setup: + supplier.get() >> 'response body content' + eventDispatcher.getDataSubscribers({ KnownAddresses.RESPONSE_BODY_RAW in it }) >> nonEmptyDsInfo + eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> + { bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE } + + when: + responseBodyDoneCB.apply(ctx, supplier) + + then: + bundle.get(KnownAddresses.RESPONSE_BODY_OBJECT) == 'response body content' + gatewayContext.isTransient == false + gatewayContext.isRasp == false + } + + void 'response body does not get published twice'() { + StoredBodySupplier supplier = Stub() + Flow flow + + given: + supplier.get() >> 'response body content' + + when: + ctx.data.setRawResBodyPublished(true) + flow = responseBodyDoneCB.apply(ctx, supplier) + + then: + flow == NoopFlow.INSTANCE + 0 * eventDispatcher.getDataSubscribers(KnownAddresses.RESPONSE_BODY_RAW) + } + + void 'response body with empty content is not published'() { + StoredBodySupplier supplier = Stub() + Flow flow + + given: + supplier.get() >> '' + + when: + flow = responseBodyDoneCB.apply(ctx, supplier) + + then: + flow == NoopFlow.INSTANCE + 1 * eventDispatcher.getDataSubscribers({ KnownAddresses.RESPONSE_BODY_RAW in it }) >> nonEmptyDsInfo + 0 * eventDispatcher.publishDataEvent(_, _, _, _) + } + + void 'response data includes response body when available'() { + setup: + eventDispatcher.getDataSubscribers({ KnownAddresses.RESPONSE_STATUS in it }) >> nonEmptyDsInfo + ctx.data.responseStatus = 200 + ctx.data.addResponseHeader('Content-Type', 'application/json') + + StoredBodySupplier supplier = Stub() + supplier.get() >> '{"status":"success"}' + ctx.data.storedResponseBodySupplier = supplier + + DataBundle bundle + + when: + respHeadersDoneCB.apply(ctx) + + then: + 1 * eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> + { dsInfo, appSecCtx, db, gwCtx -> bundle = db; NoopFlow.INSTANCE } + bundle.get(KnownAddresses.RESPONSE_STATUS) == '200' + bundle.get(KnownAddresses.RESPONSE_HEADERS_NO_COOKIES) == ['content-type': ['application/json']] + bundle.get(KnownAddresses.RESPONSE_BODY_OBJECT) == '{"status":"success"}' + } + } diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy index 42768f998b3..fb88ef9d69e 100644 --- a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy +++ b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy @@ -219,6 +219,8 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest { List command = new ArrayList<>() command.add(javaPath()) + // TODO delete when ure done testing + command.add("-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=5005") command.addAll(defaultJavaProperties) command.addAll(defaultAppSecProperties) command.addAll((String[]) ["-jar", springBootShadowJar, "--server.port=${httpPort}"]) 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..1ea2b969321 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,6 +312,29 @@ 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>() { From b47758198440c3a36383996cb86aebfd3cc88682 Mon Sep 17 00:00:00 2001 From: "sezen.leblay" Date: Mon, 19 May 2025 14:15:54 +0200 Subject: [PATCH 2/5] wip generated output streams for response --- .../AbstractServletOutputStreamWrapper.java | 48 +++++++++++++ .../HttpServletGetOutputStreamAdvice.java | 70 +++++++++++++++++++ .../Servlet31OutputStreamWrapper.java | 23 ++++++ .../Servlet31RequestBodyInstrumentation.java | 10 ++- 4 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java create mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java create mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java diff --git a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java new file mode 100644 index 00000000000..7ac7322a3bf --- /dev/null +++ b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java @@ -0,0 +1,48 @@ +package datadog.trace.instrumentation.servlet; + +import datadog.trace.api.http.StoredByteBody; +import java.io.IOException; +import javax.servlet.ServletOutputStream; + +public abstract class AbstractServletOutputStreamWrapper extends ServletOutputStream { + protected final ServletOutputStream os; + private final StoredByteBody storedByteBody; + + public AbstractServletOutputStreamWrapper(ServletOutputStream os, StoredByteBody storedByteBody) { + this.os = os; + this.storedByteBody = storedByteBody; + } + + @Override + public void write(int b) throws IOException { + os.write(b); + storedByteBody.appendData(b); + } + + @Override + public void write(byte[] b) throws IOException { + os.write(b); + if (b != null && b.length > 0) { + storedByteBody.appendData(b, 0, b.length); + } + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + os.write(b, off, len); + if (b != null && len > 0) { + storedByteBody.appendData(b, off, off + len); + } + } + + @Override + public void flush() throws IOException { + os.flush(); + } + + @Override + public void close() throws IOException { + os.close(); + storedByteBody.maybeNotifyAndBlock(); + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java new file mode 100644 index 00000000000..961220115c1 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java @@ -0,0 +1,70 @@ +package datadog.trace.instrumentation.servlet3; + +import static datadog.trace.api.gateway.Events.EVENTS; + +import datadog.trace.advice.ActiveRequestContext; +import datadog.trace.advice.RequiresRequestContext; +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.api.http.StoredBodySupplier; +import datadog.trace.api.http.StoredByteBody; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import java.nio.charset.Charset; +import java.util.function.BiFunction; +import javax.servlet.ServletOutputStream; +import javax.servlet.http.HttpServletResponse; +import net.bytebuddy.asm.Advice; + +@RequiresRequestContext(RequestContextSlot.APPSEC) +public class HttpServletGetOutputStreamAdvice { + + @Advice.OnMethodExit(suppress = Throwable.class) + static void after( + @Advice.This final HttpServletResponse resp, + @Advice.Return(readOnly = false) ServletOutputStream os, + @ActiveRequestContext RequestContext reqCtx) { + if (os == null) { + return; + } + + if (os instanceof Servlet31OutputStreamWrapper) { + return; + } + + CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); + BiFunction responseStartCb = + cbp.getCallback(EVENTS.responseBodyStart()); + BiFunction> responseEndedCb = + cbp.getCallback(EVENTS.responseBodyDone()); + if (responseStartCb == null || responseEndedCb == null) { + return; + } + + int lengthHint = 0; + String lengthHeader = resp.getHeader("content-length"); + if (lengthHeader != null) { + try { + lengthHint = Integer.parseInt(lengthHeader); + } catch (NumberFormatException nfe) { + // purposefully left blank + } + } + + String encoding = resp.getCharacterEncoding(); + Charset charset = null; + try { + if (encoding != null) { + charset = Charset.forName(encoding); + } + } catch (IllegalArgumentException iae) { + // purposefully left blank + } + + StoredByteBody storedByteBody = + new StoredByteBody(reqCtx, responseStartCb, responseEndedCb, charset, lengthHint); + + os = new Servlet31OutputStreamWrapper(os, storedByteBody); + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java new file mode 100644 index 00000000000..ba927d25776 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java @@ -0,0 +1,23 @@ +package datadog.trace.instrumentation.servlet3; + +import datadog.trace.api.http.StoredByteBody; +import datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper; +import javax.servlet.ServletOutputStream; +import javax.servlet.WriteListener; + +/** Provides additional delegation for servlet 3.1 */ +public class Servlet31OutputStreamWrapper extends AbstractServletOutputStreamWrapper { + public Servlet31OutputStreamWrapper(ServletOutputStream os, StoredByteBody storedByteBody) { + super(os, storedByteBody); + } + + @Override + public boolean isReady() { + return os.isReady(); + } + + @Override + public void setWriteListener(WriteListener writeListener) { + os.setWriteListener(writeListener); + } +} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31RequestBodyInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31RequestBodyInstrumentation.java index 76509673a04..bad7382b33e 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31RequestBodyInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31RequestBodyInstrumentation.java @@ -65,6 +65,12 @@ public void methodAdvice(MethodTransformer transformer) { .and(returns(named("java.io.BufferedReader"))) .and(isPublic()), packageName + ".HttpServletGetReaderAdvice"); + transformer.applyAdvice( + named("getOutputStream") + .and(takesNoArguments()) + .and(returns(named("javax.servlet.ServletOutputStream"))) + .and(isPublic()), + packageName + ".HttpServletGetOutputStreamAdvice"); } @Override @@ -72,7 +78,9 @@ public String[] helperClassNames() { return new String[] { "datadog.trace.instrumentation.servlet.BufferedReaderWrapper", "datadog.trace.instrumentation.servlet.AbstractServletInputStreamWrapper", - "datadog.trace.instrumentation.servlet3.Servlet31InputStreamWrapper" + "datadog.trace.instrumentation.servlet3.Servlet31InputStreamWrapper", + "datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper", + "datadog.trace.instrumentation.servlet3.Servlet31OutputStreamWrapper" }; } } From be71f96b7031ce5dd04e570aec0c85041fad5d22 Mon Sep 17 00:00:00 2001 From: "sezen.leblay" Date: Mon, 19 May 2025 16:29:45 +0200 Subject: [PATCH 3/5] wip --- .../decorator/HttpServerDecorator.java | 17 +++++++++++++---- .../akkahttp/appsec/BlockingResponseHelper.java | 6 +++++- .../JettyCommitResponseInstrumentation.java | 2 +- .../JettyCommitResponseInstrumentation.java | 2 +- .../JettyCommitResponseInstrumentation.java | 2 +- .../jetty10/JettyCommitResponseHelper.java | 2 +- .../jetty904/JettyCommitResponseHelper.java | 2 +- .../jetty93/JettyCommitResponseHelper.java | 2 +- .../jetty9421/JettyCommitResponseHelper.java | 2 +- .../server/MaybeBlockResponseHandler.java | 6 +++++- .../server/MaybeBlockResponseHandler.java | 6 +++++- .../server/MaybeBlockResponseHandler.java | 6 +++++- .../tomcat55/CommitActionInstrumentation.java | 3 ++- .../tomcat7/CommitActionInstrumentation.java | 3 ++- ...HttpServerExchangeSenderInstrumentation.java | 2 +- 15 files changed, 45 insertions(+), 18 deletions(-) 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 30532bc2ef5..64e64974b88 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 @@ -34,6 +34,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; @@ -355,6 +356,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 = @@ -365,12 +368,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; @@ -466,15 +474,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) { 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; From fe73cae45218d2a7e1620b13ef7de33d04b096a4 Mon Sep 17 00:00:00 2001 From: "sezen.leblay" Date: Mon, 19 May 2025 16:56:50 +0200 Subject: [PATCH 4/5] wip --- .../decorator/HttpServerDecorator.java | 3 +- .../appsec/gateway/AppSecRequestContext.java | 16 ++--- .../datadog/appsec/gateway/GatewayBridge.java | 7 +- .../AbstractServletOutputStreamWrapper.java | 48 ------------- .../HttpServletGetOutputStreamAdvice.java | 70 ------------------- .../Servlet31OutputStreamWrapper.java | 23 ------ .../Servlet31RequestBodyInstrumentation.java | 10 +-- .../datadog/trace/api/gateway/Events.java | 10 +-- 8 files changed, 19 insertions(+), 168 deletions(-) delete mode 100644 dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java delete mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java delete mode 100644 dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java 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 64e64974b88..d4de53578f7 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 @@ -18,7 +18,6 @@ import datadog.trace.api.gateway.IGSpanInfo; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; -import datadog.trace.api.http.StoredBodySupplier; import datadog.trace.api.naming.SpanNaming; import datadog.trace.bootstrap.ActiveSubsystems; import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; @@ -588,7 +587,7 @@ public static IGKeyClassifier create( RequestContext requestContext, TriConsumer headerCallback, Function> doneCallback, - BiFunction> bodyDoneCallback) { + BiFunction> bodyDoneCallback) { if (null == requestContext || null == headerCallback) { return null; } 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 013038cb24c..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,7 +99,7 @@ public class AppSecRequestContext implements DataBundle, Closeable { private String inferredClientIp; private volatile StoredBodySupplier storedRequestBodySupplier; - private volatile StoredBodySupplier storedResponseBodySupplier; + private volatile OutputStream storedResponseBodySupplier; private String dbType; private int responseStatus; @@ -453,7 +454,7 @@ void setStoredRequestBodySupplier(StoredBodySupplier storedRequestBodySupplier) this.storedRequestBodySupplier = storedRequestBodySupplier; } - void setStoredResponseBodySupplier(StoredBodySupplier storedResponseBodySupplier) { + void setStoredResponseBodySupplier(OutputStream storedResponseBodySupplier) { this.storedResponseBodySupplier = storedResponseBodySupplier; } @@ -594,13 +595,12 @@ public CharSequence getStoredRequestBody() { return storedRequestBodySupplier.get(); } - /** @return the portion of the body read so far, if any */ + /** @return the contents of stream */ public CharSequence getStoredResponseBody() { - StoredBodySupplier storedResponseBodySupplier = this.storedResponseBodySupplier; - if (storedResponseBodySupplier == null) { - return null; - } - return storedResponseBodySupplier.get(); + CharSequence storedResponseBody = null; + storedResponseBody = + this.storedResponseBodySupplier != null ? this.storedResponseBodySupplier.toString() : null; + return storedResponseBody; } public void reportEvents(Collection appSecEvents) { 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 56a33d9ab38..591f3b888e1 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; @@ -618,7 +619,7 @@ private Flow onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup } } - private Flow onResponseBodyDone(RequestContext ctx_, StoredBodySupplier supplier) { + private Flow onResponseBodyDone(RequestContext ctx_, OutputStream supplier) { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null || ctx.isRawResBodyPublished()) { return NoopFlow.INSTANCE; @@ -635,7 +636,7 @@ private Flow onResponseBodyDone(RequestContext ctx_, StoredBodySupplier su return NoopFlow.INSTANCE; } - CharSequence bodyContent = supplier.get(); + CharSequence bodyContent = supplier.toString(); if (bodyContent.length() == 0) { return NoopFlow.INSTANCE; } @@ -686,7 +687,7 @@ private Void onRequestBodyStart(RequestContext ctx_, StoredBodySupplier supplier return null; } - private Void onResponseBodyStart(RequestContext ctx_, StoredBodySupplier supplier) { + private Void onResponseBodyStart(RequestContext ctx_, OutputStream supplier) { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null) { return null; diff --git a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java deleted file mode 100644 index 7ac7322a3bf..00000000000 --- a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/AbstractServletOutputStreamWrapper.java +++ /dev/null @@ -1,48 +0,0 @@ -package datadog.trace.instrumentation.servlet; - -import datadog.trace.api.http.StoredByteBody; -import java.io.IOException; -import javax.servlet.ServletOutputStream; - -public abstract class AbstractServletOutputStreamWrapper extends ServletOutputStream { - protected final ServletOutputStream os; - private final StoredByteBody storedByteBody; - - public AbstractServletOutputStreamWrapper(ServletOutputStream os, StoredByteBody storedByteBody) { - this.os = os; - this.storedByteBody = storedByteBody; - } - - @Override - public void write(int b) throws IOException { - os.write(b); - storedByteBody.appendData(b); - } - - @Override - public void write(byte[] b) throws IOException { - os.write(b); - if (b != null && b.length > 0) { - storedByteBody.appendData(b, 0, b.length); - } - } - - @Override - public void write(byte[] b, int off, int len) throws IOException { - os.write(b, off, len); - if (b != null && len > 0) { - storedByteBody.appendData(b, off, off + len); - } - } - - @Override - public void flush() throws IOException { - os.flush(); - } - - @Override - public void close() throws IOException { - os.close(); - storedByteBody.maybeNotifyAndBlock(); - } -} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java deleted file mode 100644 index 961220115c1..00000000000 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletGetOutputStreamAdvice.java +++ /dev/null @@ -1,70 +0,0 @@ -package datadog.trace.instrumentation.servlet3; - -import static datadog.trace.api.gateway.Events.EVENTS; - -import datadog.trace.advice.ActiveRequestContext; -import datadog.trace.advice.RequiresRequestContext; -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.api.http.StoredBodySupplier; -import datadog.trace.api.http.StoredByteBody; -import datadog.trace.bootstrap.instrumentation.api.AgentTracer; -import java.nio.charset.Charset; -import java.util.function.BiFunction; -import javax.servlet.ServletOutputStream; -import javax.servlet.http.HttpServletResponse; -import net.bytebuddy.asm.Advice; - -@RequiresRequestContext(RequestContextSlot.APPSEC) -public class HttpServletGetOutputStreamAdvice { - - @Advice.OnMethodExit(suppress = Throwable.class) - static void after( - @Advice.This final HttpServletResponse resp, - @Advice.Return(readOnly = false) ServletOutputStream os, - @ActiveRequestContext RequestContext reqCtx) { - if (os == null) { - return; - } - - if (os instanceof Servlet31OutputStreamWrapper) { - return; - } - - CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); - BiFunction responseStartCb = - cbp.getCallback(EVENTS.responseBodyStart()); - BiFunction> responseEndedCb = - cbp.getCallback(EVENTS.responseBodyDone()); - if (responseStartCb == null || responseEndedCb == null) { - return; - } - - int lengthHint = 0; - String lengthHeader = resp.getHeader("content-length"); - if (lengthHeader != null) { - try { - lengthHint = Integer.parseInt(lengthHeader); - } catch (NumberFormatException nfe) { - // purposefully left blank - } - } - - String encoding = resp.getCharacterEncoding(); - Charset charset = null; - try { - if (encoding != null) { - charset = Charset.forName(encoding); - } - } catch (IllegalArgumentException iae) { - // purposefully left blank - } - - StoredByteBody storedByteBody = - new StoredByteBody(reqCtx, responseStartCb, responseEndedCb, charset, lengthHint); - - os = new Servlet31OutputStreamWrapper(os, storedByteBody); - } -} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java deleted file mode 100644 index ba927d25776..00000000000 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31OutputStreamWrapper.java +++ /dev/null @@ -1,23 +0,0 @@ -package datadog.trace.instrumentation.servlet3; - -import datadog.trace.api.http.StoredByteBody; -import datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper; -import javax.servlet.ServletOutputStream; -import javax.servlet.WriteListener; - -/** Provides additional delegation for servlet 3.1 */ -public class Servlet31OutputStreamWrapper extends AbstractServletOutputStreamWrapper { - public Servlet31OutputStreamWrapper(ServletOutputStream os, StoredByteBody storedByteBody) { - super(os, storedByteBody); - } - - @Override - public boolean isReady() { - return os.isReady(); - } - - @Override - public void setWriteListener(WriteListener writeListener) { - os.setWriteListener(writeListener); - } -} diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31RequestBodyInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31RequestBodyInstrumentation.java index bad7382b33e..76509673a04 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31RequestBodyInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet31RequestBodyInstrumentation.java @@ -65,12 +65,6 @@ public void methodAdvice(MethodTransformer transformer) { .and(returns(named("java.io.BufferedReader"))) .and(isPublic()), packageName + ".HttpServletGetReaderAdvice"); - transformer.applyAdvice( - named("getOutputStream") - .and(takesNoArguments()) - .and(returns(named("javax.servlet.ServletOutputStream"))) - .and(isPublic()), - packageName + ".HttpServletGetOutputStreamAdvice"); } @Override @@ -78,9 +72,7 @@ public String[] helperClassNames() { return new String[] { "datadog.trace.instrumentation.servlet.BufferedReaderWrapper", "datadog.trace.instrumentation.servlet.AbstractServletInputStreamWrapper", - "datadog.trace.instrumentation.servlet3.Servlet31InputStreamWrapper", - "datadog.trace.instrumentation.servlet.AbstractServletOutputStreamWrapper", - "datadog.trace.instrumentation.servlet3.Servlet31OutputStreamWrapper" + "datadog.trace.instrumentation.servlet3.Servlet31InputStreamWrapper" }; } } 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 1ea2b969321..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; @@ -319,8 +320,8 @@ public EventType>> shellCmd() { 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; + public EventType> responseBodyStart() { + return (EventType>) RESPONSE_BODY_START; } static final int RESPONSE_BODY_DONE_ID = 27; @@ -330,9 +331,8 @@ public EventType> responseB 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; + public EventType>> responseBodyDone() { + return (EventType>>) RESPONSE_BODY_DONE; } static final int MAX_EVENTS = nextId.get(); From a29290a8e0a4e1462005356988f814c5c6981776 Mon Sep 17 00:00:00 2001 From: "sezen.leblay" Date: Tue, 20 May 2025 08:40:46 +0200 Subject: [PATCH 5/5] wip --- .../groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy | 2 -- 1 file changed, 2 deletions(-) diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy index fb88ef9d69e..42768f998b3 100644 --- a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy +++ b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy @@ -219,8 +219,6 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest { List command = new ArrayList<>() command.add(javaPath()) - // TODO delete when ure done testing - command.add("-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=5005") command.addAll(defaultJavaProperties) command.addAll(defaultAppSecProperties) command.addAll((String[]) ["-jar", springBootShadowJar, "--server.port=${httpPort}"])