Skip to content

Commit db61f58

Browse files
jandro996jandro996
committed
Jersey support
1 parent d6d3d21 commit db61f58

File tree

20 files changed

+443
-8
lines changed

20 files changed

+443
-8
lines changed

‎dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,8 @@
190190
0 com.fasterxml.jackson.databind.util.TokenBuffer$Parser
191191
0 com.fasterxml.jackson.databind.ObjectMapper
192192
0 com.fasterxml.jackson.module.afterburner.util.MyClassLoader
193+
# Included for API Security response schema collection
194+
0 com.fasterxml.jackson.jaxrs.*
193195
2 com.github.mustachejava.*
194196
2 com.google.api.*
195197
0 com.google.api.client.http.HttpRequest

‎dd-java-agent/instrumentation/jakarta-rs-annotations-3/src/main/java/datadog/trace/instrumentation/jakarta3/MessageBodyWriterInstrumentation.java

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
package datadog.trace.instrumentation.jakarta3;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.api.gateway.Events.EVENTS;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
7+
8+
import com.google.auto.service.AutoService;
9+
import datadog.appsec.api.blocking.BlockingException;
10+
import datadog.trace.advice.ActiveRequestContext;
11+
import datadog.trace.advice.RequiresRequestContext;
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.agent.tooling.InstrumenterModule;
14+
import datadog.trace.api.gateway.BlockResponseFunction;
15+
import datadog.trace.api.gateway.CallbackProvider;
16+
import datadog.trace.api.gateway.Flow;
17+
import datadog.trace.api.gateway.RequestContext;
18+
import datadog.trace.api.gateway.RequestContextSlot;
19+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
20+
import jakarta.ws.rs.core.MediaType;
21+
import java.util.function.BiFunction;
22+
import net.bytebuddy.asm.Advice;
23+
import net.bytebuddy.description.type.TypeDescription;
24+
import net.bytebuddy.matcher.ElementMatcher;
25+
26+
@AutoService(InstrumenterModule.class)
27+
public class MessageBodyWriterInstrumentation extends InstrumenterModule.AppSec
28+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
29+
30+
public MessageBodyWriterInstrumentation() {
31+
super("jakarta-rs");
32+
}
33+
34+
@Override
35+
public String hierarchyMarkerType() {
36+
return "jakarta.ws.rs.ext.MessageBodyWriter";
37+
}
38+
39+
@Override
40+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
41+
return implementsInterface(named(hierarchyMarkerType()));
42+
}
43+
44+
@Override
45+
public void methodAdvice(MethodTransformer transformer) {
46+
transformer.applyAdvice(
47+
named("writeTo").and(takesArguments(7)), getClass().getName() + "$MessageBodyWriterAdvice");
48+
}
49+
50+
@RequiresRequestContext(RequestContextSlot.APPSEC)
51+
public static class MessageBodyWriterAdvice {
52+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
53+
static void after(
54+
@Advice.Argument(0) Object entity,
55+
@Advice.Argument(4) MediaType mediaType,
56+
@ActiveRequestContext RequestContext reqCtx,
57+
@Advice.Thrown Throwable t) {
58+
if (t != null) {
59+
return;
60+
}
61+
62+
// TODO check if this works or is better to use JSON MediaType
63+
if (!MediaType.APPLICATION_JSON_TYPE.isCompatible(mediaType)) {
64+
return;
65+
}
66+
67+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
68+
BiFunction<RequestContext, Object, Flow<Void>> callback =
69+
cbp.getCallback(EVENTS.responseBody());
70+
if (callback == null) {
71+
return;
72+
}
73+
74+
Flow<Void> flow = callback.apply(reqCtx, entity);
75+
Flow.Action action = flow.getAction();
76+
if (action instanceof Flow.Action.RequestBlockingAction) {
77+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
78+
if (blockResponseFunction == null) {
79+
return;
80+
}
81+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
82+
blockResponseFunction.tryCommitBlockingResponse(
83+
reqCtx.getTraceSegment(),
84+
rba.getStatusCode(),
85+
rba.getBlockingContentType(),
86+
rba.getExtraHeaders());
87+
88+
throw new BlockingException("Blocked request (for MessageBodyWriter)");
89+
}
90+
}
91+
}
92+
}

‎dd-java-agent/instrumentation/jax-rs-annotations-2/build.gradle

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,12 @@ muzzle {
1010
module = "javax.ws.rs-api"
1111
versions = "[,]"
1212
}
13+
pass {
14+
group = "javax.ws.rs"
15+
module = "javax.ws.rs-api"
16+
name = 'javax-message-body-writer'
17+
versions = "[,]"
18+
}
1319
}
1420

1521
apply from: "$rootDir/gradle/java.gradle"

