Create PolicyManager during bootstrap, allowing us to share initialization

prdoyle · prdoyle · commit d08d6d2ca066 · 2025-06-06T18:10:23.000-04:00
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/bootstrap/EntitlementBootstrap.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/bootstrap/EntitlementBootstrap.java
@@ -18,6 +18,7 @@
 import org.elasticsearch.core.PathUtils;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.entitlement.initialization.EntitlementInitialization;
+import org.elasticsearch.entitlement.runtime.policy.PathLookup;
 import org.elasticsearch.entitlement.runtime.policy.PathLookupImpl;
 import org.elasticsearch.entitlement.runtime.policy.Policy;
 import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
@@ -91,12 +92,9 @@ public static void bootstrap(
             settingResolver
         );
         EntitlementInitialization.initializeArgs = new EntitlementInitialization.InitializeArgs(
-            serverPolicyPatch,
-            pluginPolicies,
-            scopeResolver,
             pathLookup,
-            pluginSourcePaths,
-            suppressFailureLogPackages
+            suppressFailureLogPackages,
+            createPolicyManager(pluginPolicies, pathLookup, serverPolicyPatch, scopeResolver, pluginSourcePaths)
         );
         exportInitializationToAgent();
         loadAgent(findAgentJar(), EntitlementInitialization.class.getName());
@@ -154,5 +152,24 @@ static String findAgentJar() {
         }
     }
 
+    private static PolicyManager createPolicyManager(
+        Map<String, Policy> pluginPolicies,
+        PathLookup pathLookup,
+        Policy serverPolicyPatch,
+        Function<Class<?>, PolicyManager.PolicyScope> scopeResolver,
+        Map<String, Collection<Path>> pluginSourcePaths
+    ) {
+        FilesEntitlementsValidation.validate(pluginPolicies, pathLookup);
+
+        return new PolicyManager(
+            HardcodedEntitlements.serverPolicy(pathLookup.pidFile(), serverPolicyPatch),
+            HardcodedEntitlements.agentEntitlements(),
+            pluginPolicies,
+            scopeResolver,
+            pluginSourcePaths,
+            pathLookup
+        );
+    }
+
     private static final Logger logger = LogManager.getLogger(EntitlementBootstrap.class);
 }
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/bootstrap/FilesEntitlementsValidation.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/bootstrap/FilesEntitlementsValidation.java
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.entitlement.initialization;
+package org.elasticsearch.entitlement.bootstrap;
 
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.entitlement.runtime.policy.FileAccessTree;
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/bootstrap/HardcodedEntitlements.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/bootstrap/HardcodedEntitlements.java
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.entitlement.initialization;
+package org.elasticsearch.entitlement.bootstrap;
 
 import org.elasticsearch.core.Booleans;
 import org.elasticsearch.entitlement.runtime.policy.Policy;
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java
@@ -10,23 +10,17 @@
 package org.elasticsearch.entitlement.initialization;
 
 import org.elasticsearch.core.Booleans;
-import org.elasticsearch.core.Nullable;
 import org.elasticsearch.entitlement.bridge.EntitlementChecker;
 import org.elasticsearch.entitlement.runtime.policy.ElasticsearchEntitlementChecker;
 import org.elasticsearch.entitlement.runtime.policy.PathLookup;
-import org.elasticsearch.entitlement.runtime.policy.Policy;
 import org.elasticsearch.entitlement.runtime.policy.PolicyChecker;
 import org.elasticsearch.entitlement.runtime.policy.PolicyCheckerImpl;
 import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
 
 import java.lang.instrument.Instrumentation;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
-import java.nio.file.Path;
-import java.util.Collection;
-import java.util.Map;
 import java.util.Set;
-import java.util.function.Function;
 
 import static java.util.Objects.requireNonNull;
 
