Skip to content

Commit 414f1ea

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Extract Vert.x json body response schemas
1 parent 0078896 commit 414f1ea

File tree

16 files changed

+314
-4
lines changed

16 files changed

+314
-4
lines changed

‎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: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ public class GatewayBridge {
9696
private volatile DataSubscriberInfo initialReqDataSubInfo;
9797
private volatile DataSubscriberInfo rawRequestBodySubInfo;
9898
private volatile DataSubscriberInfo requestBodySubInfo;
99+
private volatile DataSubscriberInfo responseBodySubInfo;
99100
private volatile DataSubscriberInfo pathParamsSubInfo;
100101
private volatile DataSubscriberInfo respDataSubInfo;
101102
private volatile DataSubscriberInfo grpcServerMethodSubInfo;
@@ -135,6 +136,7 @@ public void init() {
135136
subscriptionService.registerCallback(EVENTS.requestMethodUriRaw(), this::onRequestMethodUriRaw);
136137
subscriptionService.registerCallback(EVENTS.requestBodyStart(), this::onRequestBodyStart);
137138
subscriptionService.registerCallback(EVENTS.requestBodyDone(), this::onRequestBodyDone);
139+
subscriptionService.registerCallback(EVENTS.responseBody(), this::onResponseBody);
138140
subscriptionService.registerCallback(
139141
EVENTS.requestClientSocketAddress(), this::onRequestClientSocketAddress);
140142
subscriptionService.registerCallback(
@@ -175,6 +177,7 @@ public void reset() {
175177
initialReqDataSubInfo = null;
176178
rawRequestBodySubInfo = null;
177179
requestBodySubInfo = null;
180+
responseBodySubInfo = null;
178181
pathParamsSubInfo = null;
179182
respDataSubInfo = null;
180183
grpcServerMethodSubInfo = null;
@@ -636,6 +639,40 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
636639
}
637640
}
638641

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

‎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
@@ -449,6 +450,7 @@ class GatewayBridgeSpecification extends DDSpecification {
449450
1 * ig.registerCallback(EVENTS.requestBodyStart(), _) >> { requestBodyStartCB = it[1]; null }
450451
1 * ig.registerCallback(EVENTS.requestBodyDone(), _) >> { requestBodyDoneCB = it[1]; null }
451452
1 * ig.registerCallback(EVENTS.requestBodyProcessed(), _) >> { requestBodyProcessedCB = it[1]; null }
453+
1 * ig.registerCallback(EVENTS.responseBody(), _) >> { responseBodyCB = it[1]; null }
452454
1 * ig.registerCallback(EVENTS.responseStarted(), _) >> { responseStartedCB = it[1]; null }
453455
1 * ig.registerCallback(EVENTS.responseHeader(), _) >> { respHeaderCB = it[1]; null }
454456
1 * ig.registerCallback(EVENTS.responseHeaderDone(), _) >> { respHeadersDoneCB = it[1]; null }
@@ -1314,4 +1316,17 @@ class GatewayBridgeSpecification extends DDSpecification {
13141316
0 * traceSegment.setTagTop(_, _)
13151317
}
13161318
1319+
void 'test on response body callback'() {
1320+
when:
1321+
responseBodyCB.apply(ctx, [test: 'this is a test'])
1322+
1323+
then:
1324+
1 * eventDispatcher.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT) >> nonEmptyDsInfo
1325+
1 * eventDispatcher.publishDataEvent(_, _, _, _) >> {
1326+
final bundle = it[2] as DataBundle
1327+
final body = bundle.get(KnownAddresses.RESPONSE_BODY_OBJECT)
1328+
assert body['test'] == 'this is a test'
1329+
}
1330+
}
1331+
13171332
}

‎dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayHttpServerDecorator.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@
1111
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
1212
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1313
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
14-
import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
1514
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
15+
import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
1616
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
1717
import datadog.trace.bootstrap.instrumentation.api.URIUtils;
1818
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;

‎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: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
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 io.vertx.core.json.JsonObject;
15+
import java.util.function.BiFunction;
16+
import net.bytebuddy.asm.Advice;
17+
18+
@RequiresRequestContext(RequestContextSlot.APPSEC)
19+
class RoutingContextJsonResponseAdvice {
20+
21+
@Advice.OnMethodEnter(suppress = Throwable.class)
22+
static void before(
23+
@Advice.Argument(0) Object source, @ActiveRequestContext RequestContext reqCtx) {
24+
25+
if (source == null) {
26+
return;
27+
}
28+
29+
Object object = source;
30+
if (object instanceof JsonObject) {
31+
object = ((JsonObject) object).getMap();
32+
}
33+
34+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
35+
BiFunction<RequestContext, Object, Flow<Void>> callback =
36+
cbp.getCallback(EVENTS.responseBody());
37+
if (callback == null) {
38+
return;
39+
}
40+
41+
Flow<Void> flow = callback.apply(reqCtx, object);
42+
Flow.Action action = flow.getAction();
43+
if (action instanceof Flow.Action.RequestBlockingAction) {
44+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
45+
if (blockResponseFunction == null) {
46+
return;
47+
}
48+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
49+
blockResponseFunction.tryCommitBlockingResponse(
50+
reqCtx.getTraceSegment(),
51+
rba.getStatusCode(),
52+
rba.getBlockingContentType(),
53+
rba.getExtraHeaders());
54+
55+
throw new BlockingException("Blocked request (for RoutingContext/json)");
56+
}
57+
}
58+
}

‎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)