Guard rethrowIfBlockingException call against unexpected failures

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

@@ -49,24 +49,43 @@ public Size apply(final MethodVisitor mv, final Implementation.Context context)
           //
           // Emits the following Java-equivalent code when exitOnFailure is false:
           //
-          // BlockingExceptionHandler.rethrowIfBlockingException(t);
           // try {
+          //   BlockingExceptionHandler.rethrowIfBlockingException(t);
           //   InstrumentationErrors.recordError();
           //   org.slf4j.LoggerFactory.getLogger((Class) ExceptionLogger.class)
           //     .debug("Failed to handle exception in instrumentation for <type> (" + adviceName +
           // ")", t);
           // } catch (Throwable t2) {
+          //   if (t2.getClass().getName().equals("datadog.appsec.api.blocking.BlockingException"))
+          // {
+          //     throw t2;
+          //   }
           // }
           //
           // and the same with .error(...) followed by System.exit(1) when exitOnFailure is true.
+          //
+          // rethrowIfBlockingException is inside the try/catch (rather than called bare, as
+          // before) so that an unexpected failure resolving/invoking it - e.g. a
+          // NoClassDefFoundError from a classloader that can't see the appsec module - is
+          // swallowed like any other instrumentation error instead of replacing the original
+          // exception and escaping into the instrumented method's caller. The catch handler
+          // re-throws only when the caught exception really is the BlockingException the call
+          // was meant to propagate. It compares by class name rather than using `instanceof`
+          // because `instanceof` would itself need to resolve BlockingException via the
+          // instrumented class's own classloader - the same NoClassDefFoundError risk this whole
+          // try/catch exists to guard against.
           final Label logStart = new Label();
           final Label logEnd = new Label();
           final Label eatException = new Label();
+          final Label notBlocking = new Label();
           final Label handlerExit = new Label();
 
           // Frames are only meaningful for class files in version 6 or later.
           final boolean frames = context.getClassFileVersion().isAtLeast(ClassFileVersion.JAVA_V6);
 
+          mv.visitTryCatchBlock(logStart, logEnd, eatException, "java/lang/Throwable");
+          mv.visitLabel(logStart);
+
           if (appSecEnabled) {
             // Need throwable on top for rethrowIfBlockingException.
             // stack: (top) adviceName, throwable -> top throwable
@@ -81,8 +100,6 @@ public Size apply(final MethodVisitor mv, final Implementation.Context context)
             mv.visitInsn(Opcodes.SWAP);
           }
 
