Update tests

Author: manuel-alvarez-alvarez

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java

@@ -275,6 +275,9 @@ public void increaseRaspTimeouts() {
   public boolean sampleHttpClientRequest(final long id) {
     httpClientRequestCount.incrementAndGet();
     synchronized (sampledHttpClientRequests) {
+      if (sampledHttpClientRequests.contains(id)) {
+        return true;
+      }
       if (sampledHttpClientRequests.size()
           < Config.get().getApiSecurityMaxDownstreamRequestBodyAnalysis()) {
         sampledHttpClientRequests.add(id);

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy

@@ -15,16 +15,13 @@ import datadog.trace.api.gateway.Events
 import datadog.trace.api.gateway.Flow
 import datadog.trace.api.gateway.RequestContext
 import datadog.trace.api.gateway.RequestContextSlot
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer
 import datadog.trace.bootstrap.instrumentation.api.TagContext
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.bootstrap.instrumentation.api.URIUtils
 import datadog.trace.core.DDSpan
 import datadog.trace.core.datastreams.StatsGroup
 import datadog.trace.test.util.Flaky
-import groovy.json.JsonOutput
-import groovy.json.JsonSlurper
 import spock.lang.AutoCleanup
 import spock.lang.IgnoreIf
 import spock.lang.Requires
@@ -43,7 +40,6 @@ import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_CLIENT_TA
 import static datadog.trace.api.config.TracerConfig.HEADER_TAGS
 import static datadog.trace.api.config.TracerConfig.REQUEST_HEADER_TAGS
 import static datadog.trace.api.config.TracerConfig.RESPONSE_HEADER_TAGS
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.get
 
 abstract class HttpClientTest extends VersionedNamingTestBase {

dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecHttpEngineInstrumentation.java

@@ -60,7 +60,8 @@ public static void onSendRequest(
         return;
       }
 
-      // do not include bodies in the redirect request
+      // increment the number of downstream requests but do not include request/response body
+      AppSecInterceptor.sampleRequest(ctx, span.getSpanId());
       AppSecInterceptor.onResponse(span, false, priorResponse);
       AppSecInterceptor.onRequest(span, false, userRequest.urlString(), userRequest);
     }

dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java

@@ -162,7 +162,7 @@ private static <P extends HttpClientPayload> void publish(
     }
   }
 
-  static boolean sampleRequest(final RequestContext ctx, final long requestId) {
+  public static boolean sampleRequest(final RequestContext ctx, final long requestId) {
     //  Check if the current http request was sampled
     CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
     BiFunction<RequestContext, Long, Flow<Boolean>> samplingCb =

dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecHttpEngineInstrumentation.java

@@ -60,7 +60,8 @@ public static void onSendRequest(
         return;
       }
 
-      // do not include bodies in the redirect request
+      // increment the number of downstream requests but do not include request/response body
+      AppSecInterceptor.sampleRequest(ctx, span.getSpanId());
       AppSecInterceptor.onResponse(span, false, priorResponse);
       AppSecInterceptor.onRequest(span, false, userRequest.url().toString(), userRequest);
     }

dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java

@@ -162,7 +162,7 @@ private static <P extends HttpClientPayload> void publish(
     }
   }
 
-  static boolean sampleRequest(final RequestContext ctx, final long requestId) {
+  public static boolean sampleRequest(final RequestContext ctx, final long requestId) {
     //  Check if the current http request was sampled
     CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
     BiFunction<RequestContext, Long, Flow<Boolean>> samplingCb =

dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java

@@ -14,6 +14,8 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.file.Paths;
 import java.sql.Connection;
@@ -289,7 +291,11 @@ public ResponseEntity<Map<String, Object>> apiSecurityResponse(
   public ResponseEntity<String> apiSecurityHttpClientOkHttp2(final HttpServletRequest request)
       throws IOException {
     // create an internal http request to the echo endpoint to validate the http client library
-    final String url = getEchoUrl(request);
+    String url = getEchoUrl(request);
+    final String redirect = request.getParameter("redirect");
+    if (redirect != null) {
+      url += "?redirect=true";
+    }
     Request.Builder clientRequest = new Request.Builder().url(url);
     if (requiresBody(request.getMethod())) {
       final String contentType = request.getContentType();
@@ -324,8 +330,12 @@ public ResponseEntity<String> apiSecurityHttpClientOkHttp2(final HttpServletRequ
   public ResponseEntity<String> apiSecurityHttpClientOkHttp3(final HttpServletRequest request)
       throws IOException {
     // create an internal http request to the echo endpoint to validate the http client library
-    final String url = getEchoUrl(request);
-    okhttp3.Request.Builder clientRequest = new okhttp3.Request.Builder().url(url);
+    final okhttp3.HttpUrl.Builder url = okhttp3.HttpUrl.parse(getEchoUrl(request)).newBuilder();
+    final String redirect = request.getParameter("redirect");
+    if (redirect != null) {
+      url.addQueryParameter("redirect", "true");
+    }
+    okhttp3.Request.Builder clientRequest = new okhttp3.Request.Builder().url(url.build());
     if (requiresBody(request.getMethod())) {
       final String contentType = request.getContentType();
       final byte[] data = readFully(request.getInputStream());
@@ -356,7 +366,12 @@ public ResponseEntity<String> apiSecurityHttpClientOkHttp3(final HttpServletRequ
   @RequestMapping(
       value = "/echo",
       method = {POST, GET, PUT})
-  public ResponseEntity<String> echo(final HttpServletRequest request) throws IOException {
+  public ResponseEntity<String> echo(final HttpServletRequest request)
+      throws IOException, URISyntaxException {
+    final String redirect = request.getParameter("redirect");
+    if (redirect != null) {
+      return ResponseEntity.status(HttpStatus.FOUND).location(new URI("/echo")).build();
+    }
     final String statusHeader = request.getHeader("Status");
     final int statusCode = statusHeader == null ? 200 : Integer.parseInt(statusHeader);
     ResponseEntity.BodyBuilder response = ResponseEntity.status(statusCode);

dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy

@@ -5,6 +5,7 @@ import datadog.trace.agent.test.utils.OkHttpUtils
 import datadog.trace.test.util.ThreadUtils
 import groovy.json.JsonSlurper
 import okhttp3.FormBody
+import okhttp3.HttpUrl
 import okhttp3.MediaType
 import okhttp3.Request
 import okhttp3.RequestBody
@@ -1024,12 +1025,32 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest {
     variant << httpClientDownstreamAnalysisVariants()
   }
 
-  private RootSpan assertDownstreamTrace() {
-    waitForTraceCount(2) // original + echo
+  void 'API Security downstream response redirect'() {
+    when:
+    final url =HttpUrl.parse("http://localhost:${httpPort}/api_security/http_client/${variant}")
+      .newBuilder()
+      .addQueryParameter('redirect', 'true')
+      .build()
+    final request = new Request.Builder()
+      .url(url)
+      .get()
+      .build()
+    final response = client.newCall(request).execute()
+
+    then:
+    response.code() == 200
+    assertDownstreamTrace(2)
+
+    where:
+    variant << httpClientDownstreamAnalysisVariants()
+  }
+
+  private RootSpan assertDownstreamTrace(int downstreamCount = 1) {
+    waitForTraceCount(downstreamCount + 1) // original one plus all downstream requests
 
     final rootSpans = this.rootSpans.toList()
     final span = rootSpans.find { it.getSpan().resource.contains('/api_security/http_client') }
-    span.metrics['_dd.appsec.downstream_request'] == 1
+    assert span.metrics['_dd.appsec.downstream_request'] == downstreamCount
 
     return span
   }