Skip to content

Commit 3c78ad2

Browse files
jandro996jandro996
committed
WIP
1 parent 1ecc81e commit 3c78ad2

File tree

17 files changed

+418
-8
lines changed

17 files changed

+418
-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/src/main/java/datadog/trace/instrumentation/jaxrs2/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.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 hierarchyMarkerType() {
36+
return "javax.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/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/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/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
}

‎dd-smoke-tests/jersey-2/src/test/groovy/datadog/smoketest/Jersey2AppsecSmokeTest.groovy

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package datadog.smoketest
2+
3+
import datadog.smoketest.appsec.AbstractAppSecServerSmokeTest
4+
import datadog.trace.api.Platform
5+
import okhttp3.Request
6+
import okhttp3.Response
7+
8+
class Jersey2AppsecSmokeTest extends AbstractAppSecServerSmokeTest{
9+
10+
@Override
11+
ProcessBuilder createProcessBuilder() {
12+
String jarPath = System.getProperty('datadog.smoketest.jersey2.jar.path')
13+
14+
List<String> command = []
15+
command.add(javaPath())
16+
command.addAll(defaultJavaProperties)
17+
command.addAll(defaultAppSecProperties)
18+
command.add('-Ddd.integration.grizzly.enabled=true')
19+
if (Platform.isJavaVersionAtLeast(17)) {
20+
command.addAll((String[]) ['--add-opens', 'java.base/java.lang=ALL-UNNAMED'])
21+
}
22+
command.addAll(['-jar', jarPath, Integer.toString(httpPort)])
23+
ProcessBuilder processBuilder = new ProcessBuilder(command)
24+
processBuilder.directory(new File(buildDirectory))
25+
return processBuilder
26+
}
27+
28+
void 'API Security samples only one request per endpoint'() {
29+
given:
30+
def url = "http://localhost:${httpPort}/hello/api_security/sampling/200?test=value"
31+
def request = new Request.Builder()
32+
.url(url)
33+
.addHeader('X-My-Header', "value")
34+
.get()
35+
.build()
36+
37+
when:
38+
List<Response> responses = (1..3).collect {
39+
client.newCall(request).execute()
40+
}
41+
42+
then:
43+
responses.each {
44+
assert it.code() == 200
45+
}
46+
waitForTraceCount(3)
47+
def spans = rootSpans.toList().toSorted { it.span.duration }
48+
spans.size() == 3
49+
def sampledSpans = spans.findAll {
50+
it.meta.keySet().any {
51+
it.startsWith('_dd.appsec.s.req.')
52+
}
53+
}
54+
sampledSpans.size() == 1
55+
def span = sampledSpans[0]
56+
span.meta.containsKey('_dd.appsec.s.req.query')
57+
span.meta.containsKey('_dd.appsec.s.req.params')
58+
span.meta.containsKey('_dd.appsec.s.req.headers')
59+
}
60+
61+
62+
void 'test response schema extraction'() {
63+
given:
64+
def url = "http://localhost:${httpPort}/hello/api_security/response"
65+
def request = new Request.Builder()
66+
.url(url)
67+
.get()
68+
.build()
69+
70+
when:
71+
final response = client.newCall(request).execute()
72+
waitForTraceCount(1)
73+
74+
then:
75+
response.code() == 200
76+
def span = rootSpans.first()
77+
span.meta.containsKey('_dd.appsec.s.res.headers')
78+
span.meta.containsKey('_dd.appsec.s.res.body')
79+
}
80+
}

0 commit comments

Comments
 (0)