-          mv.visitTryCatchBlock(logStart, logEnd, eatException, "java/lang/Throwable");
-          mv.visitLabel(logStart);
           // record instrumentation error
           if (detailedErrors) {
             // recordError(Throwable) needs throwable on top, then we restore.
@@ -148,12 +165,42 @@ public Size apply(final MethodVisitor mv, final Implementation.Context context)
           mv.visitLabel(logEnd);
           mv.visitJumpInsn(Opcodes.GOTO, handlerExit);
 
-          // if the runtime can't reach our ExceptionHandler or logger,
-          //   silently eat the exception
+          // If the runtime can't reach our ExceptionHandler or logger, or
+          // rethrowIfBlockingException itself failed unexpectedly, silently eat the exception -
+          // unless it's the BlockingException rethrowIfBlockingException was meant to propagate,
+          // in which case let it through.
           mv.visitLabel(eatException);
           if (frames) {
             mv.visitFrame(Opcodes.F_SAME1, 0, null, 1, new Object[] {"java/lang/Throwable"});
           }
+          if (appSecEnabled) {
+            // Compare by class name instead of `instanceof`: `instanceof` would resolve
+            // BlockingException via the instrumented class's own classloader, which can throw
+            // NoClassDefFoundError right here - uncaught - on a classloader that can't see the
+            // appsec module. A name comparison never triggers that resolution.
+            mv.visitInsn(Opcodes.DUP);
+            mv.visitMethodInsn(
+                Opcodes.INVOKEVIRTUAL,
+                "java/lang/Object",
+                "getClass",
+                "()Ljava/lang/Class;",
+                false);
+            mv.visitMethodInsn(
+                Opcodes.INVOKEVIRTUAL, "java/lang/Class", "getName", "()Ljava/lang/String;", false);
+            mv.visitLdcInsn("datadog.appsec.api.blocking.BlockingException");
+            mv.visitMethodInsn(
+                Opcodes.INVOKEVIRTUAL,
+                "java/lang/String",
+                "equals",
+                "(Ljava/lang/Object;)Z",
+                false);
+            mv.visitJumpInsn(Opcodes.IFEQ, notBlocking);
+            mv.visitInsn(Opcodes.ATHROW);
+            mv.visitLabel(notBlocking);
+            if (frames) {
+              mv.visitFrame(Opcodes.F_SAME1, 0, null, 1, new Object[] {"java/lang/Throwable"});
+            }
+          }
           mv.visitInsn(Opcodes.POP);
           // mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Throwable",
           //    "printStackTrace", "()V", false);

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

@@ -68,6 +68,14 @@ abstract class BaseExceptionHandlerTest extends DDSpecification {
       .advice(
       isMethod().and(named("blockingException")),
       BlockingExceptionAdvice.getName()))
+      .type(named(BaseExceptionHandlerTest.getName() + '$SomeOtherClass'))
+      .transform(
+      new AgentBuilder.Transformer.ForAdvice()
+      .with(new AgentBuilder.LocationStrategy.Simple(ClassFileLocator.ForClassLoader.of(BadAdvice.getClassLoader())))
+      .withExceptionHandler(ExceptionHandlers.exceptionHandlerFor(BadAdvice.getName()))
+      .advice(
+      isMethod().and(named("isInstrumented")),
+      BadAdvice.getName()))
 
     ByteBuddyAgent.install()
     transformer = builder.installOn(ByteBuddyAgent.getInstrumentation())
@@ -143,6 +151,34 @@ abstract class BaseExceptionHandlerTest extends DDSpecification {
     exitStatus.get() == 0
   }
 
+  def "exception on classloader that cannot resolve BlockingException"() {
+    setup:
+    int initLogEvents = testAppender.list.size()
+    URL[] classpath = [
+      SomeClass.getProtectionDomain().getCodeSource().getLocation(),
+      GroovyObject.getProtectionDomain().getCodeSource().getLocation(),
+    ]
+    URLClassLoader loader = new AppSecInvisibleClassLoader(
+      classpath, BaseExceptionHandlerTest.getClassLoader(), SomeOtherClass.getName())
+
+    when:
+    loader.loadClass(BlockingException.getName())
+    then:
+    thrown ClassNotFoundException
+
+    when:
+    Class<?> someClazz = loader.loadClass(SomeOtherClass.getName())
+    then:
+    someClazz.getClassLoader() == loader
+
+    when:
+    someClazz.getMethod("isInstrumented").invoke(null)
+    then:
+    noExceptionThrown()
+    testAppender.list.size() == initLogEvents + 1
+    exitStatus.get() == expectedFailureExitStatus()
+  }
+
   def "exception handler sets the correct stack size"() {
     when:
     SomeClass.smallStack()
@@ -193,6 +229,16 @@ abstract class BaseExceptionHandlerTest extends DDSpecification {
     }
   }
 
+  // Deliberately not instrumented with BlockingExceptionAdvice, unlike SomeClass: that advice's
+  // own bytecode constructs a real BlockingException, which would make the JVM verifier resolve
+  // BlockingException while linking the whole class - defeating the point of testing a
+  // classloader that can't see it.
+  static class SomeOtherClass {
+    static boolean isInstrumented() {
+      return false
+    }
+  }
+
   private static class NoExitSecurityManager extends SecurityManager {
     private final AtomicInteger status
 

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

@@ -0,0 +1,37 @@
+package datadog.trace.agent.test;
+
+import datadog.appsec.api.blocking.BlockingException;
+import java.net.URL;
+import java.net.URLClassLoader;
+
+/**
+ * A {@link URLClassLoader} that delegates to the given parent for everything - including {@code
+ * datadog.trace.bootstrap.*} and slf4j, which in production are visible from any classloader -
+ * except {@link BlockingException} (which in production is only visible if the appsec module is
+ * present) and the given target class name, which this loader defines locally so it gets
+ * instrumented as if loaded by an isolated (e.g. plugin/OSGi-style) classloader.
+ */
+final class AppSecInvisibleClassLoader extends URLClassLoader {
+  private final String isolatedClassName;
+
+  AppSecInvisibleClassLoader(URL[] classpath, ClassLoader parent, String isolatedClassName) {
+    super(classpath, parent);
+    this.isolatedClassName = isolatedClassName;
+  }
+
+  @Override
+  protected synchronized Class<?> loadClass(String name, boolean resolve)
+      throws ClassNotFoundException {
+    Class<?> found = findLoadedClass(name);
+    if (found == null) {
+      if (name.equals(BlockingException.class.getName())) {
+        throw new ClassNotFoundException(name);
+      }
+      found = name.equals(isolatedClassName) ? findClass(name) : super.loadClass(name, resolve);
+    }
+    if (resolve) {
+      resolveClass(found);
+    }
+    return found;
+  }
+}