Remove callsites for IAST sources and use bytebuddy advices

dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/iast/TaintableEnumeration.java → dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/iast/TaintableEnumeration.java

@@ -1,4 +1,4 @@
-package datadog.trace.instrumentation.servlet.http.iast;
+package datadog.trace.agent.tooling.iast;
 
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.propagation.PropagationModule;

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

@@ -284,8 +284,6 @@
 2 org.springframework.expression.*
 2 org.springframework.format.*
 2 org.springframework.http.*
-# Need for IAST: calls ServletRequest methods instrumented at callsite
-0 org.springframework.http.server.ServletServerHttpRequest
 # Need for IAST: we instrument these classes
 0 org.springframework.http.HttpHeaders
 0 org.springframework.http.ReadOnlyHttpHeaders
@@ -321,14 +319,11 @@
 2 org.springframework.validation.*
 2 org.springframework.web.*
 0 org.springframework.web.context.request.async.*
-0 org.springframework.web.context.request.*
 0 org.springframework.web.context.support.AbstractRefreshableWebApplicationContext
 0 org.springframework.web.context.support.GenericWebApplicationContext
 0 org.springframework.web.context.support.XmlWebApplicationContext
 0 org.springframework.web.reactive.*
 0 org.springframework.web.servlet.*
-# Included for IAST
-0 org.springframework.web.util.WebUtils
 # Need for IAST so propagation of tainted objects is complete in spring 2.7.5
 0 org.springframework.util.StreamUtils$NonClosingInputStream
 # Included for IAST Spring mvc unvalidated redirect and xss vulnerability

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3RequestInstrumentation.java

@@ -0,0 +1,83 @@
+package datadog.trace.instrumentation.servlet3;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.iast.IastContext;
+import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.Source;
+import datadog.trace.api.iast.SourceTypes;
+import datadog.trace.api.iast.propagation.PropagationModule;
+import java.util.Map;
+import javax.servlet.http.HttpServletRequest;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@SuppressWarnings("unused")
+@AutoService(Instrumenter.class)
+public class Servlet3RequestInstrumentation extends Instrumenter.Iast
+    implements Instrumenter.ForTypeHierarchy {
+
+  private static final String CLASS_NAME = Servlet3RequestInstrumentation.class.getName();
+  private static final ElementMatcher.Junction<? super TypeDescription> WRAPPER_CLASS =
+      named("javax.servlet.http.HttpServletRequestWrapper");
+
+  public Servlet3RequestInstrumentation() {
+    super("servlet", "servlet-3");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "javax.servlet.http.HttpServletRequest";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()))
+        .and(not(WRAPPER_CLASS))
+        .and(not(extendsClass(WRAPPER_CLASS)));
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isMethod().and(named("getParameterMap")).and(takesArguments(0)),
+        CLASS_NAME + "$GetParameterMapAdvice");
+  }
+
+  public static class GetParameterMapAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Source(SourceTypes.REQUEST_PARAMETER_VALUE)
+    public static void onExit(@Advice.Return final Map<String, String[]> parameters) {
+      if (parameters == null || parameters.isEmpty()) {
+        return;
+      }
+      final PropagationModule module = InstrumentationBridge.PROPAGATION;
+      if (module == null) {
+        return;
+      }
+      final IastContext ctx = IastContext.Provider.get();
+      for (final Map.Entry<String, String[]> entry : parameters.entrySet()) {
+        final String name = entry.getKey();
+        module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name);
+        final String[] values = entry.getValue();
+        if (values != null) {
+          for (final String value : entry.getValue()) {
+            module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name);
+          }
+        }
+      }
+    }
+
+    public static void muzzleCheck(final HttpServletRequest request) {
+      request.getParameterMap();
+    }
+  }
+}

dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/Servlet3RequestInstrumentationTest.groovy

@@ -0,0 +1,67 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.iast.InstrumentationBridge
+import datadog.trace.api.iast.SourceTypes
+import datadog.trace.api.iast.propagation.PropagationModule
+
+import javax.servlet.http.HttpServletRequest
+import javax.servlet.http.HttpServletRequestWrapper
+
+class Servlet3RequestInstrumentationTest extends AgentTestRunner {
+
+  @Override
+  protected void configurePreAgent() {
+    injectSysConfig("dd.iast.enabled", "true")
+  }
+
+  void cleanup() {
+    InstrumentationBridge.clearIastModules()
+  }
+
+  void 'test getParameterMap'() {
+    setup:
+    final iastModule = Mock(PropagationModule)
+    InstrumentationBridge.registerIastModule(iastModule)
+    final parameters = [parameter: ['header1', 'header2'] as String[]]
+    final mock = Mock(HttpServletRequest)
+    final request = suite.call(mock)
+
+    when:
+    final result = request.getParameterMap()
+
+    then:
+    result == parameters
+    1 * mock.getParameterMap() >> parameters
+    parameters.each { key, values ->
+      1 * iastModule.taint(_, key, SourceTypes.REQUEST_PARAMETER_NAME, key)
+      values.each { value ->
+        1 * iastModule.taint(_, value, SourceTypes.REQUEST_PARAMETER_VALUE, key)
+      }
+    }
+    0 * _
+
+    where:
+    suite << testSuite()
+  }
+
+  private List<Closure<? extends HttpServletRequest>> testSuite() {
+    return [
+      { HttpServletRequest request -> new CustomRequest(request: request) },
+      { HttpServletRequest request -> new CustomRequestWrapper(new CustomRequest(request: request)) },
+      { HttpServletRequest request ->
+        new HttpServletRequestWrapper(new CustomRequest(request: request))
+      }
+    ]
+  }
+
+  private static class CustomRequest implements HttpServletRequest {
+    @Delegate
+    private HttpServletRequest request
+  }
+
+  private static class CustomRequestWrapper extends HttpServletRequestWrapper {
+
+    CustomRequestWrapper(final HttpServletRequest request) {
+      super(request)
+    }
+  }
+}