Add support for endpoint discovery in spring mvc

.circleci/config.continue.yml.j2

@@ -36,7 +36,7 @@ instrumentation_modules: &instrumentation_modules "dd-java-agent/instrumentation
 debugger_modules: &debugger_modules "dd-java-agent/agent-debugger|dd-java-agent/agent-bootstrap|dd-java-agent/agent-builder|internal-api|communication|dd-trace-core"
 profiling_modules: &profiling_modules "dd-java-agent/agent-profiling"
 
-default_system_tests_commit: &default_system_tests_commit 1ef00a34ad1f83ae999887e510ef1ea1c27b151b
+default_system_tests_commit: &default_system_tests_commit 9be8b2793d687c7d9b39f3265fef27b5ec91910c
 
 parameters:
   nightly:

dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle

@@ -34,6 +34,20 @@ muzzle {
     extraDependency "javax.servlet:javax.servlet-api:3.0.1"
     extraDependency "org.springframework:spring-webmvc:3.1.0.RELEASE"
   }
+
+  pass {
+    name = 'spring-mvc-pre-5.3'
+    group = 'org.springframework'
+    module = 'spring-webmvc'
+    versions = "[3.1.0.RELEASE,5.3)"
+    skipVersions += [
+      '1.2.1',
+      '1.2.2',
+      '1.2.3',
+      '1.2.4'] // broken releases... missing dependencies
+    skipVersions += '3.2.1.RELEASE' // missing a required class.  (bad release?)
+    extraDependency "javax.servlet:javax.servlet-api:3.0.1"
+  }
 }
 
 apply from: "$rootDir/gradle/java.gradle"

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

@@ -0,0 +1,86 @@
+package datadog.trace.instrumentation.springweb;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isProtected;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
+import datadog.trace.api.telemetry.EndpointCollector;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.context.ApplicationContext;
+import org.springframework.web.method.HandlerMethod;
+import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
+import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping;
+
+@AutoService(InstrumenterModule.class)
+public class AppSecDispatcherServletInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public AppSecDispatcherServletInstrumentation() {
+    super("spring-web");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.springframework.web.servlet.DispatcherServlet";
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return not(
+        hasClassNamed(
+            "org.springframework.web.servlet.mvc.condition.PathPatternsRequestCondition"));
+  }
+
+  @Override
+  public String muzzleDirective() {
+    return "spring-mvc-pre-5.3";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".RequestMappingInfoIterator"};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isProtected())
+            .and(named("onRefresh"))
+            .and(takesArgument(0, named("org.springframework.context.ApplicationContext")))
+            .and(takesArguments(1)),
+        AppSecDispatcherServletInstrumentation.class.getName() + "$AppSecHandlerMappingAdvice");
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return super.isEnabled() && Config.get().isApiSecurityEndpointCollectionEnabled();
+  }
+
+  public static class AppSecHandlerMappingAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void afterRefresh(@Advice.Argument(0) final ApplicationContext springCtx) {
+      final RequestMappingHandlerMapping handler =
+          springCtx.getBean(RequestMappingHandlerMapping.class);
+      if (handler == null) {
+        return;
+      }
+      final Map<RequestMappingInfo, HandlerMethod> mappings = handler.getHandlerMethods();
+      if (mappings == null || mappings.isEmpty()) {
+        return;
+      }
+      EndpointCollector.get().supplier(new RequestMappingInfoIterator(mappings));
+    }
+  }
+}

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

@@ -0,0 +1,100 @@
+package datadog.trace.instrumentation.springweb;
+
+import datadog.trace.api.telemetry.Endpoint;
+import datadog.trace.api.telemetry.Endpoint.Method;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Queue;
+import java.util.Set;
+import org.springframework.web.method.HandlerMethod;
+import org.springframework.web.servlet.mvc.condition.MediaTypeExpression;
+import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
+
+public class RequestMappingInfoIterator implements Iterator<Endpoint> {
+
+  private final Map<RequestMappingInfo, HandlerMethod> mappings;
+  private final Queue<Endpoint> queue = new LinkedList<>();
+  private Iterator<Map.Entry<RequestMappingInfo, HandlerMethod>> iterator;
+  private boolean first = true;
+
+  public RequestMappingInfoIterator(final Map<RequestMappingInfo, HandlerMethod> mappings) {
+    this.mappings = mappings;
+  }
+
+  private Iterator<Map.Entry<RequestMappingInfo, HandlerMethod>> iterator() {
+    if (iterator == null) {
+      iterator = mappings.entrySet().iterator();
+    }
+    return iterator;
+  }
+
+  @Override
+  public boolean hasNext() {
+    return !queue.isEmpty() || iterator().hasNext();
+  }
+
+  @Override
+  public Endpoint next() {
+    if (queue.isEmpty()) {
+      fetchNext();
+    }
+    final Endpoint endpoint = queue.poll();
+    if (endpoint == null) {
+      throw new NoSuchElementException();
+    }
+    return endpoint;
+  }
+
+  private void fetchNext() {
+    final Iterator<Map.Entry<RequestMappingInfo, HandlerMethod>> delegate = iterator();
+    if (!delegate.hasNext()) {
+      return;
+    }
+    final Map.Entry<RequestMappingInfo, HandlerMethod> nextEntry = delegate.next();
+    final RequestMappingInfo nextInfo = nextEntry.getKey();
+    final HandlerMethod nextHandler = nextEntry.getValue();
+    final List<String> requestBody =
+        parseMediaTypes(nextInfo.getConsumesCondition().getExpressions());
+    final List<String> responseBody =
+        parseMediaTypes(nextInfo.getProducesCondition().getExpressions());
+    for (final String path : nextInfo.getPatternsCondition().getPatterns()) {
+      final List<String> methods = Method.parseMethods(nextInfo.getMethodsCondition().getMethods());
+      for (final String method : methods) {
+        Endpoint endpoint =
+            new Endpoint()
+                .type(Endpoint.Type.REST)
+                .operation(Endpoint.Operation.HTTP_REQUEST)
+                .path(path)
+                .method(method)
+                .requestBodyType(requestBody)
+                .responseBodyType(responseBody);
+        if (nextHandler != null) {
+          final Map<String, String> metadata = new HashMap<>();
+          metadata.put("handler", nextHandler.toString());
+          endpoint.metadata(metadata);
+        }
+        if (first) {
+          endpoint.first(true);
+          first = false;
+        }
+        queue.add(endpoint);
+      }
+    }
+  }
+
+  private List<String> parseMediaTypes(final Set<MediaTypeExpression> expressions) {
+    if (expressions == null || expressions.isEmpty()) {
+      return null;
+    }
+    final List<String> result = new ArrayList<>(expressions.size());
+    for (final MediaTypeExpression expression : expressions) {
+      result.add(expression.toString());
+    }
+    return result;
+  }
+}

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

