Skip to content

Commit 6e49d48

Browse files
sezen-datadogsezen-datadog
committed
HTTP response schema collection and data classification
1 parent cb7783b commit 6e49d48

File tree

14 files changed

+1054
-5
lines changed

14 files changed

+1054
-5
lines changed

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,14 +98,17 @@ 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;
104105

105106
private boolean reqDataPublished;
106107
private boolean rawReqBodyPublished;
107108
private boolean convertedReqBodyPublished;
109+
private boolean convertedResBodyPublished;
108110
private boolean respDataPublished;
111+
private boolean rawResBodyPublished;
109112
private boolean pathParamsPublished;
110113
private volatile Map<String, String> derivatives;
111114

@@ -454,6 +457,10 @@ void setStoredRequestBodySupplier(StoredBodySupplier storedRequestBodySupplier)
454457
this.storedRequestBodySupplier = storedRequestBodySupplier;
455458
}
456459

460+
void setStoredResponseBodySupplier(StoredBodySupplier storedResponseBodySupplier) {
461+
this.storedResponseBodySupplier = storedResponseBodySupplier;
462+
}
463+
457464
public String getDbType() {
458465
return dbType;
459466
}
@@ -498,18 +505,34 @@ public boolean isConvertedReqBodyPublished() {
498505
return convertedReqBodyPublished;
499506
}
500507

508+
public boolean isConvertedResBodyPublished() {
509+
return convertedResBodyPublished;
510+
}
511+
501512
public void setConvertedReqBodyPublished(boolean convertedReqBodyPublished) {
502513
this.convertedReqBodyPublished = convertedReqBodyPublished;
503514
}
504515

516+
public void setConvertedResBodyPublished(boolean convertedResBodyPublished) {
517+
this.convertedResBodyPublished = convertedResBodyPublished;
518+
}
519+
505520
public boolean isRespDataPublished() {
506521
return respDataPublished;
507522
}
508523

524+
public boolean isRawResBodyPublished() {
525+
return rawResBodyPublished;
526+
}
527+
509528
public void setRespDataPublished(boolean respDataPublished) {
510529
this.respDataPublished = respDataPublished;
511530
}
512531

