PR comments

Author: ldematte

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

@@ -122,7 +122,7 @@ public static void bootstrap(
             suppressFailureLogClasses
         );
         exportInitializationToAgent();
-        loadAgent(findAgentJar());
+        loadAgent(findAgentJar(), EntitlementInitialization.class.getName());
     }
 
     private static Path getUserHome() {
@@ -134,11 +134,11 @@ private static Path getUserHome() {
     }
 
     @SuppressForbidden(reason = "The VirtualMachine API is the only way to attach a java agent dynamically")
-    private static void loadAgent(String agentPath) {
+    static void loadAgent(String agentPath, String entitlementInitializationClassName) {
         try {
             VirtualMachine vm = VirtualMachine.attach(Long.toString(ProcessHandle.current().pid()));
             try {
-                vm.loadAgent(agentPath, EntitlementInitialization.class.getName());
+                vm.loadAgent(agentPath, entitlementInitializationClassName);
             } finally {
                 vm.detach();
             }

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

@@ -34,11 +34,11 @@ public class EntitlementInitialization {
 
     private static final Module ENTITLEMENTS_MODULE = PolicyManager.class.getModule();
 
-    private static ElasticsearchEntitlementChecker manager;
+    private static ElasticsearchEntitlementChecker checker;
 
     // Note: referenced by bridge reflectively
     public static EntitlementChecker checker() {
-        return manager;
+        return checker;
     }
 
     /**
@@ -61,18 +61,7 @@ public static EntitlementChecker checker() {
      * @param inst the JVM instrumentation class instance
      */
     public static void initialize(Instrumentation inst) throws Exception {
-        manager = initChecker();
-
-        var verifyBytecode = Booleans.parseBoolean(System.getProperty("es.entitlements.verify_bytecode", "false"));
-        if (verifyBytecode) {
-            ensureClassesSensitiveToVerificationAreInitialized();
-        }
-
-        DynamicInstrumentation.initialize(
-            inst,
-            EntitlementCheckerUtils.getVersionSpecificCheckerClass(EntitlementChecker.class, Runtime.version().feature()),
-            verifyBytecode
-        );
+        checker = initChecker(inst, createPolicyManager());
     }
 
     private static PolicyManager createPolicyManager() {
@@ -112,9 +101,7 @@ private static void ensureClassesSensitiveToVerificationAreInitialized() {
         }
     }
 
-    private static ElasticsearchEntitlementChecker initChecker() {
-        final PolicyManager policyManager = createPolicyManager();
-
+    static ElasticsearchEntitlementChecker initChecker(Instrumentation inst, PolicyManager policyManager) throws Exception {
         final Class<?> clazz = EntitlementCheckerUtils.getVersionSpecificCheckerClass(
             ElasticsearchEntitlementChecker.class,
             Runtime.version().feature()
@@ -126,10 +113,25 @@ private static ElasticsearchEntitlementChecker initChecker() {
         } catch (NoSuchMethodException e) {
             throw new AssertionError("entitlement impl is missing no arg constructor", e);
         }
+
+        ElasticsearchEntitlementChecker checker;
         try {
-            return (ElasticsearchEntitlementChecker) constructor.newInstance(policyManager);
+            checker = (ElasticsearchEntitlementChecker) constructor.newInstance(policyManager);
         } catch (IllegalAccessException | InvocationTargetException | InstantiationException e) {
             throw new AssertionError(e);
         }
+
+        var verifyBytecode = Booleans.parseBoolean(System.getProperty("es.entitlements.verify_bytecode", "false"));
+        if (verifyBytecode) {
+            ensureClassesSensitiveToVerificationAreInitialized();
+        }
+
+        DynamicInstrumentation.initialize(
+            inst,
+            EntitlementCheckerUtils.getVersionSpecificCheckerClass(EntitlementChecker.class, Runtime.version().feature()),
+            verifyBytecode
+        );
+
+        return checker;
     }
 }

test/framework/src/main/java/org/elasticsearch/entitlement/bootstrap/TestEntitlementBootstrap.java

@@ -9,18 +9,10 @@
 
 package org.elasticsearch.entitlement.bootstrap;
 
-import com.sun.tools.attach.AgentInitializationException;
-import com.sun.tools.attach.AgentLoadException;
-import com.sun.tools.attach.AttachNotSupportedException;
-import com.sun.tools.attach.VirtualMachine;
-
-import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.entitlement.initialization.TestEntitlementInitialization;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 
-import java.io.IOException;
-
 class TestEntitlementBootstrap {
 
     private static final Logger logger = LogManager.getLogger(TestEntitlementBootstrap.class);
@@ -30,20 +22,6 @@ class TestEntitlementBootstrap {
      */
     public static void bootstrap() {
         logger.debug("Loading entitlement agent");
-        loadAgent(EntitlementBootstrap.findAgentJar());
-    }
-
-    @SuppressForbidden(reason = "The VirtualMachine API is the only way to attach a java agent dynamically")
-    private static void loadAgent(String agentPath) {
-        try {
-            VirtualMachine vm = VirtualMachine.attach(Long.toString(ProcessHandle.current().pid()));
-            try {
-                vm.loadAgent(agentPath, TestEntitlementInitialization.class.getName());
-            } finally {
-                vm.detach();
-            }
-        } catch (AttachNotSupportedException | IOException | AgentLoadException | AgentInitializationException e) {
-            throw new IllegalStateException("Unable to attach entitlement agent", e);
-        }
+        EntitlementBootstrap.loadAgent(EntitlementBootstrap.findAgentJar(), TestEntitlementInitialization.class.getName());
     }
 }

test/framework/src/main/java/org/elasticsearch/entitlement/initialization/TestEntitlementInitialization.java

@@ -22,8 +22,6 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.lang.instrument.Instrumentation;
-import java.lang.reflect.Constructor;
-import java.lang.reflect.InvocationTargetException;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -37,20 +35,15 @@
  */
 public class TestEntitlementInitialization {
 
-    private static ElasticsearchEntitlementChecker manager;
+    private static ElasticsearchEntitlementChecker checker;
 
     // Note: referenced by bridge reflectively
     public static EntitlementChecker checker() {
-        return manager;
+        return checker;
     }
 
     public static void initialize(Instrumentation inst) throws Exception {
-        manager = initChecker();
-        DynamicInstrumentation.initialize(
-            inst,
-            EntitlementCheckerUtils.getVersionSpecificCheckerClass(EntitlementChecker.class, Runtime.version().feature()),
-            false
-        );
+        checker = EntitlementInitialization.initChecker(inst, createPolicyManager());
     }
 
     private record TestPluginData(String pluginName, boolean isModular, boolean isExternalPlugin) {}
@@ -125,25 +118,4 @@ private static PolicyManager createPolicyManager() {
             Set.of()
         );
     }
-
-    private static ElasticsearchEntitlementChecker initChecker() {
-        final PolicyManager policyManager = createPolicyManager();
-
-        final Class<?> clazz = EntitlementCheckerUtils.getVersionSpecificCheckerClass(
-            ElasticsearchEntitlementChecker.class,
-            Runtime.version().feature()
-        );
-
-        Constructor<?> constructor;
-        try {
-            constructor = clazz.getConstructor(PolicyManager.class);
-        } catch (NoSuchMethodException e) {
-            throw new AssertionError("entitlement impl is missing no arg constructor", e);
-        }
-        try {
-            return (ElasticsearchEntitlementChecker) constructor.newInstance(policyManager);
-        } catch (IllegalAccessException | InvocationTargetException | InstantiationException e) {
-            throw new AssertionError(e);
-        }
-    }
 }