Skip to content

Commit 45e6107

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Add endpoint discover for spring mvc
1 parent fe9f968 commit 45e6107

File tree

29 files changed

+941
-10
lines changed

29 files changed

+941
-10
lines changed

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/AppSecDispatcherServletInstrumentation.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package datadog.trace.instrumentation.springweb;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
import static net.bytebuddy.matcher.ElementMatchers.isProtected;
7+
import static net.bytebuddy.matcher.ElementMatchers.not;
8+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
9+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
10+
11+
import com.google.auto.service.AutoService;
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.agent.tooling.InstrumenterModule;
14+
import datadog.trace.api.Config;
15+
import datadog.trace.api.telemetry.EndpointCollector;
16+
import java.util.Map;
17+
import net.bytebuddy.asm.Advice;
18+
import net.bytebuddy.matcher.ElementMatcher;
19+
import org.springframework.context.ApplicationContext;
20+
import org.springframework.web.method.HandlerMethod;
21+
import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
22+
import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping;
23+
24+
@AutoService(InstrumenterModule.class)
25+
public class AppSecDispatcherServletInstrumentation extends InstrumenterModule.AppSec
26+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
27+
28+
public AppSecDispatcherServletInstrumentation() {
29+
super("spring-web");
30+
}
31+
32+
@Override
33+
public String instrumentedType() {
34+
return "org.springframework.web.servlet.DispatcherServlet";
35+
}
36+
37+
@Override
38+
public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
39+
return not(
40+
hasClassNamed(
41+
"org.springframework.web.servlet.mvc.condition.PathPatternsRequestCondition"));
42+
}
43+
44+
@Override
45+
public String[] helperClassNames() {
46+
return new String[] {packageName + ".RequestMappingInfoIterator"};
47+
}
48+
49+
@Override
50+
public void methodAdvice(MethodTransformer transformer) {
51+
transformer.applyAdvice(
52+
isMethod()
53+
.and(isProtected())
54+
.and(named("onRefresh"))
55+
.and(takesArgument(0, named("org.springframework.context.ApplicationContext")))
56+
.and(takesArguments(1)),
57+
AppSecDispatcherServletInstrumentation.class.getName() + "$AppSecHandlerMappingAdvice");
58+
}
59+
60+
public static class AppSecHandlerMappingAdvice {
61+
62+
@Advice.OnMethodExit(suppress = Throwable.class)
63+
public static void afterRefresh(@Advice.Argument(0) final ApplicationContext springCtx) {
64+
if (Config.get().isApiSecurityEndpointCollectionEnabled()) {
65+
final RequestMappingHandlerMapping handler =
66+
springCtx.getBean(RequestMappingHandlerMapping.class);
67+
if (handler == null) {
68+
return;
69+
}
70+
final Map<RequestMappingInfo, HandlerMethod> mappings = handler.getHandlerMethods();
71+
if (mappings == null || mappings.isEmpty()) {
72+
return;
73+
}
74+
EndpointCollector.get().supplier(new RequestMappingInfoIterator(mappings));
75+
}
76+
}
77+
}
78+
}

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/RequestMappingInfoIterator.java‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package datadog.trace.instrumentation.springweb;
2+
3+
import datadog.trace.api.telemetry.Endpoint;
4+
import datadog.trace.api.telemetry.Endpoint.Method;
5+
import java.util.ArrayList;
6+
import java.util.HashMap;
7+
import java.util.Iterator;
8+
import java.util.LinkedList;
9+
import java.util.List;
10+
import java.util.Map;
11+
import java.util.NoSuchElementException;
12+
import java.util.Queue;
13+
import java.util.Set;
14+
import org.springframework.web.method.HandlerMethod;
15+
import org.springframework.web.servlet.mvc.condition.MediaTypeExpression;
16+
import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
17+
18+
public class RequestMappingInfoIterator implements Iterator<Endpoint> {
19+
20+
private final Map<RequestMappingInfo, HandlerMethod> mappings;
21+
private final Queue<Endpoint> queue = new LinkedList<>();
22+
private Iterator<Map.Entry<RequestMappingInfo, HandlerMethod>> iterator;
23+
24+
public RequestMappingInfoIterator(final Map<RequestMappingInfo, HandlerMethod> mappings) {
25+
this.mappings = mappings;
26+
}
27+
28+
private Iterator<Map.Entry<RequestMappingInfo, HandlerMethod>> iterator() {
29+
if (iterator == null) {
30+
iterator = mappings.entrySet().iterator();
31+
}
32+
return iterator;
33+
}
34+
35+
@Override
36+
public boolean hasNext() {
37+
return !queue.isEmpty() || iterator().hasNext();
38+
}
39+
40+
@Override
41+
public Endpoint next() {
42+
if (queue.isEmpty()) {
43+
fetchNext();
44+
}
45+
final Endpoint endpoint = queue.poll();
46+
if (endpoint == null) {
47+
throw new NoSuchElementException();
48+
}
49+
return endpoint;
50+
}
51+
52+
private void fetchNext() {
53+
final Iterator<Map.Entry<RequestMappingInfo, HandlerMethod>> delegate = iterator();
54+
if (!delegate.hasNext()) {
55+
return;
56+
}
57+
final Map.Entry<RequestMappingInfo, HandlerMethod> nextEntry = delegate.next();
58+
final RequestMappingInfo nextInfo = nextEntry.getKey();
59+
final HandlerMethod nextHandler = nextEntry.getValue();
60+
final List<String> requestBody =
61+
parseMediaTypes(nextInfo.getConsumesCondition().getExpressions());
62+
final List<String> responseBody =
63+
parseMediaTypes(nextInfo.getProducesCondition().getExpressions());
64+
for (final String path : nextInfo.getPatternsCondition().getPatterns()) {
65+
final List<String> methods = Method.parseMethods(nextInfo.getMethodsCondition().getMethods());
66+
for (final String method : methods) {
67+
final Endpoint endpoint =
68+
new Endpoint()
69+
.type(Endpoint.Type.REST)
70+
.operation(Endpoint.Operation.HTTP_REQUEST)
71+
.path(path)
72+
.method(method)
73+
.requestBodyType(requestBody)
74+
.responseBodyType(responseBody);
75+
if (nextHandler != null) {
76+
final Map<String, String> metadata = new HashMap<>();
77+
metadata.put("handler", nextHandler.toString());
78+
endpoint.metadata(metadata);
79+
}
80+
queue.add(endpoint);
81+
}
82+
}
83+
}
84+
85+
private List<String> parseMediaTypes(final Set<MediaTypeExpression> expressions) {
86+
if (expressions == null || expressions.isEmpty()) {
87+
return null;
88+
}
89+
final List<String> result = new ArrayList<>(expressions.size());
90+
for (final MediaTypeExpression expression : expressions) {
91+
result.add(expression.toString());
92+
}
93+
return result;
94+
}
95+
}

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import datadog.trace.api.iast.IastContext
99
import datadog.trace.api.iast.InstrumentationBridge
1010
import datadog.trace.api.iast.SourceTypes
1111
import datadog.trace.api.iast.propagation.PropagationModule
12+
import datadog.trace.api.telemetry.Endpoint
1213
import datadog.trace.bootstrap.instrumentation.api.Tags
1314
import datadog.trace.core.DDSpan
1415
import datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator
@@ -19,6 +20,7 @@ import okhttp3.Response
1920
import org.springframework.boot.SpringApplication
2021
import org.springframework.boot.context.embedded.EmbeddedWebApplicationContext
2122
import org.springframework.context.ConfigurableApplicationContext
23+
import org.springframework.http.MediaType
2224
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter
2325
import org.springframework.web.servlet.view.RedirectView
2426
import test.SetupSpecHelper
@@ -129,6 +131,22 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
129131
true
130132
}
131133

