Add regression tests for AppSec response body preservation on hook failure

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());
+  }
 }

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,14 +11,19 @@
 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;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 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());
+  }
 }