Merge pull request #5393 from DataDog/glopes/tomcat-block-resp

Support block on response on tomcat

Author: cataphract

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

@@ -310,6 +310,21 @@ public AgentSpan onResponseStatus(final AgentSpan span, final int status) {
     return span;
   }
 
+  /**
+   * Whether AppSec should NOT be called during onResponse() for analysis of the status code and
+   * headers.
+   *
+   * <p>{@link #onResponse(AgentSpan, Object)} is usually called too late for AppSec to be able to
+   * alter the response, so for those modules where we support blocking on response this is <code>
+   * true</code> and AppSec has its own (earlier) hook point for processing the response (just
+   * before commit).
+   *
+   * @return whether AppSec analysis of the response is run separately from onResponse
+   */
+  protected boolean isAppSecOnResponseSeparate() {
+    return false;
+  }
+
   public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
     if (response != null) {
       final int status = status(response);
@@ -324,7 +339,9 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
         }
       }
 
-      callIGCallbackResponseAndHeaders(span, response, status);
+      if (!isAppSecOnResponseSeparate()) {
+        callIGCallbackResponseAndHeaders(span, response, status);
+      }
     }
     return span;
   }
@@ -401,30 +418,40 @@ private Flow<Void> callIGCallbackRequestHeaders(AgentSpan span, REQUEST_CARRIER
     return Flow.ResultFlow.empty();
   }
 