‎dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/MessageBodyWriterInstrumentation.java

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
package datadog.trace.instrumentation.jaxrs2;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.api.gateway.Events.EVENTS;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
7+
8+
import com.google.auto.service.AutoService;
9+
import datadog.appsec.api.blocking.BlockingException;
10+
import datadog.trace.advice.ActiveRequestContext;
11+
import datadog.trace.advice.RequiresRequestContext;
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.agent.tooling.InstrumenterModule;
14+
import datadog.trace.api.gateway.BlockResponseFunction;
15+
import datadog.trace.api.gateway.CallbackProvider;
16+
import datadog.trace.api.gateway.Flow;
17+
import datadog.trace.api.gateway.RequestContext;
18+
import datadog.trace.api.gateway.RequestContextSlot;
19+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
20+
import java.util.function.BiFunction;
21+
import javax.ws.rs.core.MediaType;
22+
import net.bytebuddy.asm.Advice;
23+
import net.bytebuddy.description.type.TypeDescription;
24+
import net.bytebuddy.matcher.ElementMatcher;
25+
26+
@AutoService(InstrumenterModule.class)
27+
public class MessageBodyWriterInstrumentation extends InstrumenterModule.AppSec
28+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
29+
30+
public MessageBodyWriterInstrumentation() {
31+
super("jax-rs");
32+
}
33+
34+
@Override
35+
public String muzzleDirective() {
36+
return "javax-message-body-writer";
37+
}
38+
39+
@Override
40+
public String hierarchyMarkerType() {
41+
return "javax.ws.rs.ext.MessageBodyWriter";
42+
}
43+
44+
@Override
45+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
46+
return implementsInterface(named(hierarchyMarkerType()));
47+
}
48+
49+
@Override
50+
public void methodAdvice(MethodTransformer transformer) {
51+
transformer.applyAdvice(
52+
named("writeTo").and(takesArguments(7)), getClass().getName() + "$MessageBodyWriterAdvice");
53+
}
54+
55+
@RequiresRequestContext(RequestContextSlot.APPSEC)
56+
public static class MessageBodyWriterAdvice {
57+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
58+
static void after(
59+
@Advice.Argument(0) Object entity,
60+
@Advice.Argument(4) MediaType mediaType,
61+
@ActiveRequestContext RequestContext reqCtx,
62+
@Advice.Thrown Throwable t) {
63+
if (t != null) {
64+
return;
65+
}
66+
67+
// TODO check if this works or is better to use JSON MediaType
68+
if (!MediaType.APPLICATION_JSON_TYPE.isCompatible(mediaType)) {
69+
return;
70+
}
71+
72+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
73+
BiFunction<RequestContext, Object, Flow<Void>> callback =
74+
cbp.getCallback(EVENTS.responseBody());
75+
if (callback == null) {
76+
return;
77+
}
78+
79+
Flow<Void> flow = callback.apply(reqCtx, entity);
80+
Flow.Action action = flow.getAction();
81+
if (action instanceof Flow.Action.RequestBlockingAction) {
82+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
83+
if (blockResponseFunction == null) {
84+
return;
85+
}
86+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
87+
blockResponseFunction.tryCommitBlockingResponse(
88+
reqCtx.getTraceSegment(),
89+
rba.getStatusCode(),
90+
rba.getBlockingContentType(),
91+
rba.getExtraHeaders());
92+
93+
throw new BlockingException("Blocked request (for MessageBodyWriter)");
94+
}
95+
}
96+
}
97+
}

