Stop AppSecInterceptor from silently retrying failed okhttp requests

dd-java-agent/instrumentation/okhttp/okhttp-2.2/build.gradle

@@ -41,6 +41,9 @@ dependencies {
   }
   testImplementation group: 'com.squareup.okhttp', name: 'okhttp', version: '2.2.0'
 
+  testImplementation libs.bundles.junit5
+  testImplementation libs.bundles.mockito
+
   testRuntimeOnly(project(':dd-java-agent:instrumentation:datadog:asm:iast-instrumenter'))
   testRuntimeOnly(project(':dd-java-agent:instrumentation:java:java-net:java-net-1.8'))
 

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

@@ -45,23 +45,32 @@ public class AppSecInterceptor implements Interceptor {
 
   @Override
   public Response intercept(final Chain chain) throws IOException {
+    Request request = chain.request();
+    final AgentSpan span = AgentTracer.activeSpan();
+    final RequestContext ctx = span == null ? null : span.getRequestContext();
+    if (ctx == null) {
+      return chain.proceed(request);
+    }
+    boolean sampled = false;
     try {
-      final AgentSpan span = AgentTracer.activeSpan();
-      final RequestContext ctx = span == null ? null : span.getRequestContext();
-      if (ctx == null) {
-        return chain.proceed(chain.request());
-      }
       final long requestId = span.getSpanId();
-      final boolean sampled = sampleRequest(ctx, requestId);
+      sampled = sampleRequest(ctx, requestId);
       final String url = span.getTag(Tags.HTTP_URL).toString();
-      final Request request = onRequest(span, sampled, url, chain.request());
-      final Response response = chain.proceed(request);
+      request = onRequest(span, sampled, url, request);
+    } catch (final BlockingException e) {
+      throw e;
+    } catch (final Exception e) {
+      LOGGER.debug("Failed to run AppSec request hooks", e);
+    }
+    // let real connection/IO failures propagate rather than swallowing and retrying the request
+    final Response response = chain.proceed(request);
+    try {
       return onResponse(span, sampled, response);
     } catch (final BlockingException e) {
       throw e;
     } catch (final Exception e) {
-      LOGGER.debug("Failed to intercept request", e);
-      return chain.proceed(chain.request());
+      LOGGER.debug("Failed to run AppSec response hooks", e);
+      return response;
     }
   }
 

dd-java-agent/instrumentation/okhttp/okhttp-2.2/src/test/java/datadog/trace/instrumentation/okhttp2/AppSecInterceptorTest.java

@@ -0,0 +1,68 @@
+package datadog.trace.instrumentation.okhttp2;
+
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.squareup.okhttp.Interceptor;
+import com.squareup.okhttp.Request;
+import datadog.trace.api.gateway.CallbackProvider;
+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 org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class AppSecInterceptorTest {
+
+  private final AgentTracer.TracerAPI originalTracer = AgentTracer.get();
+
+  private Interceptor.Chain chain;
+  private Request request;
+  private final AppSecInterceptor interceptor = new AppSecInterceptor();
+
+  @BeforeEach
+  void setup() {
+    request = new Request.Builder().url("http://example.com").build();
+
+    final RequestContext requestContext = mock(RequestContext.class);
+
+    final AgentSpan span = mock(AgentSpan.class);
+    when(span.getRequestContext()).thenReturn(requestContext);
+    when(span.getSpanId()).thenReturn(1L);
+    when(span.getTag(Tags.HTTP_URL)).thenReturn("http://example.com");
+
+    final AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class);
+    when(tracer.activeSpan()).thenReturn(span);
+    when(tracer.getCallbackProvider(any(RequestContextSlot.class)))
+        .thenReturn(CallbackProvider.CallbackProviderNoop.INSTANCE);
+    AgentTracer.forceRegister(tracer);
+
+    chain = mock(Interceptor.Chain.class);
+    when(chain.request()).thenReturn(request);
+  }
+
+  @AfterEach
+  void tearDown() {
+    AgentTracer.forceRegister(originalTracer);
+  }
+
+  @Test
+  void ioExceptionFromProceedPropagatesWithoutRetry() throws IOException {
+    final IOException failure = new IOException("boom");
+    when(chain.proceed(request)).thenThrow(failure);
+
+    final IOException thrown = assertThrows(IOException.class, () -> interceptor.intercept(chain));
+
+    assertSame(failure, thrown);
+    verify(chain, times(1)).proceed(request);
+  }
+}