-  private void callIGCallbackResponseAndHeaders(AgentSpan span, RESPONSE carrier, int status) {
+  private Flow<Void> callIGCallbackResponseAndHeaders(
+      AgentSpan span, RESPONSE carrier, int status) {
+    return callIGCallbackResponseAndHeaders(span, carrier, status, responseGetter());
+  }
+
+  public <RESP> Flow<Void> callIGCallbackResponseAndHeaders(
+      AgentSpan span,
+      RESP carrier,
+      int status,
+      AgentPropagation.ContextVisitor<RESP> contextVisitor) {
     CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
     RequestContext requestContext = span.getRequestContext();
     if (cbp == null || requestContext == null) {
-      return;
+      return Flow.ResultFlow.empty();
     }
+
     BiFunction<RequestContext, Integer, Flow<Void>> addrCallback =
         cbp.getCallback(EVENTS.responseStarted());
     if (null != addrCallback) {
       addrCallback.apply(requestContext, status);
     }
-    AgentPropagation.ContextVisitor<RESPONSE> getter = responseGetter();
-    if (getter == null) {
-      return;
+    if (contextVisitor == null) {
+      return Flow.ResultFlow.empty();
     }
     IGKeyClassifier igKeyClassifier =
         IGKeyClassifier.create(
             requestContext,
             cbp.getCallback(EVENTS.responseHeader()),
             cbp.getCallback(EVENTS.responseHeaderDone()));
     if (null != igKeyClassifier) {
-      getter.forEachKey(carrier, igKeyClassifier);
-      igKeyClassifier.done();
+      contextVisitor.forEachKey(carrier, igKeyClassifier);
+      return igKeyClassifier.done();
     }
+    return Flow.ResultFlow.empty();
   }
 
   private Flow<Void> callIGCallbackURI(
@@ -498,9 +525,9 @@ private Flow<Void> callIGCallbackAddressAndPort(
   }
 
   /** This passes the headers through to the InstrumentationGateway */
-  private static final class IGKeyClassifier implements AgentPropagation.KeyClassifier {
+  protected static final class IGKeyClassifier implements AgentPropagation.KeyClassifier {
 
-    private static IGKeyClassifier create(
+    public static IGKeyClassifier create(
         RequestContext requestContext,
         TriConsumer<RequestContext, String, String> headerCallback,
         Function<RequestContext, Flow<Void>> doneCallback) {

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

@@ -36,6 +36,11 @@ abstract class TomcatServlet3Test extends AbstractServlet3Test<Tomcat, Context>
   @Shared
   TestAccessLogValve accessLogValue
 
+  @Override
+  boolean testBlockingOnResponse() {
+    true
+  }
+
   class TomcatServer implements HttpServer {
     def port = 0
     final Tomcat server

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

@@ -58,7 +58,8 @@ dependencies {
   testImplementation project(':dd-java-agent:instrumentation:servlet-common')
   testImplementation project(':dd-java-agent:instrumentation:servlet:request-3')
 
-  // include also tomcat appsec instrumentation for intercepting urlencoded/multipart bodies
+  // include also tomcat instrumentation
+  testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-5.5')
   testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-appsec-6')
   testRuntimeOnly project(':dd-java-agent:instrumentation:tomcat-appsec-7')
 

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

@@ -4,12 +4,12 @@ import datadog.trace.agent.test.asserts.TraceAssert
 import datadog.trace.agent.test.base.HttpServer
 import datadog.trace.agent.test.base.HttpServerTest
 import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.DDTags
 import datadog.trace.api.iast.InstrumentationBridge
 import datadog.trace.api.iast.SourceTypes
 import datadog.trace.api.iast.propagation.PropagationModule
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.core.DDSpan
-import datadog.trace.instrumentation.servlet3.Servlet3Decorator
 import datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator
 import okhttp3.FormBody
 import okhttp3.Request
@@ -65,7 +65,7 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
 
   @Override
   String component() {
-    return Servlet3Decorator.DECORATE.component()
+    'tomcat-server'
   }
 
   String getServletContext() {
@@ -155,9 +155,12 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
   int spanCount(ServerEndpoint endpoint) {
     if (endpoint == REDIRECT) {
       // Spring is generates a RenderView and ResponseSpan for REDIRECT
-      return super.spanCount(endpoint) + 1
+      super.spanCount(endpoint) + 1
+    } else if (endpoint == NOT_FOUND) {
+      super.spanCount(endpoint) + 2
+    } else {
+      super.spanCount(endpoint)
     }
-    return super.spanCount(endpoint)
   }
 
   @Override
@@ -399,4 +402,37 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
       }
     }
   }
+
+
+  protected void trailingSpans(TraceAssert traceAssert, ServerEndpoint serverEndpoint) {
+    if (serverEndpoint == NOT_FOUND) {
+      traceAssert.with {
+        span {
+          spanType 'web'
+          serviceName expectedServiceName()
+          operationName 'servlet.forward'
+          resourceName 'GET /error'
+          tags {
+            "$Tags.COMPONENT" 'java-web-servlet-dispatcher'
+            "$Tags.HTTP_ROUTE" '/error'
+            'servlet.context' "/$servletContext"
+            'servlet.path' '/not-found'
+            "$DDTags.PATHWAY_HASH" String
+            defaultTags()
+          }
+        }
+        span {
+          spanType 'web'
+          serviceName expectedServiceName()
+          operationName 'spring.handler'
+          resourceName 'BasicErrorController.error'
+          tags {
+            "$Tags.COMPONENT" 'spring-web-controller'
+            "$Tags.SPAN_KIND" 'server'
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
 }

dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/dynamic/DynamicRoutingTest.groovy

@@ -1,17 +1,17 @@
 package test.dynamic
 
-
 import datadog.trace.agent.test.asserts.TraceAssert
 import datadog.trace.agent.test.base.HttpServer
 import datadog.trace.agent.test.base.HttpServerTest
 import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.DDTags
 import datadog.trace.bootstrap.instrumentation.api.Tags
-import datadog.trace.instrumentation.servlet3.Servlet3Decorator
 import datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator
 import org.springframework.boot.SpringApplication
 import org.springframework.boot.context.embedded.EmbeddedWebApplicationContext
 import org.springframework.context.ConfigurableApplicationContext
 import org.springframework.web.servlet.view.RedirectView
+import org.springframework.web.util.NestedServletException
 import test.boot.SecurityConfig
 
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
@@ -61,9 +61,14 @@ class DynamicRoutingTest extends HttpServerTest<ConfigurableApplicationContext>
     return new SpringBootServer()
   }
 
+  @Override
+  String expectedServiceName() {
+    'root-servlet'
+  }
+
   @Override
   String component() {
-    return Servlet3Decorator.DECORATE.component()
+    'tomcat-server'
   }
 
   @Override
@@ -110,15 +115,18 @@ class DynamicRoutingTest extends HttpServerTest<ConfigurableApplicationContext>
 
   @Override
   Map<String, Serializable> expectedExtraServerTags(ServerEndpoint endpoint) {
-    ["servlet.path": endpoint.path]
+    ["servlet.path": endpoint.path, 'servlet.context': '/']
   }
 
   int spanCount(ServerEndpoint endpoint) {
     if (endpoint == REDIRECT) {
       // Spring is generates a RenderView and ResponseSpan for REDIRECT
-      return super.spanCount(endpoint) + 1
+      super.spanCount(endpoint) + 1
+    }  else if (endpoint == EXCEPTION) {
+      super.spanCount(endpoint) +2
+    } else {
+      super.spanCount(endpoint)
     }
-    return super.spanCount(endpoint)
   }
 
   @Override
@@ -187,4 +195,47 @@ class DynamicRoutingTest extends HttpServerTest<ConfigurableApplicationContext>
     int firstUnderscore = x.indexOf('_')
     return firstUnderscore == -1 ? x : x.substring(0, firstUnderscore)
   }
+
+  Map<String, Serializable> expectedExtraErrorInformation(ServerEndpoint endpoint) {
+    if (endpoint == EXCEPTION) {
+      ["error.message"  : 'Request processing failed; nested exception is java.lang.Exception: controller exception',
+        "error.type" : NestedServletException.name,
+        "error.stack": String]
+    } else {
+      super.expectedExtraErrorInformation(endpoint)
+    }
+  }
+
+  protected void trailingSpans(TraceAssert traceAssert, ServerEndpoint serverEndpoint) {
+    if (serverEndpoint == EXCEPTION) {
+      traceAssert.with {
+        span {
+          spanType 'web'
+          serviceName expectedServiceName()
+          operationName 'servlet.forward'
+          resourceName 'GET /error'
+          tags {
+            "$Tags.COMPONENT" 'java-web-servlet-dispatcher'
+            "$Tags.HTTP_ROUTE" '/error'
+            'servlet.context' '/'
+            'servlet.path' serverEndpoint.path
+            "$DDTags.PATHWAY_HASH" String
+            defaultTags()
+          }
+        }
+        span {
+          spanType 'web'
+          childOfPrevious()
+          serviceName expectedServiceName()
+          operationName 'spring.handler'
+          resourceName 'BasicErrorController.error'
+          tags {
+            "$Tags.COMPONENT" 'spring-web-controller'
+            "$Tags.SPAN_KIND" 'server'
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
 }