‎dd-java-agent/instrumentation/jersey/build.gradle

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ dependencies {
5656
jersey2JettyTestRuntimeOnly group: 'javax.xml.bind', name: 'jaxb-api', version: '2.2.3'
5757
jersey2JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jetty-9')
5858
jersey2JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jersey-2-appsec')
59+
jersey2JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jax-rs-annotations-2')
5960

6061
jersey3JettyTestImplementation project(':dd-java-agent:testing'), {
6162
exclude group: 'org.eclipse.jetty', module: 'jetty-server'
@@ -72,6 +73,7 @@ dependencies {
7273
jersey3JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jetty-11')
7374
jersey3JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jersey-2-appsec')
7475
jersey3JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jersey-3-appsec')
76+
jersey3JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jakarta-rs-annotations-3')
7577
}
7678

7779
configurations.getByName('jersey3JettyTestRuntimeClasspath').resolutionStrategy {

‎dd-java-agent/instrumentation/jersey/src/jersey2JettyTest/groovy/datadog/trace/instrumentation/jersey2/ClassToConvertBodyTo.groovy

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,12 @@
11
package datadog.trace.instrumentation.jersey2
22

3+
import groovy.json.JsonBuilder
4+
35
class ClassToConvertBodyTo {
46
String a
7+
8+
@Override
9+
String toString() {
10+
new JsonBuilder([a: a]).toString()
11+
}
512
}

‎dd-java-agent/instrumentation/jersey/src/jersey2JettyTest/groovy/datadog/trace/instrumentation/jersey2/Jersey2JettyTest.groovy

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,11 @@ import javax.ws.rs.ext.ExceptionMapper
99

1010
class Jersey2JettyTest extends HttpServerTest<JettyServer> {
1111

12+
@Override
13+
boolean testResponseBodyJson() {
14+
return true
15+
}
16+
1217
@Override
1318
HttpServer server() {
1419
new JettyServer()

‎dd-java-agent/instrumentation/jersey/src/jersey2JettyTest/groovy/datadog/trace/instrumentation/jersey2/ServiceResource.groovy

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import javax.ws.rs.HeaderParam
1010
import javax.ws.rs.POST
1111
import javax.ws.rs.Path
1212
import javax.ws.rs.PathParam
13+
import javax.ws.rs.Produces
1314
import javax.ws.rs.QueryParam
1415
import javax.ws.rs.core.MediaType
1516
import javax.ws.rs.core.Response
@@ -87,10 +88,14 @@ class ServiceResource {
8788

8889
@POST
8990
@Path("body-json")
91+
@Produces(MediaType.APPLICATION_JSON)
9092
Response bodyJson(ClassToConvertBodyTo obj) {
91-
controller(BODY_JSON) {
92-
Response.status(BODY_JSON.status).entity("""{"a":"${obj.a}"}""" as String).build()
93-
}
93+
return controller(BODY_JSON, () -> {
94+
Response response = Response.status(BODY_JSON.status)
95+
.entity(obj)
96+
.build()
97+
return response
98+
})
9499
}
95100

96101
@GET

‎dd-java-agent/instrumentation/jersey/src/jersey3JettyTest/groovy/datadog/trace/instrumentation/jersey3/ClassToConvertBodyTo.groovy

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,12 @@
11
package datadog.trace.instrumentation.jersey3
22

3+
import groovy.json.JsonBuilder
4+
35
class ClassToConvertBodyTo {
46
String a
7+
8+
@Override
9+
String toString() {
10+
new JsonBuilder([a: a]).toString()
11+
}
512
}

‎dd-java-agent/instrumentation/jersey/src/jersey3JettyTest/groovy/datadog/trace/instrumentation/jersey3/Jersey3JettyTest.groovy

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,11 @@ import jakarta.ws.rs.ext.ExceptionMapper
88

99
class Jersey3JettyTest extends HttpServerTest<JettyServer> {
1010

11+
@Override
12+
boolean testResponseBodyJson() {
13+
return true
14+
}
15+
1116
@Override
1217
HttpServer server() {
1318
new JettyServer()

‎dd-java-agent/instrumentation/jersey/src/jersey3JettyTest/groovy/datadog/trace/instrumentation/jersey3/ServiceResource.groovy

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.instrumentation.jersey3
22

33
import datadog.appsec.api.blocking.Blocking
4+
import jakarta.ws.rs.Produces
45
import org.glassfish.jersey.media.multipart.FormDataParam
56

67
import jakarta.ws.rs.Consumes
@@ -87,10 +88,13 @@ class ServiceResource {
8788

8889
@POST
8990
@Path("body-json")
91+
@Produces(MediaType.APPLICATION_JSON)
9092
Response bodyJson(ClassToConvertBodyTo obj) {
91-
controller(BODY_JSON) {
92-
Response.status(BODY_JSON.status).entity("""{"a":"${obj.a}"}""" as String).build()
93-
}
93+
controller(BODY_JSON, () ->
94+
Response.status(BODY_JSON.status)
95+
.entity(obj)
96+
.build()
97+
)
9498
}
9599

96100
@GET

‎dd-smoke-tests/jersey-2/build.gradle

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ dependencies {
2020
implementation group: 'javax.xml', name: 'jaxb-api', version:'2.1'
2121
testImplementation project(':dd-smoke-tests')
2222
testImplementation(testFixtures(project(":dd-smoke-tests:iast-util")))
23+
testImplementation project(':dd-smoke-tests:appsec')
2324
}
2425

2526
tasks.withType(Test).configureEach {

‎dd-smoke-tests/jersey-2/src/main/java/com/restserver/Resource.java

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,4 +139,18 @@ public Response responseLocation(@QueryParam("param") String param) throws URISy
139139
public Response getCookie() throws SQLException {
140140
return Response.ok().cookie(new NewCookie("user-id", "7")).build();
141141
}
142+
143+
@Path("/api_security/response")
144+
@GET
145+
@Produces(MediaType.APPLICATION_JSON)
146+
public Response bodyJson() {
147+
TestEntity testEntity = new TestEntity("testing", "test");
148+
return Response.ok().entity(testEntity).build();
149+
}
150+
151+
@GET
152+
@Path("/api_security/sampling/{i}")
153+
public Response apiSecuritySamplingWithStatus(@PathParam("i") int i) {
154+
return Response.status(i).header("content-type", "text/plain").entity("Hello!\n").build();
155+
}
142156
}

0 commit comments

Comments
 (0)