@@ -9,6 +9,7 @@ import datadog.trace.api.iast.IastContext
 import datadog.trace.api.iast.InstrumentationBridge
 import datadog.trace.api.iast.SourceTypes
 import datadog.trace.api.iast.propagation.PropagationModule
+import datadog.trace.api.telemetry.Endpoint
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.core.DDSpan
 import datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator
@@ -19,6 +20,7 @@ import okhttp3.Response
 import org.springframework.boot.SpringApplication
 import org.springframework.boot.context.embedded.EmbeddedWebApplicationContext
 import org.springframework.context.ConfigurableApplicationContext
+import org.springframework.http.MediaType
 import org.springframework.web.servlet.handler.HandlerInterceptorAdapter
 import org.springframework.web.servlet.view.RedirectView
 import test.SetupSpecHelper
@@ -129,6 +131,22 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
     true
   }
 
+  @Override
+  boolean testEndpointDiscovery() {
+    true
+  }
+
+  @Override
+  void assertEndpointDiscovery(final List<?> endpoints) {
+    final discovered = endpoints.collectEntries { [(it.method): it] }  as Map<String, Endpoint>
+    assert discovered.keySet().containsAll([Endpoint.Method.POST, Endpoint.Method.PATCH, Endpoint.Method.PUT])
+    discovered.values().each {
+      assert it.requestBodyType.containsAll([MediaType.APPLICATION_JSON_VALUE])
+      assert it.responseBodyType.containsAll([MediaType.TEXT_PLAIN_VALUE])
+      assert it.metadata['handler'] == 'public org.springframework.http.ResponseEntity test.boot.TestController.discovery()'
+    }
+  }
+
   @Override
   Serializable expectedServerSpanRoute(ServerEndpoint endpoint) {
     switch (endpoint) {

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

@@ -22,6 +22,7 @@ import javax.servlet.http.HttpServletRequest
 
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_JSON
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ENDPOINT_DISCOVERY
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.FORWARDED
@@ -158,6 +159,16 @@ class TestController {
     }
   }
 
+  @RequestMapping(value = "/discovery",
+  method = [RequestMethod.POST, RequestMethod.PATCH, RequestMethod.PUT],
+  consumes = MediaType.APPLICATION_JSON_VALUE,
+  produces = MediaType.TEXT_PLAIN_VALUE)
+  ResponseEntity discovery() {
+    HttpServerTest.controller(ENDPOINT_DISCOVERY) {
+      new ResponseEntity(ENDPOINT_DISCOVERY.body, HttpStatus.valueOf(ENDPOINT_DISCOVERY.status))
+    }
+  }
+
   @ExceptionHandler
   ResponseEntity handleException(Throwable throwable) {
     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

@@ -0,0 +1,81 @@
+package datadog.trace.instrumentation.springweb;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isProtected;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
+import datadog.trace.api.telemetry.EndpointCollector;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.context.ApplicationContext;
+import org.springframework.web.method.HandlerMethod;
+import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
+import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping;
+
+@AutoService(InstrumenterModule.class)
+public class AppSecDispatcherServletWithPathPatternsInstrumentation
+    extends InstrumenterModule.AppSec
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public AppSecDispatcherServletWithPathPatternsInstrumentation() {
+    super("spring-web");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.springframework.web.servlet.DispatcherServlet";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".RequestMappingInfoWithPathPatternsIterator"};
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return hasClassNamed(
+        "org.springframework.web.servlet.mvc.condition.PathPatternsRequestCondition");
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isProtected())
+            .and(named("onRefresh"))
+            .and(takesArgument(0, named("org.springframework.context.ApplicationContext")))
+            .and(takesArguments(1)),
+        AppSecDispatcherServletWithPathPatternsInstrumentation.class.getName()
+            + "$AppSecHandlerMappingAdvice");
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return super.isEnabled() && Config.get().isApiSecurityEndpointCollectionEnabled();
+  }
+
+  public static class AppSecHandlerMappingAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void afterRefresh(@Advice.Argument(0) final ApplicationContext springCtx) {
+      final RequestMappingHandlerMapping handler =
+          springCtx.getBean(RequestMappingHandlerMapping.class);
+      if (handler == null) {
+        return;
+      }
+      final Map<RequestMappingInfo, HandlerMethod> mappings = handler.getHandlerMethods();
+      if (mappings == null || mappings.isEmpty()) {
+        return;
+      }
+      EndpointCollector.get().supplier(new RequestMappingInfoWithPathPatternsIterator(mappings));
+    }
+  }
+}