Propagate AppSec blocking exceptions from bytebuddy supressions (#7516)

Update our exception handler to propagate blocking exceptions

Author: manuel-alvarez-alvarez

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/blocking/BlockingExceptionHandler.java

@@ -0,0 +1,13 @@
+package datadog.trace.bootstrap.blocking;
+
+import datadog.appsec.api.blocking.BlockingException;
+
+public class BlockingExceptionHandler {
+
+  public static Throwable rethrowIfBlockingException(final Throwable e) {
+    if (e instanceof BlockingException) {
+      throw (BlockingException) e;
+    }
+    return e;
+  }
+}

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/blocking/BlockingExceptionHandlerTest.groovy

@@ -0,0 +1,23 @@
+package datadog.trace.bootstrap.blocking
+
+import datadog.appsec.api.blocking.BlockingException
+import spock.lang.Specification
+
+class BlockingExceptionHandlerTest extends Specification {
+
+  void 'test blocking exception handler'() {
+
+    when:
+    BlockingExceptionHandler.rethrowIfBlockingException(new BlockingException())
+
+    then:
+    thrown(BlockingException)
+
+    when:
+    def result = BlockingExceptionHandler.rethrowIfBlockingException(new UnsupportedOperationException())
+
+    then:
+    noExceptionThrown()
+    result instanceof UnsupportedOperationException
+  }
+}

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/ExceptionHandlers.java

@@ -1,6 +1,7 @@
 package datadog.trace.agent.tooling.bytebuddy;
 
 import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.api.ProductActivation;
 import datadog.trace.bootstrap.ExceptionLogger;
 import net.bytebuddy.ClassFileVersion;
 import net.bytebuddy.asm.Advice.ExceptionHandler;
@@ -24,6 +25,8 @@ public class ExceptionHandlers {
           new StackManipulation() {
             // Pops one Throwable off the stack. Maxes the stack to at least 3.
             private final Size size = new StackManipulation.Size(-1, 3);
+            private final boolean appSecEnabled =
+                InstrumenterConfig.get().getAppSecActivation() != ProductActivation.FULLY_DISABLED;
 
             @Override
             public boolean isValid() {
@@ -38,15 +41,19 @@ public Size apply(final MethodVisitor mv, final Implementation.Context context)
 
               // Writes the following bytecode if exitOnFailure is false:
               //
+              // BlockingExceptionHandler.rethrowIfBlockingException(t);
               // try {
+              //   InstrumentationErrors.incrementErrorCount();
               //   org.slf4j.LoggerFactory.getLogger((Class)ExceptionLogger.class)
               //     .debug("Failed to handle exception in instrumentation for ...", t);
               // } catch (Throwable t2) {
               // }
               //
               // And the following bytecode if exitOnFailure is true:
               //
+              // BlockingExceptionHandler.rethrowIfBlockingException(t);
               // try {
+              //   InstrumentationErrors.incrementErrorCount();
               //   org.slf4j.LoggerFactory.getLogger((Class)ExceptionLogger.class)
               //     .error("Failed to handle exception in instrumentation for ...", t);
               //   System.exit(1);
@@ -62,6 +69,16 @@ public Size apply(final MethodVisitor mv, final Implementation.Context context)
               final boolean frames =
                   context.getClassFileVersion().isAtLeast(ClassFileVersion.JAVA_V6);
 
+              if (appSecEnabled) {
+                // rethrow blocking exceptions
+                // stack: (top) throwable
+                mv.visitMethodInsn(
+                    Opcodes.INVOKESTATIC,
+                    "datadog/trace/bootstrap/blocking/BlockingExceptionHandler",
+                    "rethrowIfBlockingException",
+                    "(Ljava/lang/Throwable;)Ljava/lang/Throwable;");
+              }
+
               mv.visitTryCatchBlock(logStart, logEnd, eatException, "java/lang/Throwable");
               mv.visitLabel(logStart);
               // invoke incrementAndGet on our exception counter

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/AppSecDisabledExceptionHandlerForkedTest.groovy

@@ -0,0 +1,30 @@
+package datadog.trace.agent.test
+
+import ch.qos.logback.classic.Level
+import datadog.trace.api.Platform
+import spock.lang.IgnoreIf
+
+@IgnoreIf(reason = "SecurityManager used in the test is marked for removal and throws exceptions", value = {
+  Platform.isJavaVersionAtLeast(21)
+})
+class AppSecDisabledExceptionHandlerForkedTest extends BaseExceptionHandlerTest {
+
+  @Override
+  protected void changeConfig() {
+    injectSysConfig('appsec.enabled', 'false')
+  }
+
+  @Override
+  protected int expectedFailureExitStatus() {
+    return 0
+  }
+
+  @Override
+  protected Level expectedFailureLogLevel() {
+    return Level.DEBUG
+  }
+
+  protected boolean expectedBlockingException() {
+    return false
+  }
+}

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/BaseExceptionHandlerTest.groovy

@@ -3,10 +3,12 @@ package datadog.trace.agent.test
 import ch.qos.logback.classic.Level
 import ch.qos.logback.classic.Logger
 import ch.qos.logback.core.read.ListAppender
+import datadog.appsec.api.blocking.BlockingException
 import datadog.trace.agent.tooling.bytebuddy.ExceptionHandlers
 import datadog.trace.api.Platform
 import datadog.trace.bootstrap.ExceptionLogger
 import datadog.trace.bootstrap.InstrumentationErrors
+import datadog.trace.bootstrap.blocking.BlockingExceptionHandler
 import datadog.trace.test.util.DDSpecification
 import net.bytebuddy.agent.ByteBuddyAgent
 import net.bytebuddy.agent.builder.AgentBuilder
@@ -60,6 +62,13 @@ abstract class BaseExceptionHandlerTest extends DDSpecification {
       .advice(
       isMethod().and(namedOneOf("smallStack", "largeStack")),
       BadAdvice.NoOpAdvice.getName()))
+      .transform(
+      new AgentBuilder.Transformer.ForAdvice()
+      .with(new AgentBuilder.LocationStrategy.Simple(ClassFileLocator.ForClassLoader.of(BadAdvice.getClassLoader())))
+      .withExceptionHandler(ExceptionHandlers.defaultExceptionHandler())
+      .advice(
+      isMethod().and(named("blockingException")),
+      BlockingExceptionAdvice.getName()))
 
     ByteBuddyAgent.install()
     transformer = builder.installOn(ByteBuddyAgent.getInstrumentation())
@@ -94,6 +103,10 @@ abstract class BaseExceptionHandlerTest extends DDSpecification {
 
   abstract protected Level expectedFailureLogLevel()
 
+  protected boolean expectedBlockingException() {
+    true
+  }
+
   def "exception handler invoked"() {
     setup:
     int initLogEvents = testAppender.list.size()
@@ -113,9 +126,21 @@ abstract class BaseExceptionHandlerTest extends DDSpecification {
     int initLogEvents = testAppender.list.size()
     URL[] classpath = [
       SomeClass.getProtectionDomain().getCodeSource().getLocation(),
-      GroovyObject.getProtectionDomain().getCodeSource().getLocation()
+      GroovyObject.getProtectionDomain().getCodeSource().getLocation(),
     ]
-    URLClassLoader loader = new URLClassLoader(classpath, (ClassLoader) null)
+    URLClassLoader loader = new URLClassLoader(classpath, null, null) {
+        @Override
+        Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
+          if (name == BlockingExceptionHandler.name) {
+            return BlockingExceptionHandler
+          }
+          if (name == BlockingException.name) {
+            return BlockingException
+          }
+          return super.loadClass(name, resolve)
+        }
+      }
+
     when:
     loader.loadClass(InstrumentationErrors.getName())
     then:
@@ -140,6 +165,23 @@ abstract class BaseExceptionHandlerTest extends DDSpecification {
     exitStatus.get() == 0
   }
 
+  void 'blocking exception is rethrown'() {
+    when:
+    Throwable exception = null
+    try {
+      SomeClass.blockingException()
+    } catch (Throwable t) {
+      exception = t
+    }
+
+    then:
+    if (expectedBlockingException()) {
+      exception instanceof BlockingException
+    } else {
+      exception == null
+    }
+  }
+
   static class SomeClass {
     static boolean isInstrumented() {
       return false
@@ -157,6 +199,10 @@ abstract class BaseExceptionHandlerTest extends DDSpecification {
       Object o = new Object()
       println "large stack: $l $i $d $o"
     }
+
+    static void blockingException() {
+      // do nothing and throw from the advice
+    }
   }
 
   private static class NoExitSecurityManager extends SecurityManager {

dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/BlockingExceptionAdvice.java

@@ -0,0 +1,12 @@
+package datadog.trace.agent.test;
+
+import datadog.appsec.api.blocking.BlockingException;
+import net.bytebuddy.asm.Advice;
+
+public class BlockingExceptionAdvice {
+
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  public static void throwAnException() {
+    throw new BlockingException("You are blocked");
+  }
+}

dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java

@@ -94,6 +94,7 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
               + "datadog.trace.bootstrap.InstrumentationClassLoader:build_time,"
               + "datadog.trace.bootstrap.FieldBackedContextStores:build_time,"
               + "datadog.trace.bootstrap.benchmark.StaticEventLogger:build_time,"
+              + "datadog.trace.bootstrap.blocking.BlockingExceptionHandler:build_time,"
               + "datadog.trace.bootstrap.InstrumentationErrors:build_time,"
               + "datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState:build_time,"
               + "datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter:build_time,"

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/InstrumentationLogger.java

@@ -56,11 +56,7 @@ public Map<String, String> contextStore() {
 
   @Override
   public String[] helperClassNames() {
-    return new String[] {
-      packageName + ".JDBCDecorator",
-      packageName + ".SQLCommenter",
-      packageName + ".InstrumentationLogger",
-    };
+    return new String[] {packageName + ".JDBCDecorator", packageName + ".SQLCommenter"};
   }
 
   // prepend mode will prepend the SQL comment to the raw sql query
@@ -74,7 +70,7 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   public static class StatementAdvice {
-    @Advice.OnMethodEnter()
+    @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope onEnter(
         @Advice.Argument(value = 0, readOnly = false) String sql,
         @Advice.This final Statement statement) {
@@ -135,11 +131,6 @@ public static AgentScope onEnter(
         CallDepthThreadLocalMap.reset(Statement.class);
         // re-throw blocking exceptions
         throw e;
-      } catch (Throwable e) {
-        // suppress anything else
-        InstrumentationLogger.debug(
-            "datadog.trace.instrumentation.jdbc.StatementInstrumentation", statement.getClass(), e);
-        return null;
       }
     }