From 4ba4ebf36c60536680911b7b8453d01524cee062 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Fri, 19 Dec 2025 11:39:24 +0100 Subject: [PATCH 1/2] feat(aap): analyze okhttp client redirections as separated request --- .../agent/test/base/HttpClientTest.groovy | 146 +++++++++++------- .../AppSecHttpEngineInstrumentation.java | 68 ++++++++ .../okhttp2/AppSecInterceptor.java | 27 ++-- .../src/test/groovy/HeadersUtil.groovy | 26 +++- .../src/test/groovy/OkHttp2Test.groovy | 5 + .../AppSecHttpEngineInstrumentation.java | 68 ++++++++ .../okhttp3/AppSecInterceptor.java | 27 ++-- .../src/test/groovy/OkHttp3Test.groovy | 5 + 8 files changed, 289 insertions(+), 83 deletions(-) create mode 100644 dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecHttpEngineInstrumentation.java create mode 100644 dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecHttpEngineInstrumentation.java diff --git a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy index 229def7d5f8..70cb5e33c38 100644 --- a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy +++ b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy @@ -15,6 +15,7 @@ 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 @@ -69,7 +70,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 +96,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 +145,19 @@ abstract class HttpClientTest extends VersionedNamingTestBase { def setupSpec() { List proxyList = Collections.singletonList(new Proxy(Proxy.Type.HTTP, new InetSocketAddress(proxy.port))) proxySelector = new ProxySelector() { - @Override - List select(URI uri) { - if (uri.fragment == "proxy") { - return proxyList - } - return getDefault().select(uri) + @Override + List 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 +909,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,18 +920,19 @@ 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> - final responseHeaders = new JsonSlurper().parseText(span.getTag("downstream.response.headers") as String) as Map> + + 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: @@ -947,6 +940,46 @@ abstract class HttpClientTest extends VersionedNamingTestBase { '/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 +1103,16 @@ abstract class HttpClientTest extends VersionedNamingTestBase { false } - protected E runUnderAppSecTrace(Closure cl) { - final ddctx = new TagContext().withRequestContextDataAppSec(new IGCallbacks.Context()) + boolean testAppSecClientRedirect() { + false + } + + protected Tuple2 runUnderAppSecTrace(Closure 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 +1122,8 @@ abstract class HttpClientTest extends VersionedNamingTestBase { static class Context { boolean hasAppSecData + List requests = [] + List responses = [] } final BiFunction> httpClientBodySamplingCb = { @@ -1093,16 +1133,11 @@ abstract class HttpClientTest extends VersionedNamingTestBase { final BiFunction> 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> @@ -1111,10 +1146,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> diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecHttpEngineInstrumentation.java b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecHttpEngineInstrumentation.java new file mode 100644 index 00000000000..c878f74c928 --- /dev/null +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecHttpEngineInstrumentation.java @@ -0,0 +1,68 @@ +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; + } + + // do not include bodies in the redirect request + AppSecInterceptor.onResponse(span, false, priorResponse); + AppSecInterceptor.onRequest(span, false, userRequest.urlString(), userRequest); + } + } +} diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java index 1e926270c41..d568c7262c4 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java @@ -53,7 +53,8 @@ public Response intercept(final Chain chain) throws IOException { } final long requestId = span.getSpanId(); final boolean sampled = sampleRequest(ctx, requestId); - final Request request = onRequest(span, sampled, chain.request()); + final String url = span.getTag(Tags.HTTP_URL).toString(); + final Request request = onRequest(span, sampled, url, chain.request()); final Response response = chain.proceed(request); return onResponse(span, sampled, response); } catch (final BlockingException e) { @@ -64,7 +65,8 @@ public Response intercept(final Chain chain) throws IOException { } } - private Request onRequest(final AgentSpan span, final boolean sampled, final Request request) { + public static Request onRequest( + final AgentSpan span, final boolean sampled, final String url, final Request request) { Request result = request; CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); BiFunction> requestCb = @@ -76,7 +78,6 @@ private Request onRequest(final AgentSpan span, final boolean sampled, final Req final RequestBody requestBody = request.body(); final RequestContext ctx = span.getRequestContext(); final long requestId = span.getSpanId(); - final String url = span.getTag(Tags.HTTP_URL).toString(); final HttpClientRequest clientRequest = new HttpClientRequest(requestId, url, request.method(), mapHeaders(request.headers())); if (sampled && requestBody != null) { @@ -104,7 +105,7 @@ private Request onRequest(final AgentSpan span, final boolean sampled, final Req return result; } - private Response onResponse( + public static Response onResponse( final AgentSpan span, final boolean sampled, final Response response) { Response result = response; CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); @@ -145,7 +146,7 @@ private Response onResponse( return result; } - private

void publish( + private static

void publish( final RequestContext ctx, final P request, final BiFunction> callback) { @@ -161,7 +162,7 @@ private

void publish( } } - private boolean sampleRequest(final RequestContext ctx, final long requestId) { + 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> samplingCb = @@ -177,7 +178,7 @@ private boolean sampleRequest(final RequestContext ctx, final long requestId) { * Ensure we are only consuming payloads we can safely deserialize with a bounded size to prevent * from OOM */ - private boolean shouldProcessBody(final long contentLength, final MediaType mediaType) { + private static boolean shouldProcessBody(final long contentLength, final MediaType mediaType) { if (contentLength <= 0) { return false; // prevent from copying from unbounded source (just to be safe) } @@ -190,7 +191,8 @@ private boolean shouldProcessBody(final long contentLength, final MediaType medi return mediaType.isDeserializable(); } - private byte[] readBody(final RequestBody body, final int contentLength) throws IOException { + private static byte[] readBody(final RequestBody body, final int contentLength) + throws IOException { final ByteArrayOutputStream buffer = new ByteArrayOutputStream(contentLength); try (final BufferedSink sink = Okio.buffer(Okio.sink(buffer))) { body.writeTo(sink); @@ -198,7 +200,8 @@ private byte[] readBody(final RequestBody body, final int contentLength) throws return buffer.toByteArray(); } - private byte[] readBody(final ResponseBody body, final int contentLength) throws IOException { + private static byte[] readBody(final ResponseBody body, final int contentLength) + throws IOException { final ByteArrayOutputStream buffer = new ByteArrayOutputStream(contentLength); try (final BufferedSource source = body.source(); final Sink sink = Okio.sink(buffer)) { @@ -207,7 +210,7 @@ private byte[] readBody(final ResponseBody body, final int contentLength) throws return buffer.toByteArray(); } - private Map> mapHeaders(final Headers headers) { + private static Map> mapHeaders(final Headers headers) { if (headers == null) { return Collections.emptyMap(); } @@ -218,12 +221,12 @@ private Map> mapHeaders(final Headers headers) { return result; } - private MediaType contentType(final RequestBody body) { + private static MediaType contentType(final RequestBody body) { return MediaType.parse( body == null || body.contentType() == null ? null : body.contentType().toString()); } - private MediaType contentType(final ResponseBody body) { + private static MediaType contentType(final ResponseBody body) { return MediaType.parse( body == null || body.contentType() == null ? null : body.contentType().toString()); } diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/HeadersUtil.groovy b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/HeadersUtil.groovy index 5940d375211..0eaee8baf36 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/HeadersUtil.groovy +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/HeadersUtil.groovy @@ -1,11 +1,33 @@ +import spock.lang.Specification + class HeadersUtil { static headersToArray(Map headers) { String[] headersArr = new String[headers.size() * 2] headers.eachWithIndex { k, v, i -> - headersArr[i] = k - headersArr[i + 1] = v + headersArr[i * 2] = k + headersArr[i * 2 + 1] = v } headersArr } } + +class HeadersUtilTest extends Specification { + + void 'test headers to array'() { + setup: + final array = expected as String[] + + when: + final result = HeadersUtil.headersToArray(heaaders) + + then: + result == array + + where: + heaaders | expected + [:] | [] + ['a': 'b'] | ['a', 'b'] + ['a': 'b', 'c': 'd'] | ['a', 'b', 'c', 'd'] + } +} diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/OkHttp2Test.groovy b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/OkHttp2Test.groovy index 7c67fc888d1..b2e24739fb3 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/OkHttp2Test.groovy +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/groovy/OkHttp2Test.groovy @@ -53,6 +53,11 @@ abstract class OkHttp2Test extends HttpClientTest { boolean testAppSecClientRequest() { true } + + @Override + boolean testAppSecClientRedirect() { + true + } } @Timeout(5) diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecHttpEngineInstrumentation.java b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecHttpEngineInstrumentation.java new file mode 100644 index 00000000000..6743b93c9aa --- /dev/null +++ b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecHttpEngineInstrumentation.java @@ -0,0 +1,68 @@ +package datadog.trace.instrumentation.okhttp3; + +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 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; +import okhttp3.Request; +import okhttp3.Response; + +@AutoService(InstrumenterModule.class) +public class AppSecHttpEngineInstrumentation extends InstrumenterModule.AppSec + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + + public AppSecHttpEngineInstrumentation() { + super("okhttp", "okhttp-3"); + } + + @Override + public String instrumentedType() { + return "okhttp3.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; + } + + // do not include bodies in the redirect request + AppSecInterceptor.onResponse(span, false, priorResponse); + AppSecInterceptor.onRequest(span, false, userRequest.url().toString(), userRequest); + } + } +} diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java index 210dbd4c9e5..c411df3d52d 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java @@ -53,7 +53,8 @@ public Response intercept(final Chain chain) throws IOException { } final long requestId = span.getSpanId(); final boolean sampled = sampleRequest(ctx, requestId); - final Request request = onRequest(span, sampled, chain.request()); + final String url = span.getTag(Tags.HTTP_URL).toString(); + final Request request = onRequest(span, sampled, url, chain.request()); final Response response = chain.proceed(request); return onResponse(span, sampled, response); } catch (final BlockingException e) { @@ -64,7 +65,8 @@ public Response intercept(final Chain chain) throws IOException { } } - private Request onRequest(final AgentSpan span, final boolean sampled, final Request request) { + public static Request onRequest( + final AgentSpan span, final boolean sampled, final String url, final Request request) { Request result = request; CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); BiFunction> requestCb = @@ -76,7 +78,6 @@ private Request onRequest(final AgentSpan span, final boolean sampled, final Req final RequestBody requestBody = request.body(); final RequestContext ctx = span.getRequestContext(); final long requestId = span.getSpanId(); - final String url = span.getTag(Tags.HTTP_URL).toString(); final HttpClientRequest clientRequest = new HttpClientRequest(requestId, url, request.method(), mapHeaders(request.headers())); if (sampled && requestBody != null) { @@ -104,7 +105,7 @@ private Request onRequest(final AgentSpan span, final boolean sampled, final Req return result; } - private Response onResponse( + public static Response onResponse( final AgentSpan span, final boolean sampled, final Response response) { Response result = response; CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC); @@ -145,7 +146,7 @@ private Response onResponse( return result; } - private

void publish( + private static

void publish( final RequestContext ctx, final P request, final BiFunction> callback) { @@ -161,7 +162,7 @@ private

void publish( } } - private boolean sampleRequest(final RequestContext ctx, final long requestId) { + 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> samplingCb = @@ -177,7 +178,7 @@ private boolean sampleRequest(final RequestContext ctx, final long requestId) { * Ensure we are only consuming payloads we can safely deserialize with a bounded size to prevent * from OOM */ - private boolean shouldProcessBody(final long contentLength, final MediaType mediaType) { + private static boolean shouldProcessBody(final long contentLength, final MediaType mediaType) { if (contentLength <= 0) { return false; // prevent from copying from unbounded source (just to be safe) } @@ -190,7 +191,8 @@ private boolean shouldProcessBody(final long contentLength, final MediaType medi return mediaType.isDeserializable(); } - private byte[] readBody(final RequestBody body, final int contentLength) throws IOException { + private static byte[] readBody(final RequestBody body, final int contentLength) + throws IOException { final ByteArrayOutputStream buffer = new ByteArrayOutputStream(contentLength); try (final BufferedSink sink = Okio.buffer(Okio.sink(buffer))) { body.writeTo(sink); @@ -198,7 +200,8 @@ private byte[] readBody(final RequestBody body, final int contentLength) throws return buffer.toByteArray(); } - private byte[] readBody(final ResponseBody body, final int contentLength) throws IOException { + private static byte[] readBody(final ResponseBody body, final int contentLength) + throws IOException { final ByteArrayOutputStream buffer = new ByteArrayOutputStream(contentLength); try (final BufferedSource source = body.source(); final Sink sink = Okio.sink(buffer)) { @@ -207,7 +210,7 @@ private byte[] readBody(final ResponseBody body, final int contentLength) throws return buffer.toByteArray(); } - private Map> mapHeaders(final Headers headers) { + private static Map> mapHeaders(final Headers headers) { if (headers == null) { return Collections.emptyMap(); } @@ -218,12 +221,12 @@ private Map> mapHeaders(final Headers headers) { return result; } - private MediaType contentType(final RequestBody body) { + private static MediaType contentType(final RequestBody body) { return MediaType.parse( body == null || body.contentType() == null ? null : body.contentType().toString()); } - private MediaType contentType(final ResponseBody body) { + private static MediaType contentType(final ResponseBody body) { return MediaType.parse( body == null || body.contentType() == null ? null : body.contentType().toString()); } diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/groovy/OkHttp3Test.groovy b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/groovy/OkHttp3Test.groovy index ffaba2dd49d..903bf436c56 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/groovy/OkHttp3Test.groovy +++ b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/groovy/OkHttp3Test.groovy @@ -27,6 +27,11 @@ abstract class OkHttp3Test extends HttpClientTest { true } + @Override + boolean testAppSecClientRedirect() { + true + } + @Override boolean useStrictTraceWrites() { // TODO fix this by making sure that spans get closed properly From 709a228d2fe0263807c06aac43b08780d76ceea6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Fri, 19 Dec 2025 16:30:43 +0100 Subject: [PATCH 2/2] Update tests --- .../appsec/gateway/AppSecRequestContext.java | 3 +++ .../agent/test/base/HttpClientTest.groovy | 4 --- .../AppSecHttpEngineInstrumentation.java | 3 ++- .../okhttp2/AppSecInterceptor.java | 2 +- .../AppSecHttpEngineInstrumentation.java | 3 ++- .../okhttp3/AppSecInterceptor.java | 2 +- .../springboot/controller/WebController.java | 23 +++++++++++++--- .../appsec/SpringBootSmokeTest.groovy | 27 ++++++++++++++++--- 8 files changed, 52 insertions(+), 15 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index d4a19e7f7ec..2b197756697 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/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); diff --git a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy index 70cb5e33c38..e787feeda1f 100644 --- a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy +++ b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy @@ -15,7 +15,6 @@ 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 @@ -23,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 @@ -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 { diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecHttpEngineInstrumentation.java b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecHttpEngineInstrumentation.java index c878f74c928..8fd1e170686 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecHttpEngineInstrumentation.java +++ b/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); } diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java index d568c7262c4..7f55cc9a4fa 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/main/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptor.java @@ -162,7 +162,7 @@ private static

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> samplingCb = diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecHttpEngineInstrumentation.java b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecHttpEngineInstrumentation.java index 6743b93c9aa..0b7aa00478f 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecHttpEngineInstrumentation.java +++ b/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); } diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java index c411df3d52d..e61a78003bb 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptor.java @@ -162,7 +162,7 @@ private static

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> samplingCb = diff --git a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java index cee5f469848..59f5d9a75dc 100644 --- a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java +++ b/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> apiSecurityResponse( public ResponseEntity 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 apiSecurityHttpClientOkHttp2(final HttpServletRequ public ResponseEntity 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 apiSecurityHttpClientOkHttp3(final HttpServletRequ @RequestMapping( value = "/echo", method = {POST, GET, PUT}) - public ResponseEntity echo(final HttpServletRequest request) throws IOException { + public ResponseEntity 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); diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy index 1254535c8d7..79048bd5c18 100644 --- a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy +++ b/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 }