@@ -70,35 +64,23 @@ public static EntitlementChecker checker() {
      */
     public static void initialize(Instrumentation inst) throws Exception {
         // the checker _MUST_ be set before _any_ instrumentation is done
-        checker = initChecker(createPolicyManager());
+        checker = initChecker(initializeArgs.policyManager());
         initInstrumentation(inst);
     }
 
     /**
      * Arguments to {@link #initialize}. Since that's called in a static context from the agent,
      * we have no way to pass arguments directly, so we stuff them in here.
      *
-     * @param serverPolicyPatch
-     * @param pluginPolicies
-     * @param scopeResolver
      * @param pathLookup
-     * @param pluginSourcePaths
      * @param suppressFailureLogPackages
+     * @param policyManager
      */
-    public record InitializeArgs(
-        @Nullable Policy serverPolicyPatch,
-        Map<String, Policy> pluginPolicies,
-        Function<Class<?>, PolicyManager.PolicyScope> scopeResolver,
-        PathLookup pathLookup,
-        Map<String, Collection<Path>> pluginSourcePaths,
-        Set<Package> suppressFailureLogPackages
-    ) {
+    public record InitializeArgs(PathLookup pathLookup, Set<Package> suppressFailureLogPackages, PolicyManager policyManager) {
         public InitializeArgs {
-            requireNonNull(pluginPolicies);
-            requireNonNull(scopeResolver);
             requireNonNull(pathLookup);
-            requireNonNull(pluginSourcePaths);
             requireNonNull(suppressFailureLogPackages);
+            requireNonNull(policyManager);
         }
     }
 
@@ -111,22 +93,6 @@ private static PolicyCheckerImpl createPolicyChecker(PolicyManager policyManager
         );
     }
 
-    private static PolicyManager createPolicyManager() {
-        Map<String, Policy> pluginPolicies = initializeArgs.pluginPolicies();
-        PathLookup pathLookup = initializeArgs.pathLookup();
-
-        FilesEntitlementsValidation.validate(pluginPolicies, pathLookup);
-
-        return new PolicyManager(
-            HardcodedEntitlements.serverPolicy(pathLookup.pidFile(), initializeArgs.serverPolicyPatch()),
-            HardcodedEntitlements.agentEntitlements(),
-            pluginPolicies,
-            initializeArgs.scopeResolver(),
-            initializeArgs.pluginSourcePaths(),
-            pathLookup
-        );
-    }
-
     /**
      * 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
diff --git a/libs/entitlement/src/test/java/org/elasticsearch/entitlement/bootstrap/FilesEntitlementsValidationTests.java b/libs/entitlement/src/test/java/org/elasticsearch/entitlement/bootstrap/FilesEntitlementsValidationTests.java
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.entitlement.initialization;
+package org.elasticsearch.entitlement.bootstrap;
 
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.entitlement.runtime.policy.PathLookup;
diff --git a/test/framework/src/main/java/org/elasticsearch/entitlement/bootstrap/TestEntitlementBootstrap.java b/test/framework/src/main/java/org/elasticsearch/entitlement/bootstrap/TestEntitlementBootstrap.java
@@ -9,12 +9,30 @@
 
 package org.elasticsearch.entitlement.bootstrap;
 
-import org.elasticsearch.entitlement.initialization.TestEntitlementInitialization;
+import org.elasticsearch.bootstrap.TestBuildInfo;
+import org.elasticsearch.bootstrap.TestBuildInfoParser;
+import org.elasticsearch.bootstrap.TestScopeResolver;
+import org.elasticsearch.core.Strings;
+import org.elasticsearch.core.SuppressForbidden;
+import org.elasticsearch.entitlement.initialization.EntitlementInitialization;
 import org.elasticsearch.entitlement.runtime.policy.PathLookup;
+import org.elasticsearch.entitlement.runtime.policy.Policy;
+import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
+import org.elasticsearch.entitlement.runtime.policy.PolicyParser;
+import org.elasticsearch.entitlement.runtime.policy.TestPolicyManager;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.plugins.PluginDescriptor;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URL;
 import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.stream.Stream;
 
 public class TestEntitlementBootstrap {
@@ -24,10 +42,15 @@ public class TestEntitlementBootstrap {
     /**
      * Activates entitlement checking in tests.
      */
-    public static void bootstrap() {
-        TestEntitlementInitialization.initializeArgs = new TestEntitlementInitialization.InitializeArgs(new TestPathLookup());
+    public static void bootstrap() throws IOException {
+        TestPathLookup pathLookup = new TestPathLookup();
+        EntitlementInitialization.initializeArgs = new EntitlementInitialization.InitializeArgs(
+            pathLookup,
+            Set.of(),
+            createPolicyManager(pathLookup)
+        );
         logger.debug("Loading entitlement agent");
-        EntitlementBootstrap.loadAgent(EntitlementBootstrap.findAgentJar(), TestEntitlementInitialization.class.getName());
+        EntitlementBootstrap.loadAgent(EntitlementBootstrap.findAgentJar(), EntitlementInitialization.class.getName());
     }
 
     private record TestPathLookup() implements PathLookup {
@@ -47,4 +70,71 @@ public Stream<Path> resolveSettingPaths(BaseDir baseDir, String settingName) {
         }
 
     }
+
+    private static PolicyManager createPolicyManager(PathLookup pathLookup) throws IOException {
+
+        var pluginsTestBuildInfo = TestBuildInfoParser.parseAllPluginTestBuildInfo();
+        var serverTestBuildInfo = TestBuildInfoParser.parseServerTestBuildInfo();
+        var scopeResolver = TestScopeResolver.createScopeResolver(serverTestBuildInfo, pluginsTestBuildInfo);
+        List<String> pluginNames = pluginsTestBuildInfo.stream().map(TestBuildInfo::component).toList();
+
+        var pluginDescriptors = parsePluginsDescriptors(pluginNames);
+        var pluginsData = pluginDescriptors.stream()
+            .map(descriptor -> new TestPluginData(descriptor.getName(), descriptor.isModular(), false))
+            .toList();
+        Map<String, Policy> pluginPolicies = parsePluginsPolicies(pluginsData);
+
+        FilesEntitlementsValidation.validate(pluginPolicies, pathLookup);
+
+        return new TestPolicyManager(
+            HardcodedEntitlements.serverPolicy(null, null),
+            HardcodedEntitlements.agentEntitlements(),
+            pluginPolicies,
+            scopeResolver,
+            Map.of(),
+            pathLookup
+        );
+    }
+
+    private record TestPluginData(String pluginName, boolean isModular, boolean isExternalPlugin) {}
+
+    private static Map<String, Policy> parsePluginsPolicies(List<TestPluginData> pluginsData) {
+        Map<String, Policy> policies = new HashMap<>();
+        for (var pluginData : pluginsData) {
+            String pluginName = pluginData.pluginName();
+            var resourceName = Strings.format("META-INF/es-plugins/%s/entitlement-policy.yaml", pluginName);
+
+            var resource = EntitlementInitialization.class.getClassLoader().getResource(resourceName);
+            if (resource != null) {
+                try (var inputStream = getStream(resource)) {
+                    policies.put(pluginName, new PolicyParser(inputStream, pluginName, pluginData.isExternalPlugin()).parsePolicy());
+                } catch (IOException e) {
+                    throw new IllegalArgumentException(Strings.format("Cannot read policy for plugin [%s]", pluginName), e);
+                }
+            }
+        }
+        return policies;
+    }
+
+    private static List<PluginDescriptor> parsePluginsDescriptors(List<String> pluginNames) {
+        List<PluginDescriptor> descriptors = new ArrayList<>();
+        for (var pluginName : pluginNames) {
+            var resourceName = Strings.format("META-INF/es-plugins/%s/plugin-descriptor.properties", pluginName);
+            var resource = EntitlementInitialization.class.getClassLoader().getResource(resourceName);
+            if (resource != null) {
+                try (var inputStream = getStream(resource)) {
+                    descriptors.add(PluginDescriptor.readInternalDescriptorFromStream(inputStream));
+                } catch (IOException e) {
+                    throw new IllegalArgumentException(Strings.format("Cannot read descriptor for plugin [%s]", pluginName), e);
+                }
+            }
+        }
+        return descriptors;
+    }
+
+    @SuppressForbidden(reason = "URLs from class loader")
+    private static InputStream getStream(URL resource) throws IOException {
+        return resource.openStream();
+    }
+
 }
diff --git a/test/framework/src/main/java/org/elasticsearch/entitlement/initialization/TestEntitlementInitialization.java b/test/framework/src/main/java/org/elasticsearch/entitlement/initialization/TestEntitlementInitialization.java
