Ensure blocking exceptions are propagated for SSRF (#9790)

Author: manuel-alvarez-alvarez

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpClientDecorator.java

@@ -104,10 +104,13 @@ public AgentSpan onRequest(final AgentSpan span, final REQUEST request) {
         } else if (shouldSetResourceName()) {
           span.setResourceName(DEFAULT_RESOURCE_NAME);
         }
+      } catch (final BlockingException e) {
+        throw e;
       } catch (final Exception e) {
         log.debug("Error tagging url", e);
+      } finally {
+        ssrfIastCheck(request);
       }
-      ssrfIastCheck(request);
     }
     return span;
   }

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

@@ -56,6 +56,8 @@ public Response intercept(final Chain chain) throws IOException {
       final Request request = onRequest(span, sampled, chain.request());
       final Response response = chain.proceed(request);
       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());

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

@@ -56,6 +56,8 @@ public Response intercept(final Chain chain) throws IOException {
       final Request request = onRequest(span, sampled, chain.request());
       final Response response = chain.proceed(request);
       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());

dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java

@@ -8,6 +8,7 @@
 import com.squareup.okhttp.OkHttpClient;
 import com.squareup.okhttp.Request;
 import com.squareup.okhttp.Response;
+import datadog.appsec.api.blocking.BlockingException;
 import datadog.smoketest.appsec.springboot.service.AsyncService;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
@@ -93,7 +94,9 @@ public String sqliHeader(@RequestHeader("x-custom-header") String id) throws Exc
   public String ssrfQuery(@RequestParam("domain") final String domain) {
     try {
       new URL("http://" + domain).openStream().close();
-    } catch (Throwable e) {
+    } catch (final BlockingException e) {
+      throw e;
+    } catch (final Throwable e) {
       // ignore errors opening connection
     }
     return "EXECUTED";
@@ -105,7 +108,9 @@ public String apacheHttpClient4(@RequestParam("domain") final String domain) {
     try {
       final HttpGet request = new HttpGet("http://" + domain);
       client.execute(request);
-    } catch (Exception e) {
+    } catch (final BlockingException e) {
+      throw e;
+    } catch (final Exception e) {
       // ignore errors opening connection
     }
     client.getConnectionManager().shutdown();
@@ -118,6 +123,8 @@ public String commonsHttpClient2(@RequestParam("domain") final String domain) {
     final HttpMethod method = new GetMethod("http://" + domain);
     try {
       client.executeMethod(method);
+    } catch (final BlockingException e) {
+      throw e;
     } catch (final Exception e) {
       // ignore errors opening connection
     }
@@ -131,6 +138,8 @@ public String okHttp2(@RequestParam(value = "domain") final String domain) {
     final Request request = new Request.Builder().url("http://" + domain).build();
     try {
       client.newCall(request).execute();
+    } catch (final BlockingException e) {
+      throw e;
     } catch (final Exception e) {
       // ignore errors opening connection
     }
@@ -145,6 +154,8 @@ public String okHttp3(@RequestParam("domain") final String domain) {
     final okhttp3.Request request = new okhttp3.Request.Builder().url("http://" + domain).build();
     try {
       client.newCall(request).execute();
+    } catch (final BlockingException e) {
+      throw e;
     } catch (final Exception e) {
       // ignore errors opening connection
     }

dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/SpringBootSmokeTest.groovy

@@ -1,5 +1,6 @@
 package datadog.smoketest.appsec
 
+import datadog.appsec.api.blocking.BlockingException
 import datadog.trace.agent.test.utils.OkHttpUtils
 import datadog.trace.agent.test.utils.ThreadUtils
 import groovy.json.JsonSlurper
@@ -650,6 +651,7 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest {
     def rootSpans = this.rootSpans.toList()
     rootSpans.size() == 1
     def rootSpan = rootSpans[0]
+    assert rootSpan.meta.get('error.message').contains(BlockingException.name) // ensure the block was propagated
     assert rootSpan.meta.get('appsec.blocked') == 'true', 'appsec.blocked is not set'
     assert rootSpan.meta.get('_dd.appsec.json') != null, '_dd.appsec.json is not set'
     def trigger = null