dd-java-agent/instrumentation/okhttp/okhttp-3.0/build.gradle

@@ -38,6 +38,9 @@ dependencies {
   latestDepTestImplementation group: 'com.squareup.okhttp3', name: 'okhttp', version: '[3.11.0, 4)'
   latestDepTestImplementation group: 'com.squareup.okio', name: 'okio', version: '1.+'
 
+  testImplementation libs.bundles.junit5
+  testImplementation libs.bundles.mockito
+
   testRuntimeOnly(project(':dd-java-agent:instrumentation:datadog:asm:iast-instrumenter'))
   testRuntimeOnly(project(':dd-java-agent:instrumentation:java:java-net:java-net-1.8'))
 }

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

@@ -45,23 +45,32 @@ public class AppSecInterceptor implements Interceptor {
 
   @Override
   public Response intercept(final Chain chain) throws IOException {
+    Request request = chain.request();
+    final AgentSpan span = AgentTracer.activeSpan();
+    final RequestContext ctx = span == null ? null : span.getRequestContext();
+    if (ctx == null) {
+      return chain.proceed(request);
+    }
+    boolean sampled = false;
     try {
-      final AgentSpan span = AgentTracer.activeSpan();
-      final RequestContext ctx = span == null ? null : span.getRequestContext();
-      if (ctx == null) {
-        return chain.proceed(chain.request());
-      }
       final long requestId = span.getSpanId();
-      final boolean sampled = sampleRequest(ctx, requestId);
+      sampled = sampleRequest(ctx, requestId);
       final String url = span.getTag(Tags.HTTP_URL).toString();
-      final Request request = onRequest(span, sampled, url, chain.request());
-      final Response response = chain.proceed(request);
+      request = onRequest(span, sampled, url, request);
+    } catch (final BlockingException e) {
+      throw e;
+    } catch (final Exception e) {
+      LOGGER.debug("Failed to run AppSec request hooks", e);
+    }
+    // let real connection/IO failures propagate rather than swallowing and retrying the request
+    final Response response = chain.proceed(request);
+    try {
       return onResponse(span, sampled, response);
     } catch (final BlockingException e) {
       throw e;
     } catch (final Exception e) {
-      LOGGER.debug("Failed to intercept request", e);
-      return chain.proceed(chain.request());
+      LOGGER.debug("Failed to run AppSec response hooks", e);
+      return response;
     }
   }
 

dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/test/java/datadog/trace/instrumentation/okhttp3/AppSecInterceptorTest.java

@@ -0,0 +1,68 @@
+package datadog.trace.instrumentation.okhttp3;
+
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import datadog.trace.api.gateway.CallbackProvider;
+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.Request;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class AppSecInterceptorTest {
+
+  private final AgentTracer.TracerAPI originalTracer = AgentTracer.get();
+
+  private Interceptor.Chain chain;
+  private Request request;
+  private final AppSecInterceptor interceptor = new AppSecInterceptor();
+
+  @BeforeEach
+  void setup() {
+    request = new Request.Builder().url("http://example.com").build();
+
+    final RequestContext requestContext = mock(RequestContext.class);
+
+    final AgentSpan span = mock(AgentSpan.class);
+    when(span.getRequestContext()).thenReturn(requestContext);
+    when(span.getSpanId()).thenReturn(1L);
+    when(span.getTag(Tags.HTTP_URL)).thenReturn("http://example.com");
+
+    final AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class);
+    when(tracer.activeSpan()).thenReturn(span);
+    when(tracer.getCallbackProvider(any(RequestContextSlot.class)))
+        .thenReturn(CallbackProvider.CallbackProviderNoop.INSTANCE);
+    AgentTracer.forceRegister(tracer);
+
+    chain = mock(Interceptor.Chain.class);
+    when(chain.request()).thenReturn(request);
+  }
+
+  @AfterEach
+  void tearDown() {
+    AgentTracer.forceRegister(originalTracer);
+  }
+
+  @Test
+  void ioExceptionFromProceedPropagatesWithoutRetry() throws IOException {
+    final IOException failure = new IOException("boom");
+    when(chain.proceed(request)).thenThrow(failure);
+
+    final IOException thrown = assertThrows(IOException.class, () -> interceptor.intercept(chain));
+
+    assertSame(failure, thrown);
+    verify(chain, times(1)).proceed(request);
+  }
+}