Capture http route in more places for vertx

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/EndHandlerWrapper.java

@@ -1,6 +1,6 @@
 package datadog.trace.instrumentation.vertx_3_4.server;
 
-import static datadog.trace.instrumentation.vertx_3_4.server.RouteHandlerWrapper.HANDLER_SPAN_CONTEXT_KEY;
+import static datadog.trace.instrumentation.vertx_3_4.server.RouteUpdateHelper.HANDLER_SPAN_CONTEXT_KEY;
 import static datadog.trace.instrumentation.vertx_3_4.server.VertxDecorator.DECORATE;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteHandlerInstrumentation.java

@@ -21,6 +21,7 @@ public RouteHandlerInstrumentation() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".EndHandlerWrapper",
+      packageName + ".RouteUpdateHelper",
       packageName + ".RouteHandlerWrapper",
       packageName + ".VertxDecorator",
       packageName + ".VertxDecorator$VertxURIDataAdapter",

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteHandlerWrapper.java

@@ -4,23 +4,20 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
-import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
+import static datadog.trace.instrumentation.vertx_3_4.server.RouteUpdateHelper.HANDLER_SPAN_CONTEXT_KEY;
+import static datadog.trace.instrumentation.vertx_3_4.server.RouteUpdateHelper.PARENT_SPAN_CONTEXT_KEY;
+import static datadog.trace.instrumentation.vertx_3_4.server.RouteUpdateHelper.updateRoute;
 import static datadog.trace.instrumentation.vertx_3_4.server.VertxDecorator.DECORATE;
 import static datadog.trace.instrumentation.vertx_3_4.server.VertxDecorator.INSTRUMENTATION_NAME;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.Tags;
 import io.vertx.core.Handler;
 import io.vertx.ext.web.RoutingContext;
 import io.vertx.ext.web.impl.RouteImpl;
 import io.vertx.ext.web.impl.RouterImpl;
 
 public class RouteHandlerWrapper implements Handler<RoutingContext> {
-  static final String PARENT_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".parent";
-  static final String HANDLER_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".handler";
-  static final String ROUTE_CONTEXT_KEY = "dd." + Tags.HTTP_ROUTE;
-
   private final Handler<RoutingContext> actual;
   private final boolean spanStarter;
 
@@ -65,7 +62,8 @@ public void handle(final RoutingContext routingContext) {
 
   private void setRoute(RoutingContext routingContext) {
     final AgentSpan parentSpan = routingContext.get(PARENT_SPAN_CONTEXT_KEY);
-    if (parentSpan == null) {
+    final AgentSpan handlerSpan = routingContext.get(HANDLER_SPAN_CONTEXT_KEY);
+    if (parentSpan == null && handlerSpan == null) {
       return;
     }
 
@@ -81,22 +79,6 @@ private void setRoute(RoutingContext routingContext) {
       }
       path = mountPoint + path;
     }
-    if (method != null && path != null && shouldUpdateRoute(routingContext, parentSpan, path)) {
-      routingContext.put(ROUTE_CONTEXT_KEY, path);
-      HTTP_RESOURCE_DECORATOR.withRoute(parentSpan, method, path, true);
-    }
-  }
-
-  static boolean shouldUpdateRoute(
-      final RoutingContext routingContext, final AgentSpan span, final String path) {
-    if (span == null) {
-      return false;
-    }
-    final String currentRoute = routingContext.get(ROUTE_CONTEXT_KEY);
-    if (currentRoute != null && currentRoute.equals(path)) {
-      return false;
-    }
-    // do not override route with a "/" if it's already set (it's probably more meaningful)
-    return !path.equals("/") || span.getTag(Tags.HTTP_ROUTE) == null;
+    updateRoute(routingContext, method, path, parentSpan, handlerSpan);
   }
 }

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteImplInstrumentation.java

@@ -20,8 +20,6 @@
 public class RouteImplInstrumentation extends InstrumenterModule
     implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
 
-  private Advice.PostProcessor.Factory postProcessorFactory;
-
   public RouteImplInstrumentation() {
     super("vertx", "vertx-3.4");
   }
@@ -33,11 +31,9 @@ public Reference[] additionalMuzzleReferences() {
 
   @Override
   public boolean isApplicable(Set<TargetSystem> enabledSystems) {
-    if (enabledSystems.contains(TargetSystem.IAST)) {
-      postProcessorFactory = IastPostProcessorFactory.INSTANCE;
-      return true;
-    }
-    return enabledSystems.contains(TargetSystem.APPSEC);
+    return enabledSystems.contains(TargetSystem.TRACING)
+        || enabledSystems.contains(TargetSystem.APPSEC)
+        || enabledSystems.contains(TargetSystem.IAST);
   }
 
   @Override
@@ -50,7 +46,7 @@ public String[] knownMatchingTypes() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".PathParameterPublishingHelper",
+      packageName + ".RouteUpdateHelper", packageName + ".PathParameterPublishingHelper",
     };
   }
 
