Analyze okhttp client redirections as separated request

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

@@ -22,8 +22,6 @@ 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
@@ -42,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 {
@@ -69,7 +66,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
       }
       prefix("redirect") {
         handleDistributedRequest()
-        redirect(server.address.resolve("/success").toURL().toString())
+        redirect(server.address.resolve(request.getHeader('Location') ?: "/success").toURL().toString())
       }
       prefix("another-redirect") {
         handleDistributedRequest()
@@ -95,23 +92,21 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
         handleDistributedRequest()
         String msg = "Hello."
         response.status(200)
-          .addHeader('x-datadog-test-response-header', 'baz')
-          .send(msg)
+        .addHeader('x-datadog-test-response-header', 'baz')
+        .send(msg)
       }
       prefix("/timeout") {
         Thread.sleep(10_000)
         throw new IllegalStateException("Should never happen")
       }
       prefix("/json") {
-        if (request.getContentType() != 'application/json') {
-          response.status(400).send('Bad content type')
-        } else {
-          response
-            .status(200)
-            .addHeader('Content-Type', 'application/json')
-            .addHeader('X-AppSec-Test', 'true')
-            .sendWithType('application/json', request.body)
-        }
+        // echo if input is json
+        final responseBody = request.getContentType() == 'application/json' ? request.body : '{"goodbye": "world!"}'.bytes
+        response
+        .status(200)
+        .addHeader('Content-Type', 'application/json')
+        .addHeader('X-AppSec-Test', 'true')
+        .sendWithType('application/json', responseBody)
       }
     }
   }
