Fail startup if entitlement instrumentation failed

libs/entitlement/src/main/java/org/elasticsearch/entitlement/bootstrap/EntitlementBootstrap.java

@@ -98,6 +98,10 @@ public static void bootstrap(
         );
         exportInitializationToAgent();
         loadAgent(findAgentJar(), EntitlementInitialization.class.getName());
+
+        if (EntitlementInitialization.getError() != null) {
+            throw EntitlementInitialization.getError();
+        }
     }
 
     private static Path getUserHome() {

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

@@ -117,6 +117,10 @@ static void initialize(Instrumentation inst, Class<?> checkerInterface, boolean
             // We should have failed already in the loop above, but just in case we did not, rethrow.
             throw e;
         }
+
+        if (transformer.hadErrors()) {
+            throw new RuntimeException("Failed to transform JDK classes for entitlements");
+        }
     }
 
     private static Map<MethodKey, CheckMethod> getMethodsToInstrument(Class<?> checkerInterface) throws ClassNotFoundException,

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

@@ -16,11 +16,14 @@
 import org.elasticsearch.entitlement.runtime.policy.PolicyChecker;
 import org.elasticsearch.entitlement.runtime.policy.PolicyCheckerImpl;
 import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
 
 import java.lang.instrument.Instrumentation;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static java.util.Objects.requireNonNull;
 
@@ -32,17 +35,26 @@
  * to begin injecting our instrumentation.
  */
 public class EntitlementInitialization {
+    private static final Logger logger = LogManager.getLogger(EntitlementInitialization.class);
 
     private static final Module ENTITLEMENTS_MODULE = PolicyManager.class.getModule();
 
     public static InitializeArgs initializeArgs;
     private static ElasticsearchEntitlementChecker checker;
+    private static AtomicReference<RuntimeException> error = new AtomicReference<>();
 
     // Note: referenced by bridge reflectively
     public static EntitlementChecker checker() {
         return checker;
     }
 
+    /**
+     * Return any exception that occurred during initialization
+     */
+    public static RuntimeException getError() {
+        return error.get();
+    }
+
     /**
      * Initializes the Entitlement system:
      * <ol>
@@ -62,10 +74,16 @@ public static EntitlementChecker checker() {
      *
      * @param inst the JVM instrumentation class instance
      */
-    public static void initialize(Instrumentation inst) throws Exception {
-        // the checker _MUST_ be set before _any_ instrumentation is done
-        checker = initChecker(initializeArgs.policyManager());
-        initInstrumentation(inst);
+    public static void initialize(Instrumentation inst) {
+        try {
+            // the checker _MUST_ be set before _any_ instrumentation is done
+            checker = initChecker(initializeArgs.policyManager());
+            initInstrumentation(inst);
+        } catch (Exception e) {
+            // exceptions thrown within the agent will be swallowed, so capture it here
+            // instead so that it can be retrieved by bootstrap
+            error.set(new RuntimeException("Failed to initialize entitlements", e));
+        }
     }
 
     /**

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

@@ -9,16 +9,22 @@
 
 package org.elasticsearch.entitlement.instrumentation;
 
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
+
 import java.lang.instrument.ClassFileTransformer;
 import java.security.ProtectionDomain;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * A {@link ClassFileTransformer} that applies an {@link Instrumenter} to the appropriate classes.
  */
 public class Transformer implements ClassFileTransformer {
+    private static final Logger logger = LogManager.getLogger(Transformer.class);
     private final Instrumenter instrumenter;
     private final Set<String> classesToTransform;
+    private final AtomicBoolean hadErrors = new AtomicBoolean(false);
 
     private boolean verifyClasses;
 
@@ -33,6 +39,10 @@ public void enableClassVerification() {
         this.verifyClasses = true;
     }
 
+    public boolean hadErrors() {
+        return hadErrors.get();
+    }
+
     @Override
     public byte[] transform(
         ClassLoader loader,
@@ -42,13 +52,19 @@ public byte[] transform(
         byte[] classfileBuffer
     ) {
         if (classesToTransform.contains(className)) {
-            // System.out.println("Transforming " + className);
-            return instrumenter.instrumentClass(className, classfileBuffer, verifyClasses);
+            logger.debug("Transforming " + className);
+            try {
+                return instrumenter.instrumentClass(className, classfileBuffer, verifyClasses);
+            } catch (Throwable t) {
+                hadErrors.set(true);
+                logger.error("Failed to instrument class " + className, t);
+                // throwing an exception from a transformer results in the exception being swallowed,
+                // effectively the same as returning null anyways, so we instead log it here completely
+                return null;
+            }
         } else {
-            // System.out.println("Not transforming " + className);
+            logger.trace("Not transforming " + className);
             return null;
         }
     }
-
-    // private static final Logger LOGGER = LogManager.getLogger(Transformer.class);
 }