[Entitlements] Add an option to perform bytecode verification during instrumentation

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 @@
 
 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.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 @@ static ClassFileInfo getClassFileInfo(Class<?> clazz) throws IOException {
         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 void testNotInstrumentedTwice() throws Exception {
         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 @@ private static InstrumenterImpl createInstrumenter(Map<String, Executable> metho
     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 @@ protected void before() throws Throwable {
             .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 static void initialize(Instrumentation inst) throws Exception {
         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 static void initialize(Instrumentation inst) throws Exception {
         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 @@ private static Stream<InstrumentationService.InstrumentationInfo> pathChecks() {
         });
     }
 
+    /**
+     * 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 byte[] transform(
     ) {
         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;