Store the http.route tag value inside the iast request context in Play

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/HttpRouteHandler.java

@@ -0,0 +1,16 @@
+package com.datadog.iast;
+
+import datadog.trace.api.gateway.RequestContext;
+import datadog.trace.api.gateway.RequestContextSlot;
+import java.util.function.BiConsumer;
+
+public class HttpRouteHandler implements BiConsumer<RequestContext, String> {
+
+  @Override
+  public void accept(final RequestContext ctx, final String route) {
+    final IastRequestContext iastCtx = ctx.getData(RequestContextSlot.IAST);
+    if (iastCtx != null) {
+      iastCtx.setRoute(route);
+    }
+  }
+}

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastRequestContext.java

@@ -36,6 +36,7 @@ public class IastRequestContext implements IastContext, HasMetricCollector {
   @Nullable private volatile String xForwardedProto;
   @Nullable private volatile String contentType;
   @Nullable private volatile String authorization;
+  @Nullable private volatile String route;
 
   /**
    * Use {@link IastRequestContext#IastRequestContext(TaintedObjects)} instead as we require more
@@ -121,6 +122,15 @@ public void setAuthorization(final String authorization) {
     this.authorization = authorization;
   }
 
+  @Nullable
+  public String getRoute() {
+    return route;
+  }
+
+  public void setRoute(final String route) {
+    this.route = route;
+  }
+
   public OverheadContext getOverheadContext() {
     return overheadContext;
   }

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastSystem.java

@@ -122,6 +122,7 @@ public static void start(
     registerRequestStartedCallback(ss, addTelemetry, dependencies);
     registerRequestEndedCallback(ss, addTelemetry, dependencies);
     registerHeadersCallback(ss);
+    registerHttpRouteCallback(ss);
     registerGrpcServerRequestMessageCallback(ss);
     maybeApplySecurityControls(instrumentation);
     LOGGER.debug("IAST started");
@@ -246,6 +247,10 @@ private static void registerHeadersCallback(final SubscriptionService ss) {
     ss.registerCallback(event, handler);
   }
 
+  private static void registerHttpRouteCallback(final SubscriptionService ss) {
+    ss.registerCallback(Events.get().httpRoute(), new HttpRouteHandler());
+  }
+
   private static void registerGrpcServerRequestMessageCallback(final SubscriptionService ss) {
     ss.registerCallback(Events.get().grpcServerRequestMessage(), new GrpcRequestMessageHandler());
   }

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java

@@ -233,7 +233,7 @@ public boolean consumeQuota(
           Object methodTag = rootSpan.getTag(Tags.HTTP_METHOD);
           method = (methodTag == null) ? "" : methodTag.toString();
           Object routeTag = rootSpan.getTag(Tags.HTTP_ROUTE);
-          path = (routeTag == null) ? "" : routeTag.toString();
+          path = (routeTag == null) ? getHttpRouteFromRequestContext(span) : routeTag.toString();
         }
         if (!maybeSkipVulnerability(ctx, type, method, path)) {
           return operation.consumeQuota(ctx);
@@ -316,6 +316,19 @@ public OverheadContext getContext(@Nullable final AgentSpan span) {
       return globalContext;
     }
 
+    @Nullable
+    public String getHttpRouteFromRequestContext(@Nullable final AgentSpan span) {
+      String httpRoute = null;
+      final RequestContext requestContext = span != null ? span.getRequestContext() : null;
+      if (requestContext != null) {
+        IastRequestContext iastRequestContext = requestContext.getData(RequestContextSlot.IAST);
+        if (iastRequestContext != null) {
+          httpRoute = iastRequestContext.getRoute();
+        }
+      }
+      return httpRoute == null ? "" : httpRoute;
+    }
+
     static int computeSamplingParameter(final float pct) {
       if (pct >= 100) {
         return 100;

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/HttpRouteHandlerTest.groovy

@@ -0,0 +1,39 @@
+package com.datadog.iast
+
+import datadog.trace.api.gateway.RequestContext
+import datadog.trace.api.gateway.RequestContextSlot
+import datadog.trace.test.util.DDSpecification
+import groovy.transform.CompileDynamic
+
+@CompileDynamic
+class HttpRouteHandlerTest extends DDSpecification {
+  void 'route is set'() {
+    given:
+    final handler = new HttpRouteHandler()
+    final iastCtx = Mock(IastRequestContext)
+    final ctx = Mock(RequestContext)
+    ctx.getData(RequestContextSlot.IAST) >> iastCtx
+
+    when:
+    handler.accept(ctx, '/foo')
+
+    then:
+    1 * ctx.getData(RequestContextSlot.IAST) >> iastCtx
+    1 * iastCtx.setRoute('/foo')
+    0 * _
+  }
+
+  void 'does nothing when context missing'() {
+    given:
+    final handler = new HttpRouteHandler()
+    final ctx = Mock(RequestContext)
+    ctx.getData(RequestContextSlot.IAST) >> null
+
+    when:
+    handler.accept(ctx, '/foo')
+
+    then:
+    1 * ctx.getData(RequestContextSlot.IAST) >> null
+    0 * _
+  }
+}

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/IastSystemTest.groovy

@@ -60,6 +60,7 @@ class IastSystemTest extends DDSpecification {
     1 * ss.registerCallback(Events.get().requestStarted(), _)
     1 * ss.registerCallback(Events.get().requestEnded(), _)
     1 * ss.registerCallback(Events.get().requestHeader(), _)
+    1 * ss.registerCallback(Events.get().httpRoute(), _)
     1 * ss.registerCallback(Events.get().grpcServerRequestMessage(), _)
     0 * _
     TestLogCollector.drainCapturedLogs().any { it.message.contains('IAST is starting') }

dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayHttpServerDecorator.java

@@ -112,13 +112,21 @@ private void dispatchRoute(final AgentSpan span, final String route) {
       if (ctx == null) {
         return;
       }
+      // Send event to AppSec provider
       final CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
-      if (cbp == null) {
-        return;
+      if (cbp != null) {
+        final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
+        if (cb != null) {
+          cb.accept(ctx, route);
+        }
       }
-      final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
-      if (cb != null) {
-        cb.accept(ctx, route);
+      // Send event to IAST provider
+      final CallbackProvider cbpIast = tracer().getCallbackProvider(RequestContextSlot.IAST);
+      if (cbpIast != null) {
+        final BiConsumer<RequestContext, String> cb = cbpIast.getCallback(EVENTS.httpRoute());
+        if (cb != null) {
+          cb.accept(ctx, route);
+        }
       }
     } catch (final Throwable t) {
       LOG.debug("Failed to dispatch route", t);

dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java

@@ -46,6 +46,10 @@ public static ContextScope onEnter(@Advice.Argument(value = 0, readOnly = false)
 
     req = RequestHelper.withTag(req, "_dd_HasPlayRequestSpan", "true");
 
+    // Moved from OnMethodExit
+    // Call onRequest on return after tags are populated.
+    DECORATE.onRequest(span, req, req, (AgentSpanContext.Extracted) null);
+
     return scope;
   }
 
@@ -63,9 +67,6 @@ public static void stopTraceOnResponse(
 
     final AgentSpan playControllerSpan = spanFromContext(playControllerScope.context());
 
-    // Call onRequest on return after tags are populated.
-    DECORATE.onRequest(playControllerSpan, req, req, (AgentSpanContext.Extracted) null);
-
     if (throwable == null) {
       responseFuture.onComplete(
           new RequestCompleteCallback(playControllerSpan),

dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayHttpServerDecorator.java

@@ -112,13 +112,21 @@ private void dispatchRoute(final AgentSpan span, final String route) {
       if (ctx == null) {
         return;
       }
+      // Send event to AppSec provider
       final CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
-      if (cbp == null) {
-        return;
+      if (cbp != null) {
+        final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
+        if (cb != null) {
+          cb.accept(ctx, route);
+        }
       }
-      final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
-      if (cb != null) {
-        cb.accept(ctx, route);
+      // Send event to IAST provider
+      final CallbackProvider cbpIast = tracer().getCallbackProvider(RequestContextSlot.IAST);
+      if (cbpIast != null) {
+        final BiConsumer<RequestContext, String> cb = cbpIast.getCallback(EVENTS.httpRoute());
+        if (cb != null) {
+          cb.accept(ctx, route);
+        }
       }
     } catch (final Throwable t) {
       LOG.debug("Failed to dispatch route", t);

dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java

@@ -47,6 +47,10 @@ public static ContextScope onEnter(
 
     req = req.addAttr(HasPlayRequestSpan.KEY, HasPlayRequestSpan.INSTANCE);
 
+    // Moved from OnMethodExit
+    // Call onRequest on return after tags are populated.
+    DECORATE.onRequest(span, req, req, extractedContext);
+
     return scope;
   }
 
@@ -65,9 +69,6 @@ public static void stopTraceOnResponse(
 
     final AgentSpan playControllerSpan = spanFromContext(playControllerScope.context());
 
-    // Call onRequest on return after tags are populated.
-    DECORATE.onRequest(playControllerSpan, req, req, extractedContext);
-
     if (throwable == null) {
       responseFuture.onComplete(
           new RequestCompleteCallback(playControllerSpan),

dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayHttpServerDecorator.java

@@ -168,13 +168,21 @@ private void dispatchRoute(final AgentSpan span, final CharSequence route) {
       if (ctx == null) {
         return;
       }
-      final CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
-      if (cbp == null) {
-        return;
+      // Send event to IAST provider
+      final CallbackProvider cbpIast = tracer().getCallbackProvider(RequestContextSlot.IAST);
+      if (cbpIast != null) {
+        final BiConsumer<RequestContext, String> cb = cbpIast.getCallback(EVENTS.httpRoute());
+        if (cb != null) {
+          cb.accept(ctx, URIUtils.decode(route.toString()));
+        }
       }
-      final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
-      if (cb != null) {
-        cb.accept(ctx, URIUtils.decode(route.toString()));
+      // Send event to AppSec provider
+      final CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
+      if (cbp != null) {
+        final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
+        if (cb != null) {
+          cb.accept(ctx, URIUtils.decode(route.toString()));
+        }
       }
     } catch (final Throwable t) {
       LOG.debug("Failed to dispatch route", t);

dd-smoke-tests/iast-util/src/testFixtures/groovy/datadog/smoketest/AbstractIastServerSmokeTest.groovy

@@ -111,6 +111,16 @@ abstract class AbstractIastServerSmokeTest extends AbstractServerSmokeTest {
     }
   }
 
+  protected void hasMeta(final String name) {
+    try {
+      waitForSpan(pollingConditions()) { span ->
+        return span.meta.containsKey(name)
+      }
+    } catch (SpockTimeoutError toe) {
+      throw new AssertionError("No matching meta with name $name found")
+    }
+  }
+
   protected void hasVulnerability(@ClosureParams(value = SimpleType, options = ['datadog.smoketest.model.Vulnerability'])
     final Closure<Boolean> matcher) {
     final found = []

dd-smoke-tests/play-2.4/app/controllers/IastController.scala

@@ -0,0 +1,48 @@
+package controllers
+
+import play.api.mvc._
+
+import java.nio.charset.StandardCharsets
+import java.security.MessageDigest
+
+class IastController extends Controller {
+
+  def multipleVulns(id: String): Action[AnyContent] = Action {
+    try {
+      MessageDigest.getInstance("SHA1").digest("hash1".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("SHA-1").digest("hash2".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("MD2").digest("hash3".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("MD5").digest("hash4".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("RIPEMD128").digest("hash5".getBytes(StandardCharsets.UTF_8))
+      Ok("ok")
+    } catch {
+      case e: Exception => InternalServerError(e.getMessage)
+    }
+  }
+
+  def postMultipleVulns(id: String): Action[AnyContent] = Action {
+    try {
+      MessageDigest.getInstance("SHA1").digest("hash1".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("SHA-1").digest("hash2".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("MD2").digest("hash3".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("MD5").digest("hash4".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("RIPEMD128").digest("hash5".getBytes(StandardCharsets.UTF_8))
+      Ok("ok")
+    } catch {
+      case e: Exception => InternalServerError(e.getMessage)
+    }
+  }
+
+  def multipleVulns2(id: String): Action[AnyContent] = Action {
+    try {
+      MessageDigest.getInstance("SHA1").digest("hash1".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("SHA-1").digest("hash2".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("MD2").digest("hash3".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("MD5").digest("hash4".getBytes(StandardCharsets.UTF_8))
+      MessageDigest.getInstance("RIPEMD128").digest("hash5".getBytes(StandardCharsets.UTF_8))
+      Ok("ok")
+    } catch {
+      case e: Exception => InternalServerError(e.getMessage)
+    }
+  }
+}

dd-smoke-tests/play-2.4/build.gradle

@@ -64,6 +64,7 @@ dependencies {
 
   testImplementation project(':dd-smoke-tests')
   testImplementation project(':dd-smoke-tests:appsec')
+  testImplementation(testFixtures(project(":dd-smoke-tests:iast-util")))
 }
 
 configurations.testImplementation {

dd-smoke-tests/play-2.4/conf/routes

@@ -9,3 +9,8 @@ GET     /welcomes                           controllers.SController.doGet(id: Op
 
 # AppSec endpoints for testing
 GET     /api_security/sampling/:statusCode  controllers.AppSecController.apiSecuritySampling(statusCode: Int, test: String)
+
+# IAST Sampling endpoints
+GET  /iast/multiple_vulns/:id  controllers.IastController.multipleVulns(id: String)
+POST    /iast/multiple_vulns/:id    controllers.IastController.postMultipleVulns(id: String)
+GET     /iast/multiple_vulns-2/:id  controllers.IastController.multipleVulns2(id: String)