Skip to content

Commit ac7c355

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Extract Vert.x json body response schemas
1 parent faeb62c commit ac7c355

File tree

17 files changed

+373
-16
lines changed

17 files changed

+373
-16
lines changed

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

Lines changed: 34 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import java.lang.reflect.Method;
1212
import java.lang.reflect.Modifier;
1313
import java.util.ArrayList;
14+
import java.util.Collection;
1415
import java.util.HashMap;
1516
import java.util.Iterator;
1617
import java.util.List;
@@ -211,20 +212,42 @@ private static Object doConversion(Object obj, int depth, State state) {
211212

212213
// iterables
213214
if (obj instanceof Iterable) {
214-
List<Object> newList;
215-
if (obj instanceof List) {
216-
newList = new ArrayList<>(((List<?>) obj).size());
215+
final Iterator<?> it = ((Iterable<?>) obj).iterator();
216+
final boolean isMap = it.hasNext() && it.next() instanceof Map.Entry;
217+
// some json libraries implement objects as Iterable<Map.Entry>
218+
if (isMap) {
219+
Map<Object, Object> newMap;
220+
if (obj instanceof Collection) {
221+
newMap = new HashMap<>(((Collection<?>) obj).size());
222+
} else {
223+
newMap = new HashMap<>();
224+
}
225+
for (Map.Entry<?, ?> e : ((Iterable<Map.Entry<?, ?>>) obj)) {
226+
Object key = e.getKey();
227+
Object newKey = keyConversion(e.getKey(), state);
228+
if (newKey == null && key != null) {
229+
// probably we're out of elements anyway
230+
continue;
231+
}
232+
newMap.put(newKey, guardedConversion(e.getValue(), depth + 1, state));
233+
}
234+
return newMap;
217235
} else {
218-
newList = new ArrayList<>();
219-
}
220-
for (Object o : ((Iterable<?>) obj)) {
221-
if (state.elemsLeft <= 0) {
222-
state.listMapTooLarge = true;
223-
break;
236+
List<Object> newList;
237+
if (obj instanceof Collection) {
238+
newList = new ArrayList<>(((Collection<?>) obj).size());
239+
} else {
240+
newList = new ArrayList<>();
224241
}
225-
newList.add(guardedConversion(o, depth + 1, state));
242+
for (Object o : ((Iterable<?>) obj)) {
243+
if (state.elemsLeft <= 0) {
244+
state.listMapTooLarge = true;
245+
break;
246+
}
247+
newList.add(guardedConversion(o, depth + 1, state));
248+
}
249+
return newList;
226250
}
227-
return newList;
228251
}
229252

230253
// arrays

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
105105
private boolean reqDataPublished;
106106
private boolean rawReqBodyPublished;
107107
private boolean convertedReqBodyPublished;
108+
private boolean responseBodyPublished;
108109
private boolean respDataPublished;
109110
private boolean pathParamsPublished;
110111
private volatile Map<String, String> derivatives;
@@ -502,6 +503,14 @@ public void setConvertedReqBodyPublished(boolean convertedReqBodyPublished) {
502503
this.convertedReqBodyPublished = convertedReqBodyPublished;
503504
}
504505

506+
public boolean isResponseBodyPublished() {
507+
return responseBodyPublished;
508+
}
509+
510+
public void setResponseBodyPublished(final boolean responseBodyPublished) {
511+
this.responseBodyPublished = responseBodyPublished;
512+
}
513+
505514
public boolean isRespDataPublished() {
506515
return respDataPublished;
507516
}

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ public class GatewayBridge {
9898
private volatile DataSubscriberInfo initialReqDataSubInfo;
9999
private volatile DataSubscriberInfo rawRequestBodySubInfo;
100100
private volatile DataSubscriberInfo requestBodySubInfo;
101+
private volatile DataSubscriberInfo responseBodySubInfo;
101102
private volatile DataSubscriberInfo pathParamsSubInfo;
102103
private volatile DataSubscriberInfo respDataSubInfo;
103104
private volatile DataSubscriberInfo grpcServerMethodSubInfo;
@@ -137,6 +138,7 @@ public void init() {
137138
subscriptionService.registerCallback(EVENTS.requestMethodUriRaw(), this::onRequestMethodUriRaw);
138139
subscriptionService.registerCallback(EVENTS.requestBodyStart(), this::onRequestBodyStart);
139140
subscriptionService.registerCallback(EVENTS.requestBodyDone(), this::onRequestBodyDone);
141+
subscriptionService.registerCallback(EVENTS.responseBody(), this::onResponseBody);
140142
subscriptionService.registerCallback(
141143
EVENTS.requestClientSocketAddress(), this::onRequestClientSocketAddress);
142144
subscriptionService.registerCallback(
@@ -177,6 +179,7 @@ public void reset() {
177179
initialReqDataSubInfo = null;
178180
rawRequestBodySubInfo = null;
179181
requestBodySubInfo = null;
182+
responseBodySubInfo = null;
180183
pathParamsSubInfo = null;
181184
respDataSubInfo = null;
182185
grpcServerMethodSubInfo = null;
@@ -638,6 +641,39 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
638641
}
639642
}
640643

644+
private Flow<Void> onResponseBody(RequestContext ctx_, Object obj) {
645+
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
646+
if (ctx == null) {
647+
return NoopFlow.INSTANCE;
648+
}
649+
650+
if (ctx.isResponseBodyPublished()) {
651+
log.debug(
652+
"Response body already published; will ignore new value of type {}", obj.getClass());
653+
return NoopFlow.INSTANCE;
654+
}
655+
ctx.setResponseBodyPublished(true);
656+
657+
while (true) {
658+
DataSubscriberInfo subInfo = responseBodySubInfo;
659+
if (subInfo == null) {
660+
subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT);
661+
responseBodySubInfo = subInfo;
662+
}
663+
if (subInfo == null || subInfo.isEmpty()) {
664+
return NoopFlow.INSTANCE;
665+
}
666+
Object converted = ObjectIntrospection.convert(obj, ctx);
667+
DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.RESPONSE_BODY_OBJECT, converted);
668+
try {
669+
GatewayContext gwCtx = new GatewayContext(false);
670+
return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
671+
} catch (ExpiredSubscriberInfoException e) {
672+
responseBodySubInfo = null;
673+
}
674+
}
675+
}
676+
641677
private Flow<Void> onRequestPathParams(RequestContext ctx_, Map<String, ?> data) {
642678
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
643679
if (ctx == null || ctx.isPathParamsPublished()) {

‎dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/ObjectIntrospectionSpecification.groovy

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,10 @@ package com.datadog.appsec.event.data
22

33
import com.datadog.appsec.gateway.AppSecRequestContext
44
import com.fasterxml.jackson.databind.ObjectMapper
5-
import com.fasterxml.jackson.databind.node.ArrayNode
6-
import com.fasterxml.jackson.databind.node.ObjectNode
75
import datadog.trace.api.telemetry.WafMetricCollector
86
import datadog.trace.test.util.DDSpecification
97
import groovy.json.JsonBuilder
108
import groovy.json.JsonOutput
11-
import groovy.json.JsonSlurper
129
import spock.lang.Shared
1310

1411
import java.nio.CharBuffer
@@ -465,6 +462,17 @@ class ObjectIntrospectionSpecification extends DDSpecification {
465462
MAPPER.readTree('"unicode: \\u0041"') || 'unicode: A'
466463
}
467464

465+
void 'iterable json objects'() {
466+
setup:
467+
final map = [name: 'This is just a test', list: [1, 2, 3, 4, 5]]
468+
469+
when:
470+
final result = convert(new IterableJsonObject(map), ctx)
471+
472+
then:
473+
result == map
474+
}
475+
468476
private static int countNesting(final Map<String, Object>object, final int levels) {
469477
if (object.isEmpty()) {
470478
return levels
@@ -475,4 +483,18 @@ class ObjectIntrospectionSpecification extends DDSpecification {
475483
}
476484
return countNesting(object.values().first() as Map, levels + 1)
477485
}
486+
487+
private static class IterableJsonObject implements Iterable<Map.Entry<String, Object>> {
488+
489+
private final Map<String, Object> map
490+
491+
IterableJsonObject(Map<String, Object> map) {
492+
this.map = map
493+
}
494+
495+
@Override
496+
Iterator<Map.Entry<String, Object>> iterator() {
497+
return map.entrySet().iterator()
498+
}
499+
}
478500
}

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ class GatewayBridgeSpecification extends DDSpecification {
9999
BiFunction<RequestContext, StoredBodySupplier, Void> requestBodyStartCB
100100
BiFunction<RequestContext, StoredBodySupplier, Flow<Void>> requestBodyDoneCB
101101
BiFunction<RequestContext, Object, Flow<Void>> requestBodyProcessedCB
102+
BiFunction<RequestContext, Object, Flow<Void>> responseBodyCB
102103
BiFunction<RequestContext, Integer, Flow<Void>> responseStartedCB
103104
TriConsumer<RequestContext, String, String> respHeaderCB
104105
Function<RequestContext, Flow<Void>> respHeadersDoneCB
@@ -463,6 +464,7 @@ class GatewayBridgeSpecification extends DDSpecification {
463464
1 * ig.registerCallback(EVENTS.requestBodyStart(), _) >> { requestBodyStartCB = it[1]; null }
464465
1 * ig.registerCallback(EVENTS.requestBodyDone(), _) >> { requestBodyDoneCB = it[1]; null }
465466
1 * ig.registerCallback(EVENTS.requestBodyProcessed(), _) >> { requestBodyProcessedCB = it[1]; null }
467+
1 * ig.registerCallback(EVENTS.responseBody(), _) >> { responseBodyCB = it[1]; null }
466468
1 * ig.registerCallback(EVENTS.responseStarted(), _) >> { responseStartedCB = it[1]; null }
467469
1 * ig.registerCallback(EVENTS.responseHeader(), _) >> { respHeaderCB = it[1]; null }
468470
1 * ig.registerCallback(EVENTS.responseHeaderDone(), _) >> { respHeadersDoneCB = it[1]; null }
@@ -1340,4 +1342,17 @@ class GatewayBridgeSpecification extends DDSpecification {
13401342
arCtx.getRoute() == route
13411343
}
13421344
1345+
void 'test on response body callback'() {
1346+
when:
1347+
responseBodyCB.apply(ctx, [test: 'this is a test'])
1348+
1349+
then:
1350+
1 * eventDispatcher.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT) >> nonEmptyDsInfo
1351+
1 * eventDispatcher.publishDataEvent(_, _, _, _) >> {
1352+
final bundle = it[2] as DataBundle
1353+
final body = bundle.get(KnownAddresses.RESPONSE_BODY_OBJECT)
1354+
assert body['test'] == 'this is a test'
1355+
}
1356+
}
1357+
13431358
}

‎dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextInstrumentation.java

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package datadog.trace.instrumentation.vertx_4_0.server;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
5+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
6+
7+
import com.google.auto.service.AutoService;
8+
import datadog.trace.agent.tooling.Instrumenter;
9+
import datadog.trace.agent.tooling.InstrumenterModule;
10+
import datadog.trace.agent.tooling.muzzle.Reference;
11+
import io.vertx.ext.web.impl.RoutingContextImpl;
12+
13+
/**
14+
* @see RoutingContextImpl#getBodyAsJson(int)
15+
* @see RoutingContextImpl#getBodyAsJsonArray(int)
16+
*/
17+
@AutoService(InstrumenterModule.class)
18+
public class RoutingContextInstrumentation extends InstrumenterModule.AppSec
19+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
20+
21+
public RoutingContextInstrumentation() {
22+
super("vertx", "vertx-4.0");
23+
}
24+
25+
@Override
26+
public Reference[] additionalMuzzleReferences() {
27+
return new Reference[] {VertxVersionMatcher.HTTP_1X_SERVER_RESPONSE};
28+
}
29+
30+
@Override
31+
public String instrumentedType() {
32+
return "io.vertx.ext.web.RoutingContext";
33+
}
34+
35+
@Override
36+
public void methodAdvice(MethodTransformer transformer) {
37+
transformer.applyAdvice(
38+
named("json").and(takesArguments(1)).and(takesArgument(0, Object.class)),
39+
packageName + ".RoutingContextJsonResponseAdvice");
40+
}
41+
}

‎dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextJsonResponseAdvice.java

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
package datadog.trace.instrumentation.vertx_4_0.server;
2+
3+
import static datadog.trace.api.gateway.Events.EVENTS;
4+
5+
import datadog.appsec.api.blocking.BlockingException;
6+
import datadog.trace.advice.ActiveRequestContext;
7+
import datadog.trace.advice.RequiresRequestContext;
8+
import datadog.trace.api.gateway.BlockResponseFunction;
9+
import datadog.trace.api.gateway.CallbackProvider;
10+
import datadog.trace.api.gateway.Flow;
11+
import datadog.trace.api.gateway.RequestContext;
12+
import datadog.trace.api.gateway.RequestContextSlot;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
14+
import java.util.function.BiFunction;
15+
import net.bytebuddy.asm.Advice;
16+
17+
@RequiresRequestContext(RequestContextSlot.APPSEC)
18+
class RoutingContextJsonResponseAdvice {
19+
20+
@Advice.OnMethodEnter(suppress = Throwable.class)
21+
static void before(
22+
@Advice.Argument(0) final Object object, @ActiveRequestContext final RequestContext reqCtx) {
23+
24+
if (object == null) {
25+
return;
26+
}
27+
28+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
29+
if (cbp == null) {
30+
return;
31+
}
32+
BiFunction<RequestContext, Object, Flow<Void>> callback =
33+
cbp.getCallback(EVENTS.responseBody());
34+
if (callback == null) {
35+
return;
36+
}
37+
38+
Flow<Void> flow = callback.apply(reqCtx, object);
39+
Flow.Action action = flow.getAction();
40+
if (action instanceof Flow.Action.RequestBlockingAction) {
41+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
42+
if (blockResponseFunction == null) {
43+
return;
44+
}
45+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
46+
blockResponseFunction.tryCommitBlockingResponse(
47+
reqCtx.getTraceSegment(),
48+
rba.getStatusCode(),
49+
rba.getBlockingContentType(),
50+
rba.getExtraHeaders());
51+
52+
throw new BlockingException("Blocked request (for RoutingContext/json)");
53+
}
54+
}
55+
}

‎dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
8383
true
8484
}
8585

86+
@Override
87+
boolean testResponseBodyJson() {
88+
true
89+
}
90+
8691
@Override
8792
boolean testBlocking() {
8893
true

‎dd-java-agent/instrumentation/vertx-web-4.0/src/test/java/server/VertxTestServer.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,8 @@ public void start(final Promise<Void> startPromise) {
127127
BODY_JSON,
128128
() -> {
129129
JsonObject json = ctx.getBodyAsJson();
130-
ctx.response().setStatusCode(BODY_JSON.getStatus()).end(json.toString());
130+
ctx.response().setStatusCode(BODY_JSON.getStatus());
131+
ctx.json(json);
131132
}));
132133
router
133134
.route(QUERY_ENCODED_BOTH.getRawPath())

‎dd-java-agent/instrumentation/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
6767
true
6868
}
6969

70+
@Override
71+
boolean testResponseBodyJson() {
72+
true
73+
}
74+
7075
@Override
7176
boolean testBodyUrlencoded() {
7277
true

‎dd-java-agent/instrumentation/vertx-web-5.0/src/test/java/server/VertxTestServer.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,8 @@ public void start(final Promise<Void> startPromise) {
118118
BODY_JSON,
119119
() -> {
120120
JsonObject json = ctx.body().asJsonObject();
121-
ctx.response().setStatusCode(BODY_JSON.getStatus()).end(json.toString());
121+
ctx.response().setStatusCode(BODY_JSON.getStatus());
122+
ctx.json(json);
122123
}));
123124
router
124125
.route(QUERY_ENCODED_BOTH.getRawPath())

0 commit comments

Comments
 (0)