@@ -146,19 +141,19 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
   def setupSpec() {
     List<Proxy> proxyList = Collections.singletonList(new Proxy(Proxy.Type.HTTP, new InetSocketAddress(proxy.port)))
     proxySelector = new ProxySelector() {
-        @Override
-        List<Proxy> select(URI uri) {
-          if (uri.fragment == "proxy") {
-            return proxyList
-          }
-          return getDefault().select(uri)
+      @Override
+      List<Proxy> select(URI uri) {
+        if (uri.fragment == "proxy") {
+          return proxyList
         }
+        return getDefault().select(uri)
+      }
 
-        @Override
-        void connectFailed(URI uri, SocketAddress sa, IOException ioe) {
-          getDefault().connectFailed(uri, sa, ioe)
-        }
+      @Override
+      void connectFailed(URI uri, SocketAddress sa, IOException ioe) {
+        getDefault().connectFailed(uri, sa, ioe)
       }
+    }
 
     // Register the Instrumentation Gateway callbacks
     def ss = get().getSubscriptionService(RequestContextSlot.APPSEC)
@@ -910,16 +905,9 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
   void 'test appsec client request analysis'() {
     given:
     final url = server.address.resolve(endpoint)
-    final tags = [
-      'downstream.request.url': url.toString(),
-      'downstream.request.method': method,
-      'downstream.request.body': body,
-      'downstream.response.status': 200,
-      'downstream.response.body': body,
-    ]
 
     when:
-    final status = runUnderAppSecTrace {
+    def (ctx, status) = runUnderAppSecTrace {
       doRequest(method, url, ['Content-Type': contentType] + headers, body) {
         InputStream response ->
         assert response.text == body
@@ -928,25 +916,66 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
 
     then:
     status == 200
-    TEST_WRITER.waitForTraces(1)
-    final span = TEST_WRITER.get(0).find {
-      it.spanType == 'http'
-    }
-    tags.each {
-      assert span.getTag(it.key) == it.value
+    final request = ctx.requests.first()
+    request.method == method
+    request.url == url.toString()
+    request.body.bytes == body.bytes
+    headers.each {
+      assert request.headers[it.key] == [it.value]
     }
-    final requestHeaders = new JsonSlurper().parseText(span.getTag("downstream.request.headers") as String) as Map<String, List<String>>
-    final responseHeaders = new JsonSlurper().parseText(span.getTag("downstream.response.headers") as String) as Map<String, List<String>>
+
+    final response = ctx.responses.first()
+    response.status == 200
+    response.body.bytes == body.bytes
     headers.each {
-      assert requestHeaders[it.key] == [it.value]
-      assert responseHeaders[it.key] == [it.value]
+      assert response.headers[it.key] == [it.value]
     }
 
     where:
     endpoint | method | contentType        | headers                   | body
     '/json'  | 'POST' | 'application/json' | ['X-AppSec-Test': 'true'] | '{"hello": "world!" }'
   }
 
+  @IgnoreIf({
+    !instance.testAppSecClientRedirect()
+  })
+  void 'test appsec client redirect analysis'() {
+    given:
+    final url = server.address.resolve(endpoint)
+
+    when:
+    def (ctx, status) = runUnderAppSecTrace {
+      doRequest(method, url, ['Content-Type': contentType] + headers, requestBody)
+    }
+
+    then:
+    status == 200
+
+    def (initialRequest, redirectRequest) = ctx.requests
+    initialRequest.method == method
+    initialRequest.url == url.toString()
+    initialRequest.body.bytes == requestBody.bytes
+    headers.each {
+      assert initialRequest.headers[it.key] == [it.value]
+    }
+
+    redirectRequest.method == 'GET'
+    redirectRequest.url.toString().endsWith('/json')
+    redirectRequest.body == null
+
+    def (redirectResponse, finalResponse) = ctx.responses
+    redirectResponse.status == 302
+    redirectResponse.body == null
+    redirectResponse.headers['Location'][0].endsWith('/json')
+
+    finalResponse.status == 200
+    finalResponse.body.bytes == responseBody.bytes
+
+    where:
+    endpoint    | method | contentType        | headers                                        | requestBody            | responseBody
+    '/redirect' | 'POST' | 'application/json' | ['X-AppSec-Test': 'true', 'Location': '/json'] | '{"hello": "world!" }' | '{"goodbye": "world!"}'
+  }
+
   // parent span must be cast otherwise it breaks debugging classloading (junit loads it early)
   void clientSpan(
   TraceAssert trace,
@@ -1070,11 +1099,16 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
     false
   }
 
-  protected <E> E runUnderAppSecTrace(Closure<E> cl) {
-    final ddctx = new TagContext().withRequestContextDataAppSec(new IGCallbacks.Context())
+  boolean testAppSecClientRedirect() {
+    false
+  }
+
+  protected <E> Tuple2<IGCallbacks.Context, E> runUnderAppSecTrace(Closure<E> cl) {
+    final ctx = new IGCallbacks.Context()
+    final ddctx = new TagContext().withRequestContextDataAppSec(ctx)
     final span = TEST_TRACER.startSpan("test", "test-appsec-span", ddctx)
     try {
-      return AgentTracer.activateSpan(span).withCloseable(cl)
+      return Tuple.tuple(ctx, AgentTracer.activateSpan(span).withCloseable(cl))
     } finally {
       span.finish()
     }
@@ -1084,6 +1118,8 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
 
     static class Context {
       boolean hasAppSecData
+      List<HttpClientRequest> requests = []
+      List<HttpClientResponse> responses = []
     }
 
     final BiFunction<RequestContext, Long, Flow<Boolean>> httpClientBodySamplingCb = {
@@ -1093,16 +1129,11 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
 
     final BiFunction<RequestContext, HttpClientRequest, Flow<Void>> httpClientRequestCb = {
       RequestContext rqCtxt, HttpClientRequest req ->
-      if (req.headers?.containsKey('X-AppSec-Test')) {
-        final context = rqCtxt.getData(RequestContextSlot.APPSEC) as Context
-        if (context != null) {
-          context.hasAppSecData = true
-          activeSpan()
-          .setTag('downstream.request.url', req.url)
-          .setTag('downstream.request.method', req.method)
-          .setTag('downstream.request.headers', JsonOutput.toJson(req.headers))
-          .setTag('downstream.request.body', req.body?.text)
-        }
+      final context = rqCtxt.getData(RequestContextSlot.APPSEC) as Context
+      final boolean isAppSec = req.headers?.containsKey('X-AppSec-Test')
+      if (isAppSec || context?.hasAppSecData) {
+        context.hasAppSecData = true
+        context.requests.add(req)
       }
       Flow.ResultFlow.empty()
     } as BiFunction<RequestContext, HttpClientRequest, Flow<Void>>
@@ -1111,10 +1142,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
       RequestContext rqCtxt, HttpClientResponse res ->
       final context = rqCtxt.getData(RequestContextSlot.APPSEC) as Context
       if (context?.hasAppSecData) {
-        activeSpan()
-        .setTag('downstream.response.status', res.status)
-        .setTag('downstream.response.headers', JsonOutput.toJson(res.headers))
-        .setTag('downstream.response.body', res.body?.text)
+        context.responses.add(res)
       }
       Flow.ResultFlow.empty()
     } as BiFunction<RequestContext, HttpClientResponse, Flow<Void>>

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

@@ -0,0 +1,69 @@
+package datadog.trace.instrumentation.okhttp2;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import com.squareup.okhttp.Request;
+import com.squareup.okhttp.Response;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+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 net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public class AppSecHttpEngineInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public AppSecHttpEngineInstrumentation() {
+    super("okhttp", "okhttp-2");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "com.squareup.okhttp.internal.http.HttpEngine";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".AppSecInterceptor",
+    };
+  }
+
+  @Override
+  public void methodAdvice(final MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(named("sendRequest")).and(takesArguments(0)),
+        AppSecHttpEngineInstrumentation.class.getName() + "$SendRequestAdvice");
+  }
+
+  public static class SendRequestAdvice {
+    @Advice.OnMethodEnter
+    public static void onSendRequest(
+        @Advice.FieldValue("priorResponse") final Response priorResponse,
+        @Advice.FieldValue("userRequest") final Request userRequest) {
+      // only redirects
+      if (priorResponse == null || priorResponse.code() < 300 || priorResponse.code() >= 400) {
+        return;
+      }
+      final AgentSpan span = AgentTracer.activeSpan();
+      final RequestContext ctx = span.getRequestContext();
+      if (ctx == null) {
+        return;
+      }
+      if (ctx.getData(RequestContextSlot.APPSEC) == null) {
+        return;
+      }
+
+      // 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);
+    }
+  }
+}