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 fbca97cd8e2..9a04c87ae04 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 @@ -152,7 +152,15 @@ public static Response onResponse( } } - publish(ctx, clientResponse, responseCb); + try { + publish(ctx, clientResponse, responseCb); + } catch (final BlockingException e) { + throw e; + } catch (final Exception e) { + // the response body has already been captured into `result` above; don't lose it by + // letting the caller fall back to the original (possibly already-consumed) response + LOGGER.debug("Failed to run AppSec response hooks", e); + } return result; } diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptorTest.java b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptorTest.java index 3319774af1f..c5eea2a49c3 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptorTest.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptorTest.java @@ -1,5 +1,7 @@ package datadog.trace.instrumentation.okhttp2; +import static datadog.trace.api.gateway.Events.EVENTS; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; @@ -9,8 +11,13 @@ import static org.mockito.Mockito.when; import com.squareup.okhttp.Interceptor; +import com.squareup.okhttp.MediaType; +import com.squareup.okhttp.Protocol; import com.squareup.okhttp.Request; +import com.squareup.okhttp.Response; +import com.squareup.okhttp.ResponseBody; import datadog.trace.api.gateway.CallbackProvider; +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; @@ -65,4 +72,32 @@ void ioExceptionFromProceedPropagatesWithoutRetry() throws IOException { assertSame(failure, thrown); verify(chain, times(1)).proceed(request); } + + @Test + void responseHookFailureAfterBodyCapturePreservesCapturedBody() throws IOException { + final String json = "{\"hello\":\"world\"}"; + final Response response = + new Response.Builder() + .request(request) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(MediaType.parse("application/json"), json)) + .build(); + when(chain.proceed(request)).thenReturn(response); + + final CallbackProvider cbp = mock(CallbackProvider.class); + when(cbp.getCallback(EVENTS.httpClientSampling())) + .thenReturn((ctx, requestId) -> new Flow.ResultFlow<>(Boolean.TRUE)); + when(cbp.getCallback(EVENTS.httpClientResponse())) + .thenReturn( + (ctx, clientResponse) -> { + throw new RuntimeException("boom"); + }); + when(AgentTracer.get().getCallbackProvider(any(RequestContextSlot.class))).thenReturn(cbp); + + final Response result = interceptor.intercept(chain); + + assertEquals(json, result.body().string()); + } } 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 abb3c7aba83..218058b6e79 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 @@ -152,7 +152,15 @@ public static Response onResponse( } } - publish(ctx, clientResponse, responseCb); + try { + publish(ctx, clientResponse, responseCb); + } catch (final BlockingException e) { + throw e; + } catch (final Exception e) { + // the response body has already been captured into `result` above; don't lose it by + // letting the caller fall back to the original (possibly already-consumed) response + LOGGER.debug("Failed to run AppSec response hooks", e); + } return result; } diff --git a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptorTest.java b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptorTest.java index 56b230bb849..d56772f3de7 100644 --- a/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptorTest.java +++ b/dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptorTest.java @@ -1,5 +1,7 @@ package datadog.trace.instrumentation.okhttp3; +import static datadog.trace.api.gateway.Events.EVENTS; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; @@ -9,6 +11,7 @@ import static org.mockito.Mockito.when; import datadog.trace.api.gateway.CallbackProvider; +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; @@ -16,7 +19,11 @@ import datadog.trace.bootstrap.instrumentation.api.Tags; import java.io.IOException; import okhttp3.Interceptor; +import okhttp3.MediaType; +import okhttp3.Protocol; import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -65,4 +72,32 @@ void ioExceptionFromProceedPropagatesWithoutRetry() throws IOException { assertSame(failure, thrown); verify(chain, times(1)).proceed(request); } + + @Test + void responseHookFailureAfterBodyCapturePreservesCapturedBody() throws IOException { + final String json = "{\"hello\":\"world\"}"; + final Response response = + new Response.Builder() + .request(request) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(MediaType.parse("application/json"), json)) + .build(); + when(chain.proceed(request)).thenReturn(response); + + final CallbackProvider cbp = mock(CallbackProvider.class); + when(cbp.getCallback(EVENTS.httpClientSampling())) + .thenReturn((ctx, requestId) -> new Flow.ResultFlow<>(Boolean.TRUE)); + when(cbp.getCallback(EVENTS.httpClientResponse())) + .thenReturn( + (ctx, clientResponse) -> { + throw new RuntimeException("boom"); + }); + when(AgentTracer.get().getCallbackProvider(any(RequestContextSlot.class))).thenReturn(cbp); + + final Response result = interceptor.intercept(chain); + + assertEquals(json, result.body().string()); + } }