@@ -82,6 +78,6 @@ public void methodAdvice(MethodTransformer transformer) {
 
   @Override
   public Advice.PostProcessor.Factory postProcessor() {
-    return postProcessorFactory;
+    return IastPostProcessorFactory.INSTANCE;
   }
 }

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteMatchesAdvice.java

@@ -1,7 +1,12 @@
 package datadog.trace.instrumentation.vertx_3_4.server;
 
+import static datadog.trace.instrumentation.vertx_3_4.server.RouteUpdateHelper.HANDLER_SPAN_CONTEXT_KEY;
+import static datadog.trace.instrumentation.vertx_3_4.server.RouteUpdateHelper.PARENT_SPAN_CONTEXT_KEY;
+import static datadog.trace.instrumentation.vertx_3_4.server.RouteUpdateHelper.updateRoute;
+
 import datadog.trace.api.iast.Source;
 import datadog.trace.api.iast.SourceTypes;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import io.vertx.ext.web.RoutingContext;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -16,6 +21,23 @@ static void after(
     if (ret != 0 || t != null) {
       return;
     }
+    final AgentSpan parentSpan = ctx.get(PARENT_SPAN_CONTEXT_KEY);
+    final AgentSpan handlerSpan = ctx.get(HANDLER_SPAN_CONTEXT_KEY);
+    if (ctx.currentRoute() != null) {
+      final String method = ctx.request().rawMethod();
+      String path = ctx.currentRoute().getPath();
+      String mountPoint = ctx.mountPoint();
+      if (mountPoint != null && !mountPoint.isEmpty()) {
+        if (mountPoint.charAt(mountPoint.length() - 1) == '/'
+            && path != null
+            && !path.isEmpty()
+            && path.charAt(0) == '/') {
+          mountPoint = mountPoint.substring(0, mountPoint.length() - 1);
+        }
+        path = mountPoint + path;
+      }
+      updateRoute(ctx, method, path, parentSpan, handlerSpan);
+    }
     Map<String, String> params = ctx.pathParams();
     if (params.isEmpty()) {
       return;
@@ -35,6 +57,23 @@ static void after(
       if (!ret || t != null) {
         return;
       }
+      final AgentSpan parentSpan = ctx.get(PARENT_SPAN_CONTEXT_KEY);
+      final AgentSpan handlerSpan = ctx.get(HANDLER_SPAN_CONTEXT_KEY);
+      if (ctx.currentRoute() != null) {
+        final String method = ctx.request().rawMethod();
+        String path = ctx.currentRoute().getPath();
+        String mountPoint = ctx.mountPoint();
+        if (mountPoint != null && !mountPoint.isEmpty()) {
+          if (mountPoint.charAt(mountPoint.length() - 1) == '/'
+              && path != null
+              && !path.isEmpty()
+              && path.charAt(0) == '/') {
+            mountPoint = mountPoint.substring(0, mountPoint.length() - 1);
+          }
+          path = mountPoint + path;
+        }
+        updateRoute(ctx, method, path, parentSpan, handlerSpan);
+      }
       Map<String, String> params = ctx.pathParams();
       if (params.isEmpty()) {
         return;

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteUpdateHelper.java

@@ -0,0 +1,62 @@
+package datadog.trace.instrumentation.vertx_3_4.server;
+
+import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import io.vertx.ext.web.RoutingContext;
+
+public final class RouteUpdateHelper {
+  public static final String PARENT_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".parent";
+  public static final String HANDLER_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".handler";
+  public static final String ROUTE_CONTEXT_KEY = "dd." + Tags.HTTP_ROUTE;
+
+  private RouteUpdateHelper() {}
+
+  public static void updateRoute(
+      final RoutingContext routingContext,
+      final String method,
+      final String path,
+      final AgentSpan parentSpan,
+      final AgentSpan handlerSpan) {
+    if (method == null || path == null) {
+      return;
+    }
+    if (!shouldUpdateRoute(routingContext, parentSpan, handlerSpan, path)) {
+      return;
+    }
+
+    routingContext.put(ROUTE_CONTEXT_KEY, path);
+    if (parentSpan != null) {
+      HTTP_RESOURCE_DECORATOR.withRoute(parentSpan, method, path, true);
+    }
+    if (handlerSpan != null
+        && handlerSpan.getResourceNamePriority() < ResourceNamePriorities.HTTP_FRAMEWORK_ROUTE) {
+      HTTP_RESOURCE_DECORATOR.withRoute(handlerSpan, method, path, true);
+    }
+  }
+
+  public static boolean shouldUpdateRoute(
+      final RoutingContext routingContext,
+      final AgentSpan parentSpan,
+      final AgentSpan handlerSpan,
+      final String path) {
+    if (parentSpan == null && handlerSpan == null) {
+      return false;
+    }
+    final String currentRoute = routingContext.get(ROUTE_CONTEXT_KEY);
+    if (currentRoute != null && currentRoute.equals(path)) {
+      return false;
+    }
+    // do not override route with a "/" if it's already set (it's probably more meaningful)
+    if (!path.equals("/")) {
+      return true;
+    }
+    return !hasHttpRoute(parentSpan) && !hasHttpRoute(handlerSpan);
+  }
+
+  public static boolean hasHttpRoute(final AgentSpan span) {
+    return span != null && span.getTag(Tags.HTTP_ROUTE) != null;
+  }
+}

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-3.4/src/test/groovy/server/RouteHandlerWrapperTest.groovy

@@ -0,0 +1,47 @@
+package server
+
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan
+import datadog.trace.bootstrap.instrumentation.api.Tags
+import datadog.trace.instrumentation.vertx_3_4.server.RouteUpdateHelper
+import io.vertx.ext.web.RoutingContext
+
+class RouteHandlerWrapperTest extends InstrumentationSpecification {
+
+  void "updateRoute writes route to both spans"() {
+    given:
+    def context = Mock(RoutingContext)
+    def parentSpan = Mock(AgentSpan)
+    def handlerSpan = Mock(AgentSpan)
+
+    when:
+    RouteUpdateHelper.updateRoute(
+      context, "GET", "/items/:id", parentSpan, handlerSpan)
+
+    then:
+    1 * context.get("dd.${Tags.HTTP_ROUTE}") >> null
+    1 * context.put("dd.${Tags.HTTP_ROUTE}", "/items/:id")
+    1 * parentSpan.setTag(Tags.HTTP_ROUTE, "/items/:id")
+    1 * handlerSpan.getResourceNamePriority() >> Byte.MIN_VALUE
+    1 * handlerSpan.setTag(Tags.HTTP_ROUTE, "/items/:id")
+    0 * _
+  }
+
+  void "updateRoute does not replace root route when one exists"() {
+    given:
+    def context = Mock(RoutingContext)
+    def parentSpan = Mock(AgentSpan)
+    def handlerSpan = Mock(AgentSpan)
+
+    when:
+    RouteUpdateHelper.updateRoute(context, "GET", "/", parentSpan, handlerSpan)
+
+    then:
+    1 * context.get("dd.${Tags.HTTP_ROUTE}") >> null
+    1 * parentSpan.getTag(Tags.HTTP_ROUTE) >> "/existing"
+    1 * handlerSpan.getTag(Tags.HTTP_ROUTE) >> "/existing"
+    0 * context.put(_, _)
+    0 * parentSpan.setTag(_, _)
+    0 * handlerSpan.setTag(_, _)
+  }
+}

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/EndHandlerWrapper.java

@@ -1,6 +1,6 @@
 package datadog.trace.instrumentation.vertx_4_0.server;
 
-import static datadog.trace.instrumentation.vertx_4_0.server.RouteHandlerWrapper.HANDLER_SPAN_CONTEXT_KEY;
+import static datadog.trace.instrumentation.vertx_4_0.server.RouteUpdateHelper.HANDLER_SPAN_CONTEXT_KEY;
 import static datadog.trace.instrumentation.vertx_4_0.server.VertxDecorator.DECORATE;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerInstrumentation.java

@@ -27,6 +27,7 @@ public Reference[] additionalMuzzleReferences() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".EndHandlerWrapper",
+      packageName + ".RouteUpdateHelper",
       packageName + ".RouteHandlerWrapper",
       packageName + ".VertxDecorator",
     };

dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteHandlerWrapper.java

@@ -4,22 +4,19 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
-import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
+import static datadog.trace.instrumentation.vertx_4_0.server.RouteUpdateHelper.HANDLER_SPAN_CONTEXT_KEY;
+import static datadog.trace.instrumentation.vertx_4_0.server.RouteUpdateHelper.PARENT_SPAN_CONTEXT_KEY;
+import static datadog.trace.instrumentation.vertx_4_0.server.RouteUpdateHelper.updateRoute;
 import static datadog.trace.instrumentation.vertx_4_0.server.VertxDecorator.DECORATE;
 import static datadog.trace.instrumentation.vertx_4_0.server.VertxDecorator.INSTRUMENTATION_NAME;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.Tags;
 import io.vertx.core.Handler;
 import io.vertx.ext.web.RoutingContext;
 import io.vertx.ext.web.impl.RouteImpl;
 
 public class RouteHandlerWrapper implements Handler<RoutingContext> {
-  static final String PARENT_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".parent";
-  static final String HANDLER_SPAN_CONTEXT_KEY = AgentSpan.class.getName() + ".handler";
-  static final String ROUTE_CONTEXT_KEY = "dd." + Tags.HTTP_ROUTE;
-
   private final Handler<RoutingContext> actual;
   private final boolean spanStarter;
 
@@ -62,7 +59,8 @@ public void handle(final RoutingContext routingContext) {
 
   private void setRoute(RoutingContext routingContext) {
     final AgentSpan parentSpan = routingContext.get(PARENT_SPAN_CONTEXT_KEY);
-    if (parentSpan == null) {
+    final AgentSpan handlerSpan = routingContext.get(HANDLER_SPAN_CONTEXT_KEY);
+    if (parentSpan == null && handlerSpan == null) {
       return;
     }
 
@@ -82,22 +80,6 @@ private void setRoute(RoutingContext routingContext) {
               : mountPoint;
       path = noBackslashhMountPoint + path;
     }
-    if (method != null && path != null && shouldUpdateRoute(routingContext, parentSpan, path)) {
-      routingContext.put(ROUTE_CONTEXT_KEY, path);
-      HTTP_RESOURCE_DECORATOR.withRoute(parentSpan, method, path, true);
-    }
-  }
-
-  static boolean shouldUpdateRoute(
-      final RoutingContext routingContext, final AgentSpan span, final String path) {
-    if (span == null) {
-      return false;
-    }
-    final String currentRoute = routingContext.get(ROUTE_CONTEXT_KEY);
-    if (currentRoute != null && currentRoute.equals(path)) {
-      return false;
-    }
-    // do not override route with a "/" if it's already set (it's probably more meaningful)
-    return !path.equals("/") || span.getTag(Tags.HTTP_ROUTE) == null;
+    updateRoute(routingContext, method, path, parentSpan, handlerSpan);
   }
 }