134+
@Override
135+
boolean testEndpointDiscovery() {
136+
true
137+
}
138+
139+
@Override
140+
void assertEndpointDiscovery(final List<?> endpoints) {
141+
final discovered = endpoints.collectEntries { [(it.method): it] } as Map<String, Endpoint>
142+
assert discovered.keySet().containsAll([Endpoint.Method.POST, Endpoint.Method.PATCH, Endpoint.Method.PUT])
143+
discovered.values().each {
144+
assert it.requestBodyType.containsAll([MediaType.APPLICATION_JSON_VALUE])
145+
assert it.responseBodyType.containsAll([MediaType.TEXT_PLAIN_VALUE])
146+
assert it.metadata['handler'] == 'public org.springframework.http.ResponseEntity test.boot.TestController.discovery()'
147+
}
148+
}
149+
132150
@Override
133151
Serializable expectedServerSpanRoute(ServerEndpoint endpoint) {
134152
switch (endpoint) {

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/TestController.groovy‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import javax.servlet.http.HttpServletRequest
2222

2323
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_JSON
2424
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED
25+
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ENDPOINT_DISCOVERY
2526
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
2627
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
2728
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.FORWARDED
@@ -158,6 +159,16 @@ class TestController {
158159
}
159160
}
160161

162+
@RequestMapping(value = "/discovery",
163+
method = [RequestMethod.POST, RequestMethod.PATCH, RequestMethod.PUT],
164+
consumes = MediaType.APPLICATION_JSON_VALUE,
165+
produces = MediaType.TEXT_PLAIN_VALUE)
166+
ResponseEntity discovery() {
167+
HttpServerTest.controller(ENDPOINT_DISCOVERY) {
168+
new ResponseEntity(ENDPOINT_DISCOVERY.body, HttpStatus.valueOf(ENDPOINT_DISCOVERY.status))
169+
}
170+
}
171+
161172
@ExceptionHandler
162173
ResponseEntity handleException(Throwable throwable) {
163174
new ResponseEntity(throwable.message, HttpStatus.INTERNAL_SERVER_ERROR)

‎dd-java-agent/instrumentation/spring-webmvc-5.3/src/main/java/datadog/trace/instrumentation/springweb/AppSecDispatcherServletWithPathPatternsInstrumentation.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package datadog.trace.instrumentation.springweb;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
import static net.bytebuddy.matcher.ElementMatchers.isProtected;
7+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
8+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
9+
10+
import com.google.auto.service.AutoService;
11+
import datadog.trace.agent.tooling.Instrumenter;
12+
import datadog.trace.agent.tooling.InstrumenterModule;
13+
import datadog.trace.api.Config;
14+
import datadog.trace.api.telemetry.EndpointCollector;
15+
import java.util.Map;
16+
import net.bytebuddy.asm.Advice;
17+
import net.bytebuddy.matcher.ElementMatcher;
18+
import org.springframework.context.ApplicationContext;
19+
import org.springframework.web.method.HandlerMethod;
20+
import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
21+
import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping;
22+
23+
@AutoService(InstrumenterModule.class)
24+
public class AppSecDispatcherServletWithPathPatternsInstrumentation
25+
extends InstrumenterModule.AppSec
26+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
27+
28+
public AppSecDispatcherServletWithPathPatternsInstrumentation() {
29+
super("spring-web");
30+
}
31+
32+
@Override
33+
public String instrumentedType() {
34+
return "org.springframework.web.servlet.DispatcherServlet";
35+
}
36+
37+
@Override
38+
public String[] helperClassNames() {
39+
return new String[] {packageName + ".RequestMappingInfoWithPathPatternsIterator"};
40+
}
41+
42+
@Override
43+
public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
44+
return hasClassNamed(
45+
"org.springframework.web.servlet.mvc.condition.PathPatternsRequestCondition");
46+
}
47+
48+
@Override
49+
public void methodAdvice(MethodTransformer transformer) {
50+
transformer.applyAdvice(
51+
isMethod()
52+
.and(isProtected())
53+
.and(named("onRefresh"))
54+
.and(takesArgument(0, named("org.springframework.context.ApplicationContext")))
55+
.and(takesArguments(1)),
56+
AppSecDispatcherServletWithPathPatternsInstrumentation.class.getName()
57+
+ "$AppSecHandlerMappingAdvice");
58+
}
59+
60+
public static class AppSecHandlerMappingAdvice {
61+
62+
@Advice.OnMethodExit(suppress = Throwable.class)
63+
public static void afterRefresh(@Advice.Argument(0) final ApplicationContext springCtx) {
64+
if (Config.get().isApiSecurityEndpointCollectionEnabled()) {
65+
final RequestMappingHandlerMapping handler =
66+
springCtx.getBean(RequestMappingHandlerMapping.class);
67+
if (handler == null) {
68+
return;
69+
}
70+
final Map<RequestMappingInfo, HandlerMethod> mappings = handler.getHandlerMethods();
71+
if (mappings == null || mappings.isEmpty()) {
72+
return;
73+
}
74+
EndpointCollector.get().supplier(new RequestMappingInfoWithPathPatternsIterator(mappings));
75+
}
76+
}
77+
}
78+
}

0 commit comments

Comments
 (0)