Refactor: ScopeResolver (#126921)

* Fix: use getScopeName consistently

* Rename PolicyManagerTests method

* Refacor: simplify PluginsResolver.create

* Change PluginsResolver to ScopeResolver

* Move boot layer test to ScopeResolverTests

* [CI] Auto commit changes from spotless

* Rename PolicyScope

* Add ComponentKind enum

* Package private componentName field

---------

Co-authored-by: elasticsearchmachine <[email protected]>

Author: prdoyle

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

@@ -21,6 +21,7 @@
 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;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 
@@ -39,14 +40,14 @@ public class EntitlementBootstrap {
     public record BootstrapArgs(
         @Nullable Policy serverPolicyPatch,
         Map<String, Policy> pluginPolicies,
-        Function<Class<?>, String> pluginResolver,
+        Function<Class<?>, PolicyManager.PolicyScope> scopeResolver,
         PathLookup pathLookup,
         Map<String, Path> sourcePaths,
         Set<Class<?>> suppressFailureLogClasses
     ) {
         public BootstrapArgs {
             requireNonNull(pluginPolicies);
-            requireNonNull(pluginResolver);
+            requireNonNull(scopeResolver);
             requireNonNull(pathLookup);
             requireNonNull(sourcePaths);
             requireNonNull(suppressFailureLogClasses);
@@ -65,7 +66,7 @@ public static BootstrapArgs bootstrapArgs() {
      *
      * @param serverPolicyPatch a policy with additional entitlements to patch the embedded server layer policy
      * @param pluginPolicies a map holding policies for plugins (and modules), by plugin (or module) name.
-     * @param pluginResolver a functor to map a Java Class to the plugin it belongs to (the plugin name).
+     * @param scopeResolver a functor to map a Java Class to the component and module it belongs to.
      * @param settingResolver a functor to resolve a setting name pattern for one or more Elasticsearch settings.
      * @param dataDirs       data directories for Elasticsearch
      * @param sharedRepoDirs shared repository directories for Elasticsearch
@@ -82,7 +83,7 @@ public static BootstrapArgs bootstrapArgs() {
     public static void bootstrap(
         Policy serverPolicyPatch,
         Map<String, Policy> pluginPolicies,
-        Function<Class<?>, String> pluginResolver,
+        Function<Class<?>, PolicyManager.PolicyScope> scopeResolver,
         Function<String, Stream<String>> settingResolver,
         Path[] dataDirs,
         Path[] sharedRepoDirs,
@@ -103,7 +104,7 @@ public static void bootstrap(
         EntitlementBootstrap.bootstrapArgs = new BootstrapArgs(
             serverPolicyPatch,
             pluginPolicies,
-            pluginResolver,
+            scopeResolver,
             new PathLookupImpl(
                 getUserHome(),
                 configDir,

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

@@ -90,7 +90,6 @@
  */
 public class EntitlementInitialization {
 
-    private static final String AGENTS_PACKAGE_NAME = "co.elastic.apm.agent";
     private static final Module ENTITLEMENTS_MODULE = PolicyManager.class.getModule();
 
     private static ElasticsearchEntitlementChecker manager;
@@ -314,9 +313,8 @@ private static PolicyManager createPolicyManager() {
             serverPolicy,
             agentEntitlements,
             pluginPolicies,
-            EntitlementBootstrap.bootstrapArgs().pluginResolver(),
+            EntitlementBootstrap.bootstrapArgs().scopeResolver(),
             EntitlementBootstrap.bootstrapArgs().sourcePaths(),
-            AGENTS_PACKAGE_NAME,
             ENTITLEMENTS_MODULE,
             pathLookup,
             bootstrapArgs.suppressFailureLogClasses()

libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java

@@ -59,25 +59,66 @@
 import static java.util.zip.ZipFile.OPEN_READ;
 import static org.elasticsearch.entitlement.bridge.Util.NO_CLASS;
 import static org.elasticsearch.entitlement.runtime.policy.PathLookup.BaseDir.TEMP;
+import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ComponentKind.APM_AGENT;
+import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ComponentKind.PLUGIN;
+import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ComponentKind.SERVER;
+import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ComponentKind.UNKNOWN;
 
 public class PolicyManager {
     /**
      * Use this if you don't have a {@link ModuleEntitlements} in hand.
      */
     private static final Logger generalLogger = LogManager.getLogger(PolicyManager.class);
 
-    static final String UNKNOWN_COMPONENT_NAME = "(unknown)";
-    static final String SERVER_COMPONENT_NAME = "(server)";
-    static final String APM_AGENT_COMPONENT_NAME = "(APM agent)";
-
     static final Class<?> DEFAULT_FILESYSTEM_CLASS = PathUtils.getDefaultFileSystem().getClass();
 
     static final Set<String> MODULES_EXCLUDED_FROM_SYSTEM_MODULES = Set.of("java.desktop");
 
     /**
-     * @param componentName the plugin name; or else one of the special component names
-     *                      like {@link #SERVER_COMPONENT_NAME} or {@link #APM_AGENT_COMPONENT_NAME}.
+     * Identifies a particular entitlement {@link Scope} within a {@link Policy}.
      */
+    public record PolicyScope(ComponentKind kind, String componentName, String moduleName) {
+        public PolicyScope {
+            requireNonNull(kind);
+            requireNonNull(componentName);
+            requireNonNull(moduleName);
+            assert kind.componentName == null || kind.componentName.equals(componentName);
+        }
+
+        public static PolicyScope unknown(String moduleName) {
+            return new PolicyScope(UNKNOWN, UNKNOWN.componentName, moduleName);
+        }
+
+        public static PolicyScope server(String moduleName) {
+            return new PolicyScope(SERVER, SERVER.componentName, moduleName);
+        }
+
+        public static PolicyScope apmAgent(String moduleName) {
+            return new PolicyScope(APM_AGENT, APM_AGENT.componentName, moduleName);
+        }
+
+        public static PolicyScope plugin(String componentName, String moduleName) {
+            return new PolicyScope(PLUGIN, componentName, moduleName);
+        }
+    }
+
+    public enum ComponentKind {
+        UNKNOWN("(unknown)"),
+        SERVER("(server)"),
+        APM_AGENT("(APM agent)"),
+        PLUGIN(null);
+
+        /**
+         * If this kind corresponds to a single component, this is that component's name;
+         * otherwise null.
+         */
+        final String componentName;
+
+        ComponentKind(String componentName) {
+            this.componentName = componentName;
+        }
+    }
+
     record ModuleEntitlements(
         String componentName,
         Map<Class<? extends Entitlement>, List<Entitlement>> entitlementsByType,
@@ -132,7 +173,7 @@ ModuleEntitlements policyEntitlements(String componentName, Path componentPath,
     private final Map<String, List<Entitlement>> serverEntitlements;
     private final List<Entitlement> apmAgentEntitlements;
     private final Map<String, Map<String, List<Entitlement>>> pluginsEntitlements;
-    private final Function<Class<?>, String> pluginResolver;
+    private final Function<Class<?>, PolicyScope> scopeResolver;
     private final PathLookup pathLookup;
     private final Set<Class<?>> mutedClasses;
 
@@ -168,10 +209,6 @@ private static Set<Module> findSystemLayerModules() {
         .collect(Collectors.toUnmodifiableSet());
 
     private final Map<String, Path> sourcePaths;
-    /**
-     * The package name containing classes from the APM agent.
-     */
-    private final String apmAgentPackageName;
 
     /**
      * Frames originating from this module are ignored in the permission logic.
@@ -189,9 +226,8 @@ public PolicyManager(
         Policy serverPolicy,
         List<Entitlement> apmAgentEntitlements,
         Map<String, Policy> pluginPolicies,
-        Function<Class<?>, String> pluginResolver,
+        Function<Class<?>, PolicyScope> scopeResolver,
         Map<String, Path> sourcePaths,
-        String apmAgentPackageName,
         Module entitlementsModule,
         PathLookup pathLookup,
         Set<Class<?>> suppressFailureLogClasses
@@ -201,18 +237,17 @@ public PolicyManager(
         this.pluginsEntitlements = requireNonNull(pluginPolicies).entrySet()
             .stream()
             .collect(toUnmodifiableMap(Map.Entry::getKey, e -> buildScopeEntitlementsMap(e.getValue())));
-        this.pluginResolver = pluginResolver;
+        this.scopeResolver = scopeResolver;
         this.sourcePaths = sourcePaths;
-        this.apmAgentPackageName = apmAgentPackageName;
         this.entitlementsModule = entitlementsModule;
         this.pathLookup = requireNonNull(pathLookup);
         this.mutedClasses = suppressFailureLogClasses;
 
         List<ExclusiveFileEntitlement> exclusiveFileEntitlements = new ArrayList<>();
         for (var e : serverEntitlements.entrySet()) {
-            validateEntitlementsPerModule(SERVER_COMPONENT_NAME, e.getKey(), e.getValue(), exclusiveFileEntitlements);
+            validateEntitlementsPerModule(SERVER.componentName, e.getKey(), e.getValue(), exclusiveFileEntitlements);
         }
-        validateEntitlementsPerModule(APM_AGENT_COMPONENT_NAME, ALL_UNNAMED, apmAgentEntitlements, exclusiveFileEntitlements);
+        validateEntitlementsPerModule(APM_AGENT.componentName, ALL_UNNAMED, apmAgentEntitlements, exclusiveFileEntitlements);
         for (var p : pluginsEntitlements.entrySet()) {
             for (var m : p.getValue().entrySet()) {
                 validateEntitlementsPerModule(p.getKey(), m.getKey(), m.getValue(), exclusiveFileEntitlements);
@@ -607,50 +642,40 @@ ModuleEntitlements getEntitlements(Class<?> requestingClass) {
     }
 
     private ModuleEntitlements computeEntitlements(Class<?> requestingClass) {
-        Module requestingModule = requestingClass.getModule();
-        if (isServerModule(requestingModule)) {
-            return getModuleScopeEntitlements(
-                serverEntitlements,
-                requestingModule.getName(),
-                SERVER_COMPONENT_NAME,
-                getComponentPathFromClass(requestingClass)
-            );
-        }
+        var policyScope = scopeResolver.apply(requestingClass);
+        var componentName = policyScope.componentName();
+        var moduleName = policyScope.moduleName();
 
-        // plugins
-        var pluginName = pluginResolver.apply(requestingClass);
-        if (pluginName != null) {
-            var pluginEntitlements = pluginsEntitlements.get(pluginName);
-            if (pluginEntitlements == null) {
-                return defaultEntitlements(pluginName, sourcePaths.get(pluginName), requestingModule.getName());
-            } else {
+        switch (policyScope.kind()) {
+            case SERVER -> {
                 return getModuleScopeEntitlements(
-                    pluginEntitlements,
-                    getScopeName(requestingModule),
-                    pluginName,
-                    sourcePaths.get(pluginName)
+                    serverEntitlements,
+                    moduleName,
+                    SERVER.componentName,
+                    getComponentPathFromClass(requestingClass)
                 );
             }
-        }
-
-        if (requestingModule.isNamed() == false && requestingClass.getPackageName().startsWith(apmAgentPackageName)) {
-            // The APM agent is the only thing running non-modular in the system classloader
-            return policyEntitlements(
-                APM_AGENT_COMPONENT_NAME,
-                getComponentPathFromClass(requestingClass),
-                ALL_UNNAMED,
-                apmAgentEntitlements
-            );
-        }
-
-        return defaultEntitlements(UNKNOWN_COMPONENT_NAME, null, requestingModule.getName());
-    }
-
-    private static String getScopeName(Module requestingModule) {
-        if (requestingModule.isNamed() == false) {
-            return ALL_UNNAMED;
-        } else {
-            return requestingModule.getName();
+            case APM_AGENT -> {
+                // The APM agent is the only thing running non-modular in the system classloader
+                return policyEntitlements(
+                    APM_AGENT.componentName,
+                    getComponentPathFromClass(requestingClass),
+                    ALL_UNNAMED,
+                    apmAgentEntitlements
+                );
+            }
+            case UNKNOWN -> {
+                return defaultEntitlements(UNKNOWN.componentName, null, moduleName);
+            }
+            default -> {
+                assert policyScope.kind() == PLUGIN;
+                var pluginEntitlements = pluginsEntitlements.get(componentName);
+                if (pluginEntitlements == null) {
+                    return defaultEntitlements(componentName, sourcePaths.get(componentName), moduleName);
+                } else {
+                    return getModuleScopeEntitlements(pluginEntitlements, moduleName, componentName, sourcePaths.get(componentName));
+                }
+            }
         }
     }
 
@@ -686,10 +711,6 @@ private ModuleEntitlements getModuleScopeEntitlements(
         return policyEntitlements(componentName, componentPath, scopeName, entitlements);
     }
 
-    private static boolean isServerModule(Module requestingModule) {
-        return requestingModule.isNamed() && requestingModule.getLayer() == ModuleLayer.boot();
-    }
-
     /**
      * Walks the stack to determine which class should be checked for entitlements.
      *