[Entitlements] Add an option to perform bytecode verification during instrumentation (#124404)

Using ASM CheckClassAdapter was key to diagnose the issue we had with incorrect signatures for some check methods.
In this PR I polished up the code I used to pinpoint the issue, and made it available via a system property so it can be turned on if we need it (and it's always on for Entitlements IT tests too).

It is also turned on in case we get VerifyErrors during retransformClasses early in the Entitlement agent bootstrap: retransformClasses runs in the native part of the JVM, so the VerifyError it produces is not so readable (e.g. it lacks a full stack trace and a description); in case this happens, we re-apply the transformation with verification turned on to get a meaningful error before dying.

libs/entitlement/asm-provider/build.gradle

@@ -14,9 +14,15 @@ dependencies {
   compileOnly project(':libs:core')
   compileOnly project(':libs:logging')
   implementation 'org.ow2.asm:asm:9.7.1'
+  implementation 'org.ow2.asm:asm-util:9.7.1'
+  implementation 'org.ow2.asm:asm-tree:9.7.1'
+  implementation 'org.ow2.asm:asm-analysis:9.7.1'
   testImplementation project(":test:framework")
   testImplementation project(":libs:entitlement:bridge")
-  testImplementation 'org.ow2.asm:asm-util:9.7.1'
+}
+
+tasks.named("dependencyLicenses").configure {
+  mapping from: /asm-.*/, to: 'asm'
 }
 
 tasks.named('test').configure {

libs/entitlement/asm-provider/src/main/java/module-info.java

@@ -12,6 +12,7 @@ import org.elasticsearch.entitlement.instrumentation.impl.InstrumentationService
 
 module org.elasticsearch.entitlement.instrumentation {
     requires org.objectweb.asm;
+    requires org.objectweb.asm.util;
     requires org.elasticsearch.entitlement;
 
     requires static org.elasticsearch.base; // for SuppressForbidden

libs/entitlement/asm-provider/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumenterImpl.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.entitlement.instrumentation.impl;
 
+import org.elasticsearch.core.Strings;
 import org.elasticsearch.entitlement.instrumentation.CheckMethod;
 import org.elasticsearch.entitlement.instrumentation.EntitlementInstrumented;
 import org.elasticsearch.entitlement.instrumentation.Instrumenter;
@@ -24,9 +25,12 @@ import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.RecordComponentVisitor;
 import org.objectweb.asm.Type;
+import org.objectweb.asm.util.CheckClassAdapter;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.util.Map;
 import java.util.stream.Stream;
 
@@ -63,6 +67,7 @@ public class InstrumenterImpl implements Instrumenter {
     }
 
     public static InstrumenterImpl create(Class<?> checkerClass, Map<MethodKey, CheckMethod> checkMethods) {
+
         Type checkerClassType = Type.getType(checkerClass);
         String handleClass = checkerClassType.getInternalName() + "Handle";
         String getCheckerClassMethodDescriptor = Type.getMethodDescriptor(checkerClassType);
@@ -82,13 +87,54 @@ public class InstrumenterImpl implements Instrumenter {
         return new ClassFileInfo(fileName, originalBytecodes);
     }
 
+    private enum VerificationPhase {
+        BEFORE_INSTRUMENTATION,
+        AFTER_INSTRUMENTATION
+    }
+
+    private static String verify(byte[] classfileBuffer) {
+        ClassReader reader = new ClassReader(classfileBuffer);
+        StringWriter stringWriter = new StringWriter();
+        PrintWriter printWriter = new PrintWriter(stringWriter);
+        CheckClassAdapter.verify(reader, false, printWriter);
+        return stringWriter.toString();
+    }
+
+    private static void verifyAndLog(byte[] classfileBuffer, String className, VerificationPhase phase) {
+        try {
+            String result = verify(classfileBuffer);
+            if (result.isEmpty() == false) {
+                logger.error(Strings.format("Bytecode verification (%s) for class [%s] failed: %s", phase, className, result));
+            } else {
+                logger.info("Bytecode verification ({}) for class [{}] passed", phase, className);
+            }
+        } catch (ClassCircularityError e) {
+            // Apparently, verification during instrumentation is challenging for class resolution and loading
+            // Treat this not as an error, but as "inconclusive"
+            logger.warn(Strings.format("Cannot perform bytecode verification (%s) for class [%s]", phase, className), e);
+        } catch (IllegalArgumentException e) {
+            // The ASM CheckClassAdapter in some cases throws this instead of printing the error
+            logger.error(Strings.format("Bytecode verification (%s) for class [%s] failed", phase, className), e);
+        }
+    }
+
     @Override
-    public byte[] instrumentClass(String className, byte[] classfileBuffer) {
+    public byte[] instrumentClass(String className, byte[] classfileBuffer, boolean verify) {
+        if (verify) {
+            verifyAndLog(classfileBuffer, className, VerificationPhase.BEFORE_INSTRUMENTATION);
+        }
+
         ClassReader reader = new ClassReader(classfileBuffer);
         ClassWriter writer = new ClassWriter(reader, COMPUTE_FRAMES | COMPUTE_MAXS);
         ClassVisitor visitor = new EntitlementClassVisitor(Opcodes.ASM9, writer, className);
         reader.accept(visitor, 0);
-        return writer.toByteArray();
+        var outBytes = writer.toByteArray();
+
+        if (verify) {
+            verifyAndLog(outBytes, className, VerificationPhase.AFTER_INSTRUMENTATION);
+        }
+
+        return outBytes;
     }
 
     class EntitlementClassVisitor extends ClassVisitor {

libs/entitlement/asm-provider/src/test/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumenterTests.java

@@ -227,9 +227,9 @@ public class InstrumenterTests extends ESTestCase {
         var instrumenter = createInstrumenter(Map.of("checkSomeStaticMethod", targetMethod));
 
         var loader1 = instrumentTestClass(instrumenter);
-        byte[] instrumentedTwiceBytecode = instrumenter.instrumentClass(TestClassToInstrument.class.getName(), loader1.testClassBytes);
-        logger.trace(() -> Strings.format("Bytecode after 2nd instrumentation:\n%s", bytecode2text(instrumentedTwiceBytecode)));
-        var loader2 = new TestLoader(TestClassToInstrument.class.getName(), instrumentedTwiceBytecode);
+        byte[] instrumentedTwiceBytes = instrumenter.instrumentClass(TestClassToInstrument.class.getName(), loader1.testClassBytes, true);
+        logger.trace(() -> Strings.format("Bytecode after 2nd instrumentation:\n%s", bytecode2text(instrumentedTwiceBytes)));
+        var loader2 = new TestLoader(TestClassToInstrument.class.getName(), instrumentedTwiceBytes);
 
         assertStaticMethodThrows(loader2, targetMethod, 123);
         assertEquals(1, TestEntitlementCheckerHolder.checkerInstance.checkSomeStaticMethodIntCallCount);
@@ -307,7 +307,7 @@ public class InstrumenterTests extends ESTestCase {
     private static TestLoader instrumentTestClass(InstrumenterImpl instrumenter) throws IOException {
         var clazz = TestClassToInstrument.class;
         ClassFileInfo initial = getClassFileInfo(clazz);
-        byte[] newBytecode = instrumenter.instrumentClass(Type.getInternalName(clazz), initial.bytecodes());
+        byte[] newBytecode = instrumenter.instrumentClass(Type.getInternalName(clazz), initial.bytecodes(), true);
         if (logger.isTraceEnabled()) {
             logger.trace("Bytecode after instrumentation:\n{}", bytecode2text(newBytecode));
         }

libs/entitlement/qa/src/javaRestTest/java/org/elasticsearch/entitlement/qa/EntitlementsTestRule.java

@@ -84,6 +84,7 @@ class EntitlementsTestRule implements TestRule {
             .module("entitled", spec -> buildEntitlements(spec, "org.elasticsearch.entitlement.qa.entitled", ENTITLED_POLICY))
             .module(ENTITLEMENT_TEST_PLUGIN_NAME, spec -> setupEntitlements(spec, modular, policyBuilder))
             .systemProperty("es.entitlements.enabled", "true")
+            .systemProperty("es.entitlements.verify_bytecode", "true")
             .systemProperty("es.entitlements.testdir", () -> testDir.getRoot().getAbsolutePath())
             .systemProperties(spec -> tempDirSystemPropertyProvider.get(testDir.getRoot().toPath()))
             .setting("xpack.security.enabled", "false")

libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java

@@ -102,6 +102,11 @@ public class EntitlementInitialization {
         manager = initChecker();
 
         var latestCheckerInterface = getVersionSpecificCheckerClass(EntitlementChecker.class);
+        var verifyBytecode = Booleans.parseBoolean(System.getProperty("es.entitlements.verify_bytecode", "false"));
+
+        if (verifyBytecode) {
+            ensureClassesSensitiveToVerificationAreInitialized();
+        }
 
         Map<MethodKey, CheckMethod> checkMethods = new HashMap<>(INSTRUMENTATION_SERVICE.lookupMethods(latestCheckerInterface));
         Stream.of(
@@ -124,8 +129,23 @@ public class EntitlementInitialization {
         var classesToTransform = checkMethods.keySet().stream().map(MethodKey::className).collect(Collectors.toSet());
 
         Instrumenter instrumenter = INSTRUMENTATION_SERVICE.newInstrumenter(latestCheckerInterface, checkMethods);
-        inst.addTransformer(new Transformer(instrumenter, classesToTransform), true);
-        inst.retransformClasses(findClassesToRetransform(inst.getAllLoadedClasses(), classesToTransform));
+        var transformer = new Transformer(instrumenter, classesToTransform, verifyBytecode);
+        inst.addTransformer(transformer, true);
+
+        var classesToRetransform = findClassesToRetransform(inst.getAllLoadedClasses(), classesToTransform);
+        try {
+            inst.retransformClasses(classesToRetransform);
+        } catch (VerifyError e) {
+            // Turn on verification and try to retransform one class at the time to get detailed diagnostic
+            transformer.enableClassVerification();
+
+            for (var classToRetransform : classesToRetransform) {
+                inst.retransformClasses(classToRetransform);
+            }
+
+            // We should have failed already in the loop above, but just in case we did not, rethrow.
+            throw e;
+        }
     }
 
     private static Class<?>[] findClassesToRetransform(Class<?>[] loadedClasses, Set<String> classesToTransform) {
@@ -431,6 +451,24 @@ public class EntitlementInitialization {
         });
     }
 
+    /**
+     * If bytecode verification is enabled, ensure these classes get loaded before transforming/retransforming them.
+     * For these classes, the order in which we transform and verify them matters. Verification during class transformation is at least an
+     * unforeseen (if not unsupported) scenario: we are loading a class, and while we are still loading it (during transformation) we try
+     * to verify it. This in turn leads to more classes loading (for verification purposes), which could turn into those classes to be
+     * transformed and undergo verification. In order to avoid circularity errors as much as possible, we force a partial order.
+     */
+    private static void ensureClassesSensitiveToVerificationAreInitialized() {
+        var classesToInitialize = Set.of("sun.net.www.protocol.http.HttpURLConnection");
+        for (String className : classesToInitialize) {
+            try {
+                Class.forName(className);
+            } catch (ClassNotFoundException unexpected) {
+                throw new AssertionError(unexpected);
+            }
+        }
+    }
+
     /**
      * Returns the "most recent" checker class compatible with the current runtime Java version.
      * For checkers, we have (optionally) version specific classes, each with a prefix (e.g. Java23).

libs/entitlement/src/main/java/org/elasticsearch/entitlement/instrumentation/Instrumenter.java

@@ -10,5 +10,5 @@
 package org.elasticsearch.entitlement.instrumentation;
 
 public interface Instrumenter {
-    byte[] instrumentClass(String className, byte[] classfileBuffer);
+    byte[] instrumentClass(String className, byte[] classfileBuffer, boolean verify);
 }

libs/entitlement/src/main/java/org/elasticsearch/entitlement/instrumentation/Transformer.java

@@ -20,12 +20,19 @@ public class Transformer implements ClassFileTransformer {
     private final Instrumenter instrumenter;
     private final Set<String> classesToTransform;
 
-    public Transformer(Instrumenter instrumenter, Set<String> classesToTransform) {
+    private boolean verifyClasses;
+
+    public Transformer(Instrumenter instrumenter, Set<String> classesToTransform, boolean verifyClasses) {
         this.instrumenter = instrumenter;
         this.classesToTransform = classesToTransform;
+        this.verifyClasses = verifyClasses;
         // TODO: Should warn if any MethodKey doesn't match any methods
     }
 
+    public void enableClassVerification() {
+        this.verifyClasses = true;
+    }
+
     @Override
     public byte[] transform(
         ClassLoader loader,
@@ -36,7 +43,7 @@ public class Transformer implements ClassFileTransformer {
     ) {
         if (classesToTransform.contains(className)) {
             // System.out.println("Transforming " + className);
-            return instrumenter.instrumentClass(className, classfileBuffer);
+            return instrumenter.instrumentClass(className, classfileBuffer, verifyClasses);
         } else {
             // System.out.println("Not transforming " + className);
             return classfileBuffer;