Keep the request's resource when Undertow forwards internally

dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java

@@ -18,6 +18,7 @@
 import io.undertow.server.HttpServerExchange;
 import io.undertow.servlet.handlers.ServletPathMatch;
 import io.undertow.servlet.handlers.ServletRequestContext;
+import javax.servlet.DispatcherType;
 import javax.servlet.ServletRequest;
 import javax.servlet.http.MappingMatch;
 import net.bytebuddy.asm.Advice;
@@ -64,7 +65,8 @@ public static class DispatchAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(
         @Advice.Argument(0) final HttpServerExchange exchange,
-        @Advice.Argument(1) final ServletRequestContext servletRequestContext) {
+        @Advice.Argument(1) final ServletRequestContext servletRequestContext,
+        @Advice.Argument(3) final DispatcherType dispatcherType) {
       ContextContinuation continuation = exchange.getAttachment(DATADOG_UNDERTOW_CONTINUATION);
       if (continuation != null) {
         AgentSpan undertowSpan = spanFromContext(continuation.context());
@@ -73,21 +75,27 @@ public static void enter(
         undertowSpan.setSpanName(SERVLET_REQUEST);
 
         undertowSpan.setTag(SERVLET_CONTEXT, request.getServletContext().getContextPath());
-        String relativePath = exchange.getRelativePath();
+        // Forward, include, error and async dispatches happen inside the original request, so they
+        // must not rename it: the root span keeps describing what the client asked for.
+        if (dispatcherType == DispatcherType.REQUEST
+            || UndertowDecorator.UNDERTOW_DISPATCH_RENAME_RESOURCE) {
+          String relativePath = exchange.getRelativePath();
 
-        ServletPathMatch servletPathMatch = servletRequestContext.getServletPathMatch();
-        if (UndertowDecorator.UNDERTOW_LEGACY_TRACING
-            && servletPathMatch != null
-            && servletPathMatch.getMappingMatch() != MappingMatch.DEFAULT) {
-          // Set the route unless the mapping match is default, this way we prevent setting route
-          // for a non-existing resource. Otherwise, it'd set a non-existing resource name with
-          // higher priority than 404 resource, so it wouldn't be able to set resource 404 later in
-          // the onResponse instrumentation.
-          HTTP_RESOURCE_DECORATOR.withRoute(
-              undertowSpan, exchange.getRequestMethod().toString(), relativePath, false);
+          ServletPathMatch servletPathMatch = servletRequestContext.getServletPathMatch();
+          if (UndertowDecorator.UNDERTOW_LEGACY_TRACING
+              && servletPathMatch != null
+              && servletPathMatch.getMappingMatch() != MappingMatch.DEFAULT) {
+            // Set the route unless the mapping match is default, this way we prevent setting route
+            // for a non-existing resource. Otherwise, it'd set a non-existing resource name with
+            // higher priority than 404 resource, so it wouldn't be able to set resource 404 later
+            // in
+            // the onResponse instrumentation.
+            HTTP_RESOURCE_DECORATOR.withRoute(
+                undertowSpan, exchange.getRequestMethod().toString(), relativePath, false);
+          }
+          // The servlet.path tag is expected even for a non-existing resource.
+          undertowSpan.setTag(SERVLET_PATH, relativePath);
         }
-        // The servlet.path tag is expected even for a non-existing resource.
-        undertowSpan.setTag(SERVLET_PATH, relativePath);
       }
     }
   }

dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy

@@ -5,6 +5,8 @@ import datadog.trace.agent.test.naming.TestingGenericHttpNamingConventions
 import datadog.trace.bootstrap.blocking.BlockingActionHelper
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.core.DDSpan
+import datadog.trace.instrumentation.undertow.ForwardingServlet
+import datadog.trace.instrumentation.undertow.IncludingServlet
 import io.undertow.Handlers
 import io.undertow.Undertow
 import io.undertow.UndertowOptions
@@ -14,6 +16,7 @@ import io.undertow.servlet.api.ServletContainer
 import io.undertow.servlet.api.ServletInfo
 import io.undertow.websockets.jsr.WebSocketDeploymentInfo
 import okhttp3.MediaType
+import okhttp3.Request
 import okhttp3.RequestBody
 import spock.lang.IgnoreIf
 
@@ -45,6 +48,8 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.WEBSOC
 
 abstract class UndertowServletTest extends HttpServerTest<Undertow> {
   private static final CONTEXT = "ctx"
+  private static final String FORWARDING_PATH = "/dispatch/forward"
+  private static final String INCLUDING_PATH = "/dispatch/include"
 
   private final static String MULTIPART_CONTENT_TYPE =
   'multipart/form-data; charset=utf-8; boundary=------------------------943d3207457896a3'
@@ -70,6 +75,8 @@ abstract class UndertowServletTest extends HttpServerTest<Undertow> {
         .setDeploymentName("servletContext.war")
         .addServlet(new ServletInfo("SuccessServlet", SuccessServlet).addMapping(SUCCESS.getPath()))
         .addServlet(new ServletInfo("ForwardedServlet", ForwardedServlet).addMapping(FORWARDED.getPath()))
+        .addServlet(new ServletInfo("ForwardingServlet", ForwardingServlet).addMapping(FORWARDING_PATH))
+        .addServlet(new ServletInfo("IncludingServlet", IncludingServlet).addMapping(INCLUDING_PATH))
         .addServlet(new ServletInfo("QueryEncodedBothServlet", QueryEncodedBothServlet).addMapping(QUERY_ENCODED_BOTH.getPath()))
         .addServlet(new ServletInfo("QueryEncodedServlet", QueryEncodedServlet).addMapping(QUERY_ENCODED_QUERY.getPath()))
         .addServlet(new ServletInfo("QueryParamServlet", QueryServlet).addMapping(QUERY_PARAM.getPath()))
@@ -290,6 +297,10 @@ abstract class UndertowServletTest extends HttpServerTest<Undertow> {
     true
   }
 
+  def dispatchRenamesResource() {
+    false
+  }
+
   @Override
   Serializable expectedServerSpanRoute(ServerEndpoint endpoint) {
     if (!generateHttpRoute()) {
@@ -331,6 +342,36 @@ abstract class UndertowServletTest extends HttpServerTest<Undertow> {
     body = null
   }
 
+  def "internal #dispatch keeps the original request's resource and servlet path"() {
+    setup:
+    def request = new Request.Builder().url(address.resolve(path.substring(1)).toURL()).get().build()
+
+    when:
+    def response = client.newCall(request).execute()
+
+    then:
+    response.code() == SUCCESS.status
+    response.body().string() == SUCCESS.body
+
+    when:
+    TEST_WRITER.waitForTraces(1)
+    DDSpan serverSpan = TEST_WRITER.flatten().find {
+      it.parentId == 0
+    }
+
+    then:
+    // with the opt-out, a forward renames the request to its target; an include leaves the path alone
+    def namedPath = dispatchRenamesResource() && dispatch == "forward" ? SUCCESS.path : path
+    serverSpan.resourceName.toString() == (generateHttpRoute() ? "GET $namedPath" : "GET /$CONTEXT$path")
+    serverSpan.tags["servlet.path"] == namedPath
+    serverSpan.tags[Tags.HTTP_ROUTE] == (generateHttpRoute() ? namedPath : null)
+
+    where:
+    dispatch  | path
+    "forward" | FORWARDING_PATH
+    "include" | INCLUDING_PATH
+  }
+
   // Pins the block-telemetry-3 wiring added to FormDataParserInstrumentation (urlencoded body,
   // DoParseAdvice.after) and MultiPartUploadHandlerInstrumentation (multipart body,
   // ParseBlockingAdvice.after): both now check the boolean returned by
@@ -422,3 +463,16 @@ class UndertowServletNoHttpRouteForkedTest extends UndertowServletTest implement
 
 class UndertowServletV1ForkedTest extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV1 {
 }
+
+class UndertowServletDispatchRenameForkedTest extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 {
+  @Override
+  def dispatchRenamesResource() {
+    true
+  }
+
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig("trace.undertow.dispatch.rename-resource", "true")
+  }
+}

dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/java/datadog/trace/instrumentation/undertow/ForwardingServlet.java

@@ -0,0 +1,18 @@
+package datadog.trace.instrumentation.undertow;
+
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS;
+
+import java.io.IOException;
+import javax.servlet.ServletException;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+/** Forwards to the success servlet through an internal {@code RequestDispatcher}. */
+public class ForwardingServlet extends HttpServlet {
+  @Override
+  protected void doGet(final HttpServletRequest req, final HttpServletResponse resp)
+      throws ServletException, IOException {
+    req.getRequestDispatcher(SUCCESS.getPath()).forward(req, resp);
+  }
+}

dd-java-agent/instrumentation/undertow/undertow-2.0/src/test/java/datadog/trace/instrumentation/undertow/IncludingServlet.java

@@ -0,0 +1,18 @@
+package datadog.trace.instrumentation.undertow;
+
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS;
+
+import java.io.IOException;
+import javax.servlet.ServletException;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+/** Includes to the success servlet through an internal {@code RequestDispatcher}. */
+public class IncludingServlet extends HttpServlet {
+  @Override
+  protected void doGet(final HttpServletRequest req, final HttpServletResponse resp)
+      throws ServletException, IOException {
+    req.getRequestDispatcher(SUCCESS.getPath()).include(req, resp);
+  }
+}

dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java

@@ -16,6 +16,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import io.undertow.server.HttpServerExchange;
 import io.undertow.servlet.handlers.ServletRequestContext;
+import jakarta.servlet.DispatcherType;
 import jakarta.servlet.ServletRequest;
 import net.bytebuddy.asm.Advice;
 
@@ -56,7 +57,8 @@ public static class DispatchAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(
         @Advice.Argument(0) final HttpServerExchange exchange,
-        @Advice.Argument(1) final ServletRequestContext servletRequestContext) {
+        @Advice.Argument(1) final ServletRequestContext servletRequestContext,
+        @Advice.Argument(3) final DispatcherType dispatcherType) {
       ContextContinuation continuation = exchange.getAttachment(DATADOG_UNDERTOW_CONTINUATION);
       if (continuation != null) {
         AgentSpan undertowSpan = spanFromContext(continuation.context());
@@ -65,7 +67,12 @@ public static void enter(
         undertowSpan.setSpanName(SERVLET_REQUEST);
 
         undertowSpan.setTag(SERVLET_CONTEXT, request.getServletContext().getContextPath());
-        undertowSpan.setTag(SERVLET_PATH, exchange.getRelativePath());
+        // Forward, include, error and async dispatches happen inside the original request, so they
+        // must not overwrite the original servlet path.
+        if (dispatcherType == DispatcherType.REQUEST
+            || UndertowDecorator.UNDERTOW_DISPATCH_RENAME_RESOURCE) {
+          undertowSpan.setTag(SERVLET_PATH, exchange.getRelativePath());
+        }
       }
     }
   }

dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/groovy/UndertowServletTest.groovy

@@ -1,6 +1,9 @@
 import datadog.trace.agent.test.asserts.TraceAssert
 import datadog.trace.agent.test.base.HttpServerTest
 import datadog.trace.agent.test.base.WebsocketServer
+import datadog.trace.core.DDSpan
+import datadog.trace.instrumentation.undertow.ForwardingServlet
+import datadog.trace.instrumentation.undertow.IncludingServlet
 import io.undertow.Handlers
 import io.undertow.Undertow
 import io.undertow.UndertowOptions
@@ -10,6 +13,7 @@ import io.undertow.servlet.api.ServletContainer
 import io.undertow.servlet.api.ServletInfo
 import io.undertow.websockets.jsr.WebSocketDeploymentInfo
 import jakarta.servlet.MultipartConfigElement
+import okhttp3.Request
 
 import java.nio.ByteBuffer
 
@@ -30,6 +34,8 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_B
 
 class UndertowServletTest extends HttpServerTest<Undertow> {
   private static final CONTEXT = "ctx"
+  private static final String FORWARDING_PATH = "/dispatch/forward"
+  private static final String INCLUDING_PATH = "/dispatch/include"
 
   class UndertowServer implements WebsocketServer {
     def port = 0
@@ -46,6 +52,8 @@ class UndertowServletTest extends HttpServerTest<Undertow> {
         .setDeploymentName("servletContext.war")
         .addServlet(new ServletInfo("SuccessServlet", SuccessServlet).addMapping(SUCCESS.getPath()))
         .addServlet(new ServletInfo("ForwardedServlet", ForwardedServlet).addMapping(FORWARDED.getPath()))
+        .addServlet(new ServletInfo("ForwardingServlet", ForwardingServlet).addMapping(FORWARDING_PATH))
+        .addServlet(new ServletInfo("IncludingServlet", IncludingServlet).addMapping(INCLUDING_PATH))
         .addServlet(new ServletInfo("QueryEncodedBothServlet", QueryEncodedBothServlet).addMapping(QUERY_ENCODED_BOTH.getPath()))
         .addServlet(new ServletInfo("QueryEncodedServlet", QueryEncodedServlet).addMapping(QUERY_ENCODED_QUERY.getPath()))
         .addServlet(new ServletInfo("QueryParamServlet", QueryServlet).addMapping(QUERY_PARAM.getPath()))
@@ -246,4 +254,49 @@ class UndertowServletTest extends HttpServerTest<Undertow> {
       throw new UnsupportedOperationException("responseSpan not implemented for " + endpoint)
     }
   }
+
+  def dispatchRenamesResource() {
+    false
+  }
+
+  def "internal #dispatch keeps the original request's resource and servlet path"() {
+    setup:
+    def request = new Request.Builder().url(address.resolve(path.substring(1)).toURL()).get().build()
+
+    when:
+    def response = client.newCall(request).execute()
+
+    then:
+    response.code() == SUCCESS.status
+    response.body().string() == SUCCESS.body
+
+    when:
+    TEST_WRITER.waitForTraces(1)
+    DDSpan serverSpan = TEST_WRITER.flatten().find {
+      it.parentId == 0
+    }
+
+    then:
+    // with the opt-out, a forward overwrites servlet.path with its target; an include leaves it alone
+    serverSpan.resourceName.toString() == "GET /$CONTEXT$path"
+    serverSpan.tags["servlet.path"] == (dispatchRenamesResource() && dispatch == "forward" ? SUCCESS.path : path)
+
+    where:
+    dispatch  | path
+    "forward" | FORWARDING_PATH
+    "include" | INCLUDING_PATH
+  }
+}
+
+class UndertowServletDispatchRenameForkedTest extends UndertowServletTest {
+  @Override
+  def dispatchRenamesResource() {
+    true
+  }
+
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig("trace.undertow.dispatch.rename-resource", "true")
+  }
 }

dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/java/datadog/trace/instrumentation/undertow/ForwardingServlet.java

@@ -0,0 +1,18 @@
+package datadog.trace.instrumentation.undertow;
+
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS;
+
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServlet;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import java.io.IOException;
+
+/** Forwards to the success servlet through an internal {@code RequestDispatcher}. */
+public class ForwardingServlet extends HttpServlet {
+  @Override
+  protected void doGet(final HttpServletRequest req, final HttpServletResponse resp)
+      throws ServletException, IOException {
+    req.getRequestDispatcher(SUCCESS.getPath()).forward(req, resp);
+  }
+}

dd-java-agent/instrumentation/undertow/undertow-2.2/src/test/java/datadog/trace/instrumentation/undertow/IncludingServlet.java

@@ -0,0 +1,18 @@
+package datadog.trace.instrumentation.undertow;
+
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS;
+
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServlet;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import java.io.IOException;
+
+/** Includes to the success servlet through an internal {@code RequestDispatcher}. */
+public class IncludingServlet extends HttpServlet {
+  @Override
+  protected void doGet(final HttpServletRequest req, final HttpServletResponse resp)
+      throws ServletException, IOException {
+    req.getRequestDispatcher(SUCCESS.getPath()).include(req, resp);
+  }
+}

dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java

@@ -43,6 +43,8 @@ public class UndertowDecorator
       UTF8BytesString.create(DECORATE.operationName());
   public static final boolean UNDERTOW_LEGACY_TRACING =
       Config.get().isLegacyTracingEnabled(true, "undertow");
+  public static final boolean UNDERTOW_DISPATCH_RENAME_RESOURCE =
+      Config.get().isUndertowDispatchRenameResource();
 
   @Override
   protected String[] instrumentationNames() {

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -99,6 +99,8 @@ public final class ConfigDefaults {
 
   static final boolean DEFAULT_CLIENT_IP_ENABLED = false;
 
+  static final boolean DEFAULT_UNDERTOW_DISPATCH_RENAME_RESOURCE = false;
+
   static final int DEFAULT_CLOCK_SYNC_PERIOD = 30; // seconds
 
   static final boolean DEFAULT_TRACE_LAMBDA_SNAPSTART_CLOCK_RESYNC_ENABLED = false;

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -154,6 +154,13 @@ public final class TraceInstrumentationConfig {
   public static final String SERVLET_PRINCIPAL_ENABLED = "trace.servlet.principal.enabled";
   public static final String SERVLET_ASYNC_TIMEOUT_ERROR = "trace.servlet.async-timeout.error";
 
+  /**
+   * Restores the old Undertow behavior of renaming the request (resource, route and servlet.path)
+   * to the target of an internal forward, include, error or async dispatch.
+   */
+  public static final String UNDERTOW_DISPATCH_RENAME_RESOURCE =
+      "trace.undertow.dispatch.rename-resource";
+
   public static final String SERVLET_ROOT_CONTEXT_SERVICE_NAME =
       "trace.servlet.root-context.service.name";