[Entitlements] Add test entitlement bootstrap and initialization classes

libs/entitlement/agent/src/main/java/org/elasticsearch/entitlement/agent/EntitlementAgent.java

@@ -25,10 +25,16 @@
  */
 public class EntitlementAgent {
 
+    /**
+     * The agent main method
+     * @param agentArgs arguments passed to the agent.For our agent, this is the class to load and use for Entitlement Initialization.
+     *                  See e.g. {@code EntitlementsBootstrap#loadAgent}
+     * @param inst      The {@link Instrumentation} instance to use for injecting Entitlements checks
+     */
     public static void agentmain(String agentArgs, Instrumentation inst) {
         final Class<?> initClazz;
         try {
-            initClazz = Class.forName("org.elasticsearch.entitlement.initialization.EntitlementInitialization");
+            initClazz = Class.forName(agentArgs);
         } catch (ClassNotFoundException e) {
             throw new AssertionError("entitlement agent does could not find EntitlementInitialization", e);
         }

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

@@ -138,7 +138,7 @@ private static void loadAgent(String agentPath) {
         try {
             VirtualMachine vm = VirtualMachine.attach(Long.toString(ProcessHandle.current().pid()));
             try {
-                vm.loadAgent(agentPath);
+                vm.loadAgent(agentPath, EntitlementInitialization.class.getName());
             } finally {
                 vm.detach();
             }
@@ -154,7 +154,7 @@ private static void exportInitializationToAgent() {
         EntitlementInitialization.class.getModule().addExports(initPkg, unnamedModule);
     }
 
-    private static String findAgentJar() {
+    public static String findAgentJar() {
         String propertyName = "es.entitlement.agentJar";
         String propertyValue = System.getProperty(propertyName);
         if (propertyValue != null) {

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

@@ -44,7 +44,7 @@
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
-class DynamicInstrumentation {
+public class DynamicInstrumentation {
 
     interface InstrumentationInfoFactory {
         InstrumentationService.InstrumentationInfo of(String methodName, Class<?>... parameterTypes) throws ClassNotFoundException,
@@ -93,7 +93,7 @@ InstrumentationService.InstrumentationInfo of(String methodName, Class<?>... par
      * @param checkerInterface the interface to use to find methods to instrument and to use in the injected instrumentation code
      * @param verifyBytecode   whether we should perform bytecode verification before and after instrumenting each method
      */
-    static void initialize(Instrumentation inst, Class<?> checkerInterface, boolean verifyBytecode) throws ClassNotFoundException,
+    public static void initialize(Instrumentation inst, Class<?> checkerInterface, boolean verifyBytecode) throws ClassNotFoundException,
         NoSuchMethodException, UnmodifiableClassException {
 
         var checkMethods = getMethodsToInstrument(checkerInterface);

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

@@ -9,15 +9,15 @@
 
 package org.elasticsearch.entitlement.initialization;
 
-class EntitlementCheckerUtils {
+public class EntitlementCheckerUtils {
 
     /**
      * Returns the "most recent" checker class compatible with the provided runtime Java version.
      * For checkers, we have (optionally) version specific classes, each with a prefix (e.g. Java23).
      * The mapping cannot be automatic, as it depends on the actual presence of these classes in the final Jar (see
      * the various mainXX source sets).
      */
-    static Class<?> getVersionSpecificCheckerClass(Class<?> baseClass, int javaVersion) {
+    public static Class<?> getVersionSpecificCheckerClass(Class<?> baseClass, int javaVersion) {
         String packageName = baseClass.getPackageName();
         String baseClassName = baseClass.getSimpleName();
 

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

@@ -27,9 +27,9 @@
 import static org.elasticsearch.entitlement.runtime.policy.entitlements.FilesEntitlement.Mode.READ;
 import static org.elasticsearch.entitlement.runtime.policy.entitlements.FilesEntitlement.Mode.READ_WRITE;
 
-class FilesEntitlementsValidation {
+public class FilesEntitlementsValidation {
 
-    static void validate(Map<String, Policy> pluginPolicies, PathLookup pathLookup) {
+    public static void validate(Map<String, Policy> pluginPolicies, PathLookup pathLookup) {
         Set<Path> readAccessForbidden = new HashSet<>();
         pathLookup.getBaseDirPaths(PLUGINS).forEach(p -> readAccessForbidden.add(p.toAbsolutePath().normalize()));
         pathLookup.getBaseDirPaths(MODULES).forEach(p -> readAccessForbidden.add(p.toAbsolutePath().normalize()));

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

@@ -42,7 +42,7 @@
 import static org.elasticsearch.entitlement.runtime.policy.entitlements.FilesEntitlement.Mode.READ;
 import static org.elasticsearch.entitlement.runtime.policy.entitlements.FilesEntitlement.Mode.READ_WRITE;
 
-class HardcodedEntitlements {
+public class HardcodedEntitlements {
 
     private static List<Scope> createServerEntitlements(Path pidFile) {
 
@@ -179,7 +179,7 @@ private static List<Scope> createServerEntitlements(Path pidFile) {
         return serverScopes;
     }
 
-    static Policy serverPolicy(Path pidFile, Policy serverPolicyPatch) {
+    public static Policy serverPolicy(Path pidFile, Policy serverPolicyPatch) {
         var serverScopes = createServerEntitlements(pidFile);
         return new Policy(
             "server",
@@ -190,7 +190,7 @@ static Policy serverPolicy(Path pidFile, Policy serverPolicyPatch) {
     // agents run without a module, so this is a special hack for the apm agent
     // this should be removed once https://github.com/elastic/elasticsearch/issues/109335 is completed
     // See also modules/apm/src/main/plugin-metadata/entitlement-policy.yaml
-    static List<Entitlement> agentEntitlements() {
+    public static List<Entitlement> agentEntitlements() {
         return List.of(
             new CreateClassLoaderEntitlement(),
             new ManageThreadsEntitlement(),

test/framework/build.gradle

@@ -16,6 +16,7 @@ dependencies {
   api project(':libs:ssl-config')
   api project(":server")
   api project(":libs:cli")
+  api project(":libs:entitlement:bridge")
   api "com.carrotsearch.randomizedtesting:randomizedtesting-runner:${versions.randomizedrunner}"
   api "junit:junit:${versions.junit}"
   api "org.hamcrest:hamcrest:${versions.hamcrest}"

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.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.bootstrap.EntitlementBootstrap;
+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);
+
+    /**
+     * Activates entitlement checking in tests.
+     */
+    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);
+        }
+    }
+}

test/framework/src/main/java/org/elasticsearch/bootstrap/TestEntitlementInitialization.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.bootstrap;
+
+import org.elasticsearch.entitlement.bridge.EntitlementChecker;
+import org.elasticsearch.entitlement.initialization.DynamicInstrumentation;
+import org.elasticsearch.entitlement.initialization.EntitlementCheckerUtils;
+import org.elasticsearch.entitlement.initialization.FilesEntitlementsValidation;
+import org.elasticsearch.entitlement.initialization.HardcodedEntitlements;
+import org.elasticsearch.entitlement.runtime.api.ElasticsearchEntitlementChecker;
+import org.elasticsearch.entitlement.runtime.policy.PathLookup;
+import org.elasticsearch.entitlement.runtime.policy.Policy;
+import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
+
+import java.lang.instrument.Instrumentation;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Test-only version of {@code EntitlementInitialization}
+ */
+public class TestEntitlementInitialization {
+
+    private static final Module ENTITLEMENTS_MODULE = PolicyManager.class.getModule();
+
+    private static ElasticsearchEntitlementChecker manager;
+
+    // Note: referenced by bridge reflectively
+    public static EntitlementChecker checker() {
+        return manager;
+    }
+
+    public static void initialize(Instrumentation inst) throws Exception {
+        manager = initChecker();
+        DynamicInstrumentation.initialize(
+            inst,
+            EntitlementCheckerUtils.getVersionSpecificCheckerClass(EntitlementChecker.class, Runtime.version().feature()),
+            false
+        );
+    }
+
+    private static PolicyManager createPolicyManager() {
+
+        // TODO: parse policies. Locate them using help from TestBuildInfo
+        Map<String, Policy> pluginPolicies = Map.of();
+
+        // TODO: create here the test pathLookup
+        PathLookup pathLookup = null;
+
+        FilesEntitlementsValidation.validate(pluginPolicies, pathLookup);
+
+        return new PolicyManager(
+            HardcodedEntitlements.serverPolicy(null, null),
+            HardcodedEntitlements.agentEntitlements(),
+            pluginPolicies,
+            null, // TestScopeResolver.createScopeResolver
+            Map.of(), // TODO: a map that always return nulls? Replace with functor
+            ENTITLEMENTS_MODULE, // TODO: this will need to change -- encapsulate it when we extract isTriviallyAllowed
+            pathLookup,
+            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);
+        }
+    }
+}