Skip to content

Commit 7c044fd

Browse files
sezen-datadogsezen-datadog
committed
wip
Signed-off-by: sezen.leblay <sezen.leblay@datadoghq.com>
1 parent 19a036c commit 7c044fd

File tree

4 files changed

+133
-12
lines changed

4 files changed

+133
-12
lines changed

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,8 @@ public void init() {
135135
subscriptionService.registerCallback(EVENTS.requestMethodUriRaw(), this::onRequestMethodUriRaw);
136136
subscriptionService.registerCallback(EVENTS.requestBodyStart(), this::onRequestBodyStart);
137137
subscriptionService.registerCallback(EVENTS.requestBodyDone(), this::onRequestBodyDone);
138+
subscriptionService.registerCallback(EVENTS.responseBodyStart(), this::onResponseBodyStart);
139+
subscriptionService.registerCallback(EVENTS.responseBodyDone(), this::onResponseBodyDone);
138140
subscriptionService.registerCallback(
139141
EVENTS.requestClientSocketAddress(), this::onRequestClientSocketAddress);
140142
subscriptionService.registerCallback(
@@ -164,6 +166,10 @@ public void init() {
164166
subscriptionService.registerCallback(
165167
EVENTS.requestBodyProcessed(), this::onRequestBodyProcessed);
166168
}
169+
if (additionalIGEvents.contains(EVENTS.responseBodyProcessed())) {
170+
subscriptionService.registerCallback(
171+
EVENTS.responseBodyProcessed(), this::onResponseBodyProcessed);
172+
}
167173
}
168174

169175
/**
@@ -596,6 +602,51 @@ private Flow<Void> onRequestBodyProcessed(RequestContext ctx_, Object obj) {
596602
}
597603
}
598604

605+
private Flow<Void> onResponseBodyProcessed(RequestContext ctx_, Object obj) {
606+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
607+
if (ctx == null) {
608+
return NoopFlow.INSTANCE;
609+
}
610+
611+
if (ctx.isConvertedResBodyPublished()) {
612+
log.debug(
613+
"Response body already published; will ignore new value of type {}", obj.getClass());
614+
return NoopFlow.INSTANCE;
615+
}
616+
ctx.setConvertedResBodyPublished(true);
617+
618+
while (true) {
619+
DataSubscriberInfo subInfo = responseBodySubInfo;
620+
if (subInfo == null) {
621+
subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT);
622+
responseBodySubInfo = subInfo;
623+
}
624+
if (subInfo == null || subInfo.isEmpty()) {
625+
return NoopFlow.INSTANCE;
626+
}
627+
Object converted =
628+
ObjectIntrospection.convert(
629+
obj,
630+
ctx,
631+
() -> {
632+
if (Config.get().isAppSecRaspCollectRequestBody()) {
633+
ctx_.getTraceSegment()
634+
.setTagTop("_dd.appsec.rasp.response_body_size.exceeded", true);
635+
}
636+
});
637+
if (Config.get().isAppSecRaspCollectResponseBody()) {
638+
ctx.setProcessedResponseBody(converted);
639+
}
640+
DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.RESPONSE_BODY_OBJECT, converted);
641+
try {
642+
GatewayContext gwCtx = new GatewayContext(false);
643+
return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
644+
} catch (ExpiredSubscriberInfoException e) {
645+
responseBodySubInfo = null;
646+
}
647+
}
648+
}
649+
599650
private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier supplier) {
600651
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
601652
if (ctx == null || ctx.isRawReqBodyPublished()) {
@@ -627,6 +678,37 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
627678
}
628679
}
629680

681+
private Flow<Void> onResponseBodyDone(RequestContext ctx_, StoredBodySupplier supplier) {
682+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
683+
if (ctx == null || ctx.isRawResBodyPublished()) {
684+
return NoopFlow.INSTANCE;
685+
}
686+
ctx.setRawResBodyPublished(true);
687+
688+
while (true) {
689+
DataSubscriberInfo subInfo = rawResponseBodySubInfo;
690+
if (subInfo == null) {
691+
subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_RAW);
692+
rawResponseBodySubInfo = subInfo;
693+
}
694+
if (subInfo == null || subInfo.isEmpty()) {
695+
return NoopFlow.INSTANCE;
696+
}
697+
698+
CharSequence bodyContent = supplier.get();
699+
if (bodyContent == null || bodyContent.length() == 0) {
700+
return NoopFlow.INSTANCE;
701+
}
702+
DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.RESPONSE_BODY_RAW, bodyContent);
703+
try {
704+
GatewayContext gwCtx = new GatewayContext(false);
705+
return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
706+
} catch (ExpiredSubscriberInfoException e) {
707+
rawResponseBodySubInfo = null;
708+
}
709+
}
710+
}
711+
630712
private Flow<Void> onRequestPathParams(RequestContext ctx_, Map<String, ?> data) {
631713
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
632714
if (ctx == null || ctx.isPathParamsPublished()) {
@@ -663,6 +745,16 @@ private Void onRequestBodyStart(RequestContext ctx_, StoredBodySupplier supplier
663745
return null;
664746
}
665747

748+
private Void onResponseBodyStart(RequestContext ctx_, StoredBodySupplier supplier) {
749+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
750+
if (ctx == null) {
751+
return null;
752+
}
753+
754+
ctx.setStoredResponseBodySupplier(supplier);
755+
return null;
756+
}
757+
666758
private Flow<AppSecRequestContext> onRequestStarted() {
667759
if (!AppSecSystem.isActive()) {
668760
return RequestContextSupplier.EMPTY;

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HttpMessageConverterInstrumentation.java

Lines changed: 1 addition & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -137,20 +137,9 @@ public static void before(
137137
return;
138138
}
139139

140-
// TODO: A dedicated responseBodyProcessed event should be added to the Events class
141-
// For now, we're using requestBodyProcessed as a placeholder, but this is not semantically
142-
// correct
143-
// The proper solution would be to:
144-
// 1. Add RESPONSE_BODY_PROCESSED_ID = 26 to Events.java
145-
// 2. Add corresponding responseBodyProcessed() method to Events.java
146-
// 3. Update GatewayBridge to handle the new event type
147-
// 4. Replace the callback below with the proper responseBodyProcessed event
148-
149140
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
150141
BiFunction<RequestContext, Object, Flow<Void>> callback =
151-
cbp.getCallback(
152-
EVENTS
153-
.requestBodyProcessed()); // TEMPORARY: Using requestBodyProcessed as placeholder
142+
cbp.getCallback(EVENTS.responseBodyProcessed());
154143
if (callback == null) {
155144
return;
156145
}

‎internal-api/src/main/java/datadog/trace/api/gateway/Events.java

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -312,6 +312,40 @@ public EventType<BiFunction<RequestContext, String, Flow<Void>>> shellCmd() {
312312
return (EventType<BiFunction<RequestContext, String, Flow<Void>>>) SHELL_CMD;
313313
}
314314

315+
static final int RESPONSE_BODY_START_ID = 26;
316+
317+
@SuppressWarnings("rawtypes")
318+
private static final EventType RESPONSE_BODY_START =
319+
new ET<>("response.body.started", RESPONSE_BODY_START_ID);
320+
/** The response body has started being read */
321+
@SuppressWarnings("unchecked")
322+
public EventType<BiFunction<RequestContext, StoredBodySupplier, Void>> responseBodyStart() {
323+
return (EventType<BiFunction<RequestContext, StoredBodySupplier, Void>>) RESPONSE_BODY_START;
324+
}
325+
326+
static final int RESPONSE_BODY_DONE_ID = 27;
327+
328+
@SuppressWarnings("rawtypes")
329+
private static final EventType RESPONSE_BODY_DONE =
330+
new ET<>("response.body.done", RESPONSE_BODY_DONE_ID);
331+
/** The response body is done being read */
332+
@SuppressWarnings("unchecked")
333+
public EventType<BiFunction<RequestContext, StoredBodySupplier, Flow<Void>>> responseBodyDone() {
334+
return (EventType<BiFunction<RequestContext, StoredBodySupplier, Flow<Void>>>)
335+
RESPONSE_BODY_DONE;
336+
}
337+
338+
static final int RESPONSE_BODY_CONVERTED_ID = 28;
339+
340+
@SuppressWarnings("rawtypes")
341+
private static final EventType RESPONSE_BODY_CONVERTED =
342+
new ET<>("response.body.converted", RESPONSE_BODY_CONVERTED_ID);
343+
/** The response body has been converted by the framework */
344+
@SuppressWarnings("unchecked")
345+
public EventType<BiFunction<RequestContext, Object, Flow<Void>>> responseBodyProcessed() {
346+
return (EventType<BiFunction<RequestContext, Object, Flow<Void>>>) RESPONSE_BODY_CONVERTED;
347+
}
348+
315349
static final int MAX_EVENTS = nextId.get();
316350

317351
private static final class ET<T> extends EventType<T> {

‎internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@
2222
import static datadog.trace.api.gateway.Events.REQUEST_PATH_PARAMS_ID;
2323
import static datadog.trace.api.gateway.Events.REQUEST_SESSION_ID;
2424
import static datadog.trace.api.gateway.Events.REQUEST_STARTED_ID;
25+
import static datadog.trace.api.gateway.Events.RESPONSE_BODY_CONVERTED_ID;
26+
import static datadog.trace.api.gateway.Events.RESPONSE_BODY_DONE_ID;
27+
import static datadog.trace.api.gateway.Events.RESPONSE_BODY_START_ID;
2528
import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_DONE_ID;
2629
import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_ID;
2730
import static datadog.trace.api.gateway.Events.RESPONSE_STARTED_ID;
@@ -316,6 +319,7 @@ public boolean equals(Object obj) {
316319
}
317320
};
318321
case REQUEST_BODY_START_ID:
322+
case RESPONSE_BODY_START_ID:
319323
return (C)
320324
new BiFunction<RequestContext, StoredBodySupplier, Void>() {
321325
@Override
@@ -330,6 +334,7 @@ public Void apply(RequestContext ctx, StoredBodySupplier storedBodySupplier) {
330334
}
331335
};
332336
case REQUEST_BODY_DONE_ID:
337+
case RESPONSE_BODY_DONE_ID:
333338
return (C)
334339
new BiFunction<RequestContext, StoredBodySupplier, Flow<Void>>() {
335340
@Override
@@ -346,6 +351,7 @@ public Flow<Void> apply(RequestContext ctx, StoredBodySupplier storedBodySupplie
346351
case GRPC_SERVER_REQUEST_MESSAGE_ID:
347352
case GRAPHQL_SERVER_REQUEST_MESSAGE_ID:
348353
case REQUEST_BODY_CONVERTED_ID:
354+
case RESPONSE_BODY_CONVERTED_ID:
349355
return (C)
350356
new BiFunction<RequestContext, Object, Flow<Void>>() {
351357
@Override

0 commit comments

Comments
 (0)