Skip to content

Commit a2809f5

Browse files
sezen-datadogsezen-datadog
committed
HTTP response schema collection and data classification
1 parent 50e8b5a commit a2809f5

File tree

7 files changed

+115
-7
lines changed

7 files changed

+115
-7
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import datadog.trace.api.gateway.IGSpanInfo;
1919
import datadog.trace.api.gateway.RequestContext;
2020
import datadog.trace.api.gateway.RequestContextSlot;
21+
import datadog.trace.api.http.StoredBodySupplier;
2122
import datadog.trace.api.naming.SpanNaming;
2223
import datadog.trace.bootstrap.ActiveSubsystems;
2324
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
@@ -435,7 +436,8 @@ private Flow<Void> callIGCallbackRequestHeaders(AgentSpan span, REQUEST_CARRIER
435436
IGKeyClassifier.create(
436437
requestContext,
437438
cbp.getCallback(EVENTS.requestHeader()),
438-
cbp.getCallback(EVENTS.requestHeaderDone()));
439+
cbp.getCallback(EVENTS.requestHeaderDone()),
440+
null);
439441
if (null != igKeyClassifier) {
440442
getter.forEachKey(carrier, igKeyClassifier);
441443
return igKeyClassifier.done();
@@ -491,7 +493,8 @@ public <RESP> Flow<Void> callIGCallbackResponseAndHeaders(
491493
IGKeyClassifier.create(
492494
requestContext,
493495
cbp.getCallback(EVENTS.responseHeader()),
494-
cbp.getCallback(EVENTS.responseHeaderDone()));
496+
cbp.getCallback(EVENTS.responseHeaderDone()),
497+
cbp.getCallback(EVENTS.responseBodyDone()));
495498
if (null != igKeyClassifier) {
496499
contextVisitor.forEachKey(carrier, igKeyClassifier);
497500
return igKeyClassifier.done();
@@ -575,7 +578,8 @@ protected static final class IGKeyClassifier implements AgentPropagation.KeyClas
575578
public static IGKeyClassifier create(
576579
RequestContext requestContext,
577580
TriConsumer<RequestContext, String, String> headerCallback,
578-
Function<RequestContext, Flow<Void>> doneCallback) {
581+
Function<RequestContext, Flow<Void>> doneCallback,
582+
BiFunction<RequestContext, StoredBodySupplier, Flow<Void>> bodyDoneCallback) {
579583
if (null == requestContext || null == headerCallback) {
580584
return null;
581585
}

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ public interface KnownAddresses {
4848
Address<Object> RESPONSE_BODY_OBJECT = new Address<>("server.response.body");
4949

5050
/** First chars of HTTP response body */
51-
Address<String> RESPONSE_BODY_RAW = new Address<>("server.response.body.raw");
51+
Address<CharSequence> RESPONSE_BODY_RAW = new Address<>("server.response.body.raw");
5252

5353
/** Reponse headers excluding cookies */
5454
Address<Map<String, List<String>>> RESPONSE_HEADERS_NO_COOKIES =

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
9898
private String inferredClientIp;
9999

100100
private volatile StoredBodySupplier storedRequestBodySupplier;
101+
private volatile StoredBodySupplier storedResponseBodySupplier;
101102
private String dbType;
102103

103104
private int responseStatus;
@@ -106,6 +107,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
106107
private boolean rawReqBodyPublished;
107108
private boolean convertedReqBodyPublished;
108109
private boolean respDataPublished;
110+
private boolean rawResBodyPublished;
109111
private boolean pathParamsPublished;
110112
private volatile Map<String, String> derivatives;
111113

@@ -451,6 +453,10 @@ void setStoredRequestBodySupplier(StoredBodySupplier storedRequestBodySupplier)
451453
this.storedRequestBodySupplier = storedRequestBodySupplier;
452454
}
453455

456+
void setStoredResponseBodySupplier(StoredBodySupplier storedResponseBodySupplier) {
457+
this.storedResponseBodySupplier = storedResponseBodySupplier;
458+
}
459+
454460
public String getDbType() {
455461
return dbType;
456462
}
@@ -503,10 +509,18 @@ public boolean isRespDataPublished() {
503509
return respDataPublished;
504510
}
505511

512+
public boolean isRawResBodyPublished() {
513+
return rawResBodyPublished;
514+
}
515+
506516
public void setRespDataPublished(boolean respDataPublished) {
507517
this.respDataPublished = respDataPublished;
508518
}
509519

520+
public void setRawResBodyPublished(boolean rawResBodyPublished) {
521+
this.rawResBodyPublished = rawResBodyPublished;
522+
}
523+
510524
/**
511525
* Updates the current used usr.id
512526
*
@@ -580,6 +594,15 @@ public CharSequence getStoredRequestBody() {
580594
return storedRequestBodySupplier.get();
581595
}
582596

597+
/** @return the portion of the body read so far, if any */
598+
public CharSequence getStoredResponseBody() {
599+
StoredBodySupplier storedResponseBodySupplier = this.storedResponseBodySupplier;
600+
if (storedResponseBodySupplier == null) {
601+
return null;
602+
}
603+
return storedResponseBodySupplier.get();
604+
}
605+
583606
public void reportEvents(Collection<AppSecEvent> appSecEvents) {
584607
for (AppSecEvent event : appSecEvents) {
585608
StandardizedLogging.attackDetected(log, event);

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

Lines changed: 55 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ public class GatewayBridge {
9494
// subscriber cache
9595
private volatile DataSubscriberInfo initialReqDataSubInfo;
9696
private volatile DataSubscriberInfo rawRequestBodySubInfo;
97+
private volatile DataSubscriberInfo rawResponseBodySubInfo;
9798
private volatile DataSubscriberInfo requestBodySubInfo;
9899
private volatile DataSubscriberInfo pathParamsSubInfo;
99100
private volatile DataSubscriberInfo respDataSubInfo;
@@ -141,6 +142,8 @@ public void init() {
141142
subscriptionService.registerCallback(EVENTS.responseStarted(), this::onResponseStarted);
142143
subscriptionService.registerCallback(EVENTS.responseHeader(), this::onResponseHeader);
143144
subscriptionService.registerCallback(EVENTS.responseHeaderDone(), this::onResponseHeaderDone);
145+
subscriptionService.registerCallback(EVENTS.responseBodyStart(), this::onResponseBodyStart);
146+
subscriptionService.registerCallback(EVENTS.responseBodyDone(), this::onResponseBodyDone);
144147
subscriptionService.registerCallback(EVENTS.grpcServerMethod(), this::onGrpcServerMethod);
145148
subscriptionService.registerCallback(
146149
EVENTS.grpcServerRequestMessage(), this::onGrpcServerRequestMessage);
@@ -602,7 +605,7 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
602605
}
603606

604607
CharSequence bodyContent = supplier.get();
605-
if (bodyContent == null || bodyContent.length() == 0) {
608+
if (bodyContent.length() == 0) {
606609
return NoopFlow.INSTANCE;
607610
}
608611
DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.REQUEST_BODY_RAW, bodyContent);
@@ -615,6 +618,38 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
615618
}
616619
}
617620

621+
private Flow<Void> onResponseBodyDone(RequestContext ctx_, StoredBodySupplier supplier) {
622+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
623+
if (ctx == null || ctx.isRawResBodyPublished()) {
624+
return NoopFlow.INSTANCE;
625+
}
626+
ctx.setRawResBodyPublished(true);
627+
628+
while (true) {
629+
DataSubscriberInfo subInfo = rawResponseBodySubInfo;
630+
if (subInfo == null) {
631+
subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_RAW);
632+
rawResponseBodySubInfo = subInfo;
633+
}
634+
if (subInfo == null || subInfo.isEmpty()) {
635+
return NoopFlow.INSTANCE;
636+
}
637+
638+
CharSequence bodyContent = supplier.get();
639+
if (bodyContent.length() == 0) {
640+
return NoopFlow.INSTANCE;
641+
}
642+
DataBundle bundle =
643+
new SingletonDataBundle<>(KnownAddresses.RESPONSE_BODY_OBJECT, bodyContent);
644+
try {
645+
GatewayContext gwCtx = new GatewayContext(false);
646+
return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
647+
} catch (ExpiredSubscriberInfoException e) {
648+
rawResponseBodySubInfo = null;
649+
}
650+
}
651+
}
652+
618653
private Flow<Void> onRequestPathParams(RequestContext ctx_, Map<String, ?> data) {
619654
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
620655
if (ctx == null || ctx.isPathParamsPublished()) {
@@ -651,6 +686,16 @@ private Void onRequestBodyStart(RequestContext ctx_, StoredBodySupplier supplier
651686
return null;
652687
}
653688

689+
private Void onResponseBodyStart(RequestContext ctx_, StoredBodySupplier supplier) {
690+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
691+
if (ctx == null) {
692+
return null;
693+
}
694+
695+
ctx.setStoredResponseBodySupplier(supplier);
696+
return null;
697+
}
698+
654699
private Flow<AppSecRequestContext> onRequestStarted() {
655700
if (!AppSecSystem.isActive()) {
656701
return RequestContextSupplier.EMPTY;
@@ -962,8 +1007,12 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
9621007

9631008
MapDataBundle bundle =
9641009
MapDataBundle.of(
965-
KnownAddresses.RESPONSE_STATUS, String.valueOf(ctx.getResponseStatus()),
966-
KnownAddresses.RESPONSE_HEADERS_NO_COOKIES, ctx.getResponseHeaders());
1010+
KnownAddresses.RESPONSE_STATUS,
1011+
String.valueOf(ctx.getResponseStatus()),
1012+
KnownAddresses.RESPONSE_HEADERS_NO_COOKIES,
1013+
ctx.getResponseHeaders(),
1014+
KnownAddresses.RESPONSE_BODY_OBJECT,
1015+
ctx.getStoredResponseBody());
9671016

9681017
while (true) {
9691018
DataSubscriberInfo subInfo = respDataSubInfo;
@@ -1058,6 +1107,9 @@ private static class IGAppSecEventDependencies {
10581107
KnownAddresses.REQUEST_BODY_RAW, l(EVENTS.requestBodyStart(), EVENTS.requestBodyDone()));
10591108
DATA_DEPENDENCIES.put(KnownAddresses.REQUEST_PATH_PARAMS, l(EVENTS.requestPathParams()));
10601109
DATA_DEPENDENCIES.put(KnownAddresses.REQUEST_BODY_OBJECT, l(EVENTS.requestBodyProcessed()));
1110+
DATA_DEPENDENCIES.put(
1111+
KnownAddresses.RESPONSE_BODY_RAW,
1112+
l(EVENTS.responseBodyStart(), EVENTS.responseBodyDone()));
10611113
}
10621114

10631115
private static Collection<datadog.trace.api.gateway.EventType<?>> l(

‎dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,8 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest {
219219

220220
List<String> command = new ArrayList<>()
221221
command.add(javaPath())
222+
// TODO delete when ure done testing
223+
command.add("-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=5005")
222224
command.addAll(defaultJavaProperties)
223225
command.addAll(defaultAppSecProperties)
224226
command.addAll((String[]) ["-jar", springBootShadowJar, "--server.port=${httpPort}"])

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -312,6 +312,29 @@ 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 request 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 request 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+
315338
static final int MAX_EVENTS = nextId.get();
316339

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

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
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_DONE_ID;
26+
import static datadog.trace.api.gateway.Events.RESPONSE_BODY_START_ID;
2527
import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_DONE_ID;
2628
import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_ID;
2729
import static datadog.trace.api.gateway.Events.RESPONSE_STARTED_ID;
@@ -315,6 +317,7 @@ public boolean equals(Object obj) {
315317
return callback.equals(obj);
316318
}
317319
};
320+
case RESPONSE_BODY_START_ID:
318321
case REQUEST_BODY_START_ID:
319322
return (C)
320323
new BiFunction<RequestContext, StoredBodySupplier, Void>() {
@@ -329,6 +332,7 @@ public Void apply(RequestContext ctx, StoredBodySupplier storedBodySupplier) {
329332
}
330333
}
331334
};
335+
case RESPONSE_BODY_DONE_ID:
332336
case REQUEST_BODY_DONE_ID:
333337
return (C)
334338
new BiFunction<RequestContext, StoredBodySupplier, Flow<Void>>() {

0 commit comments

Comments
 (0)