532+
public void setRawResBodyPublished(boolean rawResBodyPublished) {
533+
this.rawResBodyPublished = rawResBodyPublished;
534+
}
535+
513536
/**
514537
* Updates the current used usr.id
515538
*
@@ -583,6 +606,14 @@ public CharSequence getStoredRequestBody() {
583606
return storedRequestBodySupplier.get();
584607
}
585608

609+
/** @return the contents of stream */
610+
public CharSequence getStoredResponseBody() {
611+
CharSequence storedResponseBody = null;
612+
storedResponseBody =
613+
this.storedResponseBodySupplier != null ? this.storedResponseBodySupplier.toString() : null;
614+
return storedResponseBody;
615+
}
616+
586617
public void reportEvents(Collection<AppSecEvent> appSecEvents) {
587618
for (AppSecEvent event : appSecEvents) {
588619
StandardizedLogging.attackDetected(log, event);

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

Lines changed: 95 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,9 @@ public class GatewayBridge {
9595
// subscriber cache
9696
private volatile DataSubscriberInfo initialReqDataSubInfo;
9797
private volatile DataSubscriberInfo rawRequestBodySubInfo;
98+
private volatile DataSubscriberInfo rawResponseBodySubInfo;
9899
private volatile DataSubscriberInfo requestBodySubInfo;
100+
private volatile DataSubscriberInfo responseBodySubInfo;
99101
private volatile DataSubscriberInfo pathParamsSubInfo;
100102
private volatile DataSubscriberInfo respDataSubInfo;
101103
private volatile DataSubscriberInfo grpcServerMethodSubInfo;
@@ -142,6 +144,8 @@ public void init() {
142144
subscriptionService.registerCallback(EVENTS.responseStarted(), this::onResponseStarted);
143145
subscriptionService.registerCallback(EVENTS.responseHeader(), this::onResponseHeader);
144146
subscriptionService.registerCallback(EVENTS.responseHeaderDone(), this::onResponseHeaderDone);
147+
subscriptionService.registerCallback(EVENTS.responseBodyStart(), this::onResponseBodyStart);
148+
subscriptionService.registerCallback(EVENTS.responseBodyDone(), this::onResponseBodyDone);
145149
subscriptionService.registerCallback(EVENTS.grpcServerMethod(), this::onGrpcServerMethod);
146150
subscriptionService.registerCallback(
147151
EVENTS.grpcServerRequestMessage(), this::onGrpcServerRequestMessage);
@@ -164,6 +168,10 @@ public void init() {
164168
subscriptionService.registerCallback(
165169
EVENTS.requestBodyProcessed(), this::onRequestBodyProcessed);
166170
}
171+
if (additionalIGEvents.contains(EVENTS.responseBodyProcessed())) {
172+
subscriptionService.registerCallback(
173+
EVENTS.responseBodyProcessed(), this::onResponseBodyProcessed);
174+
}
167175
}
168176

169177
/**
@@ -173,7 +181,9 @@ public void init() {
173181
public void reset() {
174182
initialReqDataSubInfo = null;
175183
rawRequestBodySubInfo = null;
184+
rawResponseBodySubInfo = null;
176185
requestBodySubInfo = null;
186+
responseBodySubInfo = null;
177187
pathParamsSubInfo = null;
178188
respDataSubInfo = null;
179189
grpcServerMethodSubInfo = null;
@@ -596,6 +606,40 @@ private Flow<Void> onRequestBodyProcessed(RequestContext ctx_, Object obj) {
596606
}
597607
}
598608

609+
private Flow<Void> onResponseBodyProcessed(RequestContext ctx_, Object obj) {
610+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
611+
if (ctx == null) {
612+
return NoopFlow.INSTANCE;
613+
}
614+
615+
if (ctx.isConvertedResBodyPublished()) {
616+
log.debug(
617+
"Response body already published; will ignore new value of type {}", obj.getClass());
618+
return NoopFlow.INSTANCE;
619+
}
620+
ctx.setConvertedResBodyPublished(true);
621+
622+
while (true) {
623+
DataSubscriberInfo subInfo = responseBodySubInfo;
624+
if (subInfo == null) {
625+
subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT);
626+
responseBodySubInfo = subInfo;
627+
}
628+
if (subInfo == null || subInfo.isEmpty()) {
629+
return NoopFlow.INSTANCE;
630+
}
631+
DataBundle bundle =
632+
new SingletonDataBundle<>(
633+
KnownAddresses.RESPONSE_BODY_OBJECT, ObjectIntrospection.convert(obj, ctx));
634+
try {
635+
GatewayContext gwCtx = new GatewayContext(false);
636+
return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
637+
} catch (ExpiredSubscriberInfoException e) {
638+
responseBodySubInfo = null;
639+
}
640+
}
641+
}
642+
599643
private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier supplier) {
600644
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
601645
if (ctx == null || ctx.isRawReqBodyPublished()) {
@@ -614,7 +658,7 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
614658
}
615659

616660
CharSequence bodyContent = supplier.get();
617-
if (bodyContent == null || bodyContent.length() == 0) {
661+
if (bodyContent.length() == 0) {
618662
return NoopFlow.INSTANCE;
619663
}
620664
DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.REQUEST_BODY_RAW, bodyContent);
@@ -627,6 +671,38 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
627671
}
628672
}
629673

674+
private Flow<Void> onResponseBodyDone(RequestContext ctx_, StoredBodySupplier supplier) {
675+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
676+
if (ctx == null || ctx.isRawResBodyPublished()) {
677+
return NoopFlow.INSTANCE;
678+
}
679+
ctx.setRawResBodyPublished(true);
680+
681+
while (true) {
682+
DataSubscriberInfo subInfo = responseBodySubInfo;
683+
if (subInfo == null) {
684+
subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT);
685+
responseBodySubInfo = subInfo;
686+
}
687+
if (subInfo == null || subInfo.isEmpty()) {
688+
return NoopFlow.INSTANCE;
689+
}
690+
691+
CharSequence bodyContent = supplier.toString();
692+
if (bodyContent.length() == 0) {
693+
return NoopFlow.INSTANCE;
694+
}
695+
DataBundle bundle =
696+
new SingletonDataBundle<>(KnownAddresses.RESPONSE_BODY_OBJECT, bodyContent);
697+
try {
698+
GatewayContext gwCtx = new GatewayContext(false);
699+
return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
700+
} catch (ExpiredSubscriberInfoException e) {
701+
responseBodySubInfo = null;
702+
}
703+
}
704+
}
705+
630706
private Flow<Void> onRequestPathParams(RequestContext ctx_, Map<String, ?> data) {
631707
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
632708
if (ctx == null || ctx.isPathParamsPublished()) {
@@ -663,6 +739,16 @@ private Void onRequestBodyStart(RequestContext ctx_, StoredBodySupplier supplier
663739
return null;
664740
}
665741

742+
private Void onResponseBodyStart(RequestContext ctx_, StoredBodySupplier supplier) {
743+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
744+
if (ctx == null) {
745+
return null;
746+
}
747+
748+
ctx.setStoredResponseBodySupplier(supplier);
749+
return null;
750+
}
751+
666752
private Flow<AppSecRequestContext> onRequestStarted() {
667753
if (!AppSecSystem.isActive()) {
668754
return RequestContextSupplier.EMPTY;
@@ -1032,8 +1118,10 @@ private Flow<Void> maybePublishResponseData(AppSecRequestContext ctx) {
10321118

10331119
MapDataBundle bundle =
10341120
MapDataBundle.of(
1035-
KnownAddresses.RESPONSE_STATUS, String.valueOf(ctx.getResponseStatus()),
1036-
KnownAddresses.RESPONSE_HEADERS_NO_COOKIES, ctx.getResponseHeaders());
1121+
KnownAddresses.RESPONSE_STATUS,
1122+
String.valueOf(ctx.getResponseStatus()),
1123+
KnownAddresses.RESPONSE_HEADERS_NO_COOKIES,
1124+
ctx.getResponseHeaders());
10371125

10381126
while (true) {
10391127
DataSubscriberInfo subInfo = respDataSubInfo;
@@ -1128,6 +1216,10 @@ private static class IGAppSecEventDependencies {
11281216
KnownAddresses.REQUEST_BODY_RAW, l(EVENTS.requestBodyStart(), EVENTS.requestBodyDone()));
11291217
DATA_DEPENDENCIES.put(KnownAddresses.REQUEST_PATH_PARAMS, l(EVENTS.requestPathParams()));
11301218
DATA_DEPENDENCIES.put(KnownAddresses.REQUEST_BODY_OBJECT, l(EVENTS.requestBodyProcessed()));
1219+
DATA_DEPENDENCIES.put(
1220+
KnownAddresses.RESPONSE_BODY_RAW,
1221+
l(EVENTS.responseBodyStart(), EVENTS.responseBodyDone()));
1222+
DATA_DEPENDENCIES.put(KnownAddresses.RESPONSE_BODY_OBJECT, l(EVENTS.responseBodyProcessed()));
11311223
}
11321224

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

‎dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import com.datadog.appsec.report.AppSecEvent
1111
import com.datadog.appsec.report.AppSecEventWrapper
1212
import datadog.trace.api.ProductTraceSource
1313
import datadog.trace.api.config.GeneralConfig
14-
import static datadog.trace.api.config.IastConfig.IAST_DEDUPLICATION_ENABLED
1514
import datadog.trace.api.function.TriConsumer
1615
import datadog.trace.api.function.TriFunction
1716
import datadog.trace.api.gateway.BlockResponseFunction
@@ -114,6 +113,8 @@ class GatewayBridgeSpecification extends DDSpecification {
114113
BiFunction<RequestContext, String, Flow<Void>> shellCmdCB
115114
BiFunction<RequestContext, String, Flow<Void>> userCB
116115
TriFunction<RequestContext, LoginEvent, String, Flow<Void>> loginEventCB
116+
BiFunction<RequestContext, StoredBodySupplier, Void> responseBodyStartCB
117+
BiFunction<RequestContext, StoredBodySupplier, Flow<Void>> responseBodyDoneCB
117118

118119
WafMetricCollector wafMetricCollector = Mock(WafMetricCollector)
119120

@@ -452,6 +453,8 @@ class GatewayBridgeSpecification extends DDSpecification {
452453
1 * ig.registerCallback(EVENTS.responseStarted(), _) >> { responseStartedCB = it[1]; null }
453454
1 * ig.registerCallback(EVENTS.responseHeader(), _) >> { respHeaderCB = it[1]; null }
454455
1 * ig.registerCallback(EVENTS.responseHeaderDone(), _) >> { respHeadersDoneCB = it[1]; null }
456+
1 * ig.registerCallback(EVENTS.responseBodyStart(), _) >> { responseBodyStartCB = it[1]; null }
457+
1 * ig.registerCallback(EVENTS.responseBodyDone(), _) >> { responseBodyDoneCB = it[1]; null }
455458
1 * ig.registerCallback(EVENTS.grpcServerMethod(), _) >> { grpcServerMethodCB = it[1]; null }
456459
1 * ig.registerCallback(EVENTS.grpcServerRequestMessage(), _) >> { grpcServerRequestMessageCB = it[1]; null }
457460
1 * ig.registerCallback(EVENTS.graphqlServerRequestMessage(), _) >> { graphqlServerRequestMessageCB = it[1]; null }
@@ -991,7 +994,7 @@ class GatewayBridgeSpecification extends DDSpecification {
991994
getTraceSegment() >> traceSegment
992995
}
993996
final spanInfo = Mock(AgentSpan) {
994-
getTags() >> ['http.route':'/']
997+
getTags() >> ['http.route': '/']
995998
}
996999

9971000
when:

0 commit comments

Comments
 (0)