[Entitlements] Integrate PluginsLoader with PolicyManager

libs/entitlement/src/main/java/module-info.java

@@ -17,6 +17,7 @@
     requires static org.elasticsearch.entitlement.bridge; // At runtime, this will be in java.base
 
     exports org.elasticsearch.entitlement.runtime.api;
+    exports org.elasticsearch.entitlement.runtime.policy;
     exports org.elasticsearch.entitlement.instrumentation;
     exports org.elasticsearch.entitlement.bootstrap to org.elasticsearch.server;
     exports org.elasticsearch.entitlement.initialization to java.base;

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

@@ -18,6 +18,8 @@
 import org.elasticsearch.entitlement.instrumentation.MethodKey;
 import org.elasticsearch.entitlement.instrumentation.Transformer;
 import org.elasticsearch.entitlement.runtime.api.ElasticsearchEntitlementChecker;
+import org.elasticsearch.entitlement.runtime.policy.CreateClassLoaderEntitlement;
+import org.elasticsearch.entitlement.runtime.policy.ExitVMEntitlement;
 import org.elasticsearch.entitlement.runtime.policy.Policy;
 import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
 import org.elasticsearch.entitlement.runtime.policy.PolicyParser;
@@ -86,9 +88,11 @@ private static Class<?> internalNameToClass(String internalName) {
     private static PolicyManager createPolicyManager() throws IOException {
         Map<String, Policy> pluginPolicies = createPluginPolicies(EntitlementBootstrap.bootstrapArgs().pluginData());
 
-        // TODO: What should the name be?
         // TODO(ES-10031): Decide what goes in the elasticsearch default policy and extend it
-        var serverPolicy = new Policy("server", List.of());
+        var serverPolicy = new Policy(
+            "server",
+            List.of(new Scope("org.elasticsearch.server", List.of(new ExitVMEntitlement(), new CreateClassLoaderEntitlement())))
+        );
         return new PolicyManager(serverPolicy, pluginPolicies, EntitlementBootstrap.bootstrapArgs().pluginResolver());
     }
 

libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/api/ElasticsearchEntitlementChecker.java

@@ -10,7 +10,6 @@
 package org.elasticsearch.entitlement.runtime.api;
 
 import org.elasticsearch.entitlement.bridge.EntitlementChecker;
-import org.elasticsearch.entitlement.runtime.policy.FlagEntitlementType;
 import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
 
 import java.net.URL;
@@ -30,27 +29,27 @@ public ElasticsearchEntitlementChecker(PolicyManager policyManager) {
 
     @Override
     public void check$java_lang_System$exit(Class<?> callerClass, int status) {
-        policyManager.checkFlagEntitlement(callerClass, FlagEntitlementType.SYSTEM_EXIT);
+        policyManager.checkExitVM(callerClass);
     }
 
     @Override
     public void check$java_net_URLClassLoader$(Class<?> callerClass, URL[] urls) {
-        policyManager.checkFlagEntitlement(callerClass, FlagEntitlementType.CREATE_CLASSLOADER);
+        policyManager.checkCreateClassLoader(callerClass);
     }
 
     @Override
     public void check$java_net_URLClassLoader$(Class<?> callerClass, URL[] urls, ClassLoader parent) {
-        policyManager.checkFlagEntitlement(callerClass, FlagEntitlementType.CREATE_CLASSLOADER);
+        policyManager.checkCreateClassLoader(callerClass);
     }
 
     @Override
     public void check$java_net_URLClassLoader$(Class<?> callerClass, URL[] urls, ClassLoader parent, URLStreamHandlerFactory factory) {
-        policyManager.checkFlagEntitlement(callerClass, FlagEntitlementType.CREATE_CLASSLOADER);
+        policyManager.checkCreateClassLoader(callerClass);
     }
 
     @Override
     public void check$java_net_URLClassLoader$(Class<?> callerClass, String name, URL[] urls, ClassLoader parent) {
-        policyManager.checkFlagEntitlement(callerClass, FlagEntitlementType.CREATE_CLASSLOADER);
+        policyManager.checkCreateClassLoader(callerClass);
     }
 
     @Override
@@ -61,6 +60,6 @@ public ElasticsearchEntitlementChecker(PolicyManager policyManager) {
         ClassLoader parent,
         URLStreamHandlerFactory factory
     ) {
-        policyManager.checkFlagEntitlement(callerClass, FlagEntitlementType.CREATE_CLASSLOADER);
+        policyManager.checkCreateClassLoader(callerClass);
     }
 }

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

@@ -12,5 +12,4 @@
 public class CreateClassLoaderEntitlement implements Entitlement {
     @ExternalEntitlement
     public CreateClassLoaderEntitlement() {}
-
 }

libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/FlagEntitlementType.java → libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/ExitVMEntitlement.java

@@ -9,7 +9,7 @@
 
 package org.elasticsearch.entitlement.runtime.policy;
 
-public enum FlagEntitlementType {
-    SYSTEM_EXIT,
-    CREATE_CLASSLOADER;
-}
+/**
+ * Internal policy type (not-parseable -- not available to plugins).
+ */
+public class ExitVMEntitlement implements Entitlement {}

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

@@ -20,6 +20,9 @@ public class FileEntitlement implements Entitlement {
     public static final int READ_ACTION = 0x1;
     public static final int WRITE_ACTION = 0x2;
 
+    public static final String READ = "read";
+    public static final String WRITE = "write";
+
     private final String path;
     private final int actions;
 
@@ -29,12 +32,12 @@ public FileEntitlement(String path, List<String> actionsList) {
         int actionsInt = 0;
 
         for (String actionString : actionsList) {
-            if ("read".equals(actionString)) {
+            if (READ.equals(actionString)) {
                 if ((actionsInt & READ_ACTION) == READ_ACTION) {
                     throw new IllegalArgumentException("file action [read] specified multiple times");
                 }
                 actionsInt |= READ_ACTION;
-            } else if ("write".equals(actionString)) {
+            } else if (WRITE.equals(actionString)) {
                 if ((actionsInt & WRITE_ACTION) == WRITE_ACTION) {
                     throw new IllegalArgumentException("file action [write] specified multiple times");
                 }

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

@@ -17,17 +17,45 @@
 
 import java.lang.module.ModuleFinder;
 import java.lang.module.ModuleReference;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.IdentityHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class PolicyManager {
     private static final Logger logger = LogManager.getLogger(ElasticsearchEntitlementChecker.class);
 
+    static class ModuleEntitlements {
+        public static final ModuleEntitlements NONE = new ModuleEntitlements(List.of());
+        private final IdentityHashMap<Class<? extends Entitlement>, List<Entitlement>> entitlementsByType;
+
+        ModuleEntitlements(List<Entitlement> entitlements) {
+            this.entitlementsByType = entitlements.stream()
+                .collect(Collectors.toMap(Entitlement::getClass, e -> new ArrayList<>(List.of(e)), (a, b) -> {
+                    a.addAll(b);
+                    return a;
+                }, IdentityHashMap::new));
+        }
+
+        public boolean hasEntitlement(Class<? extends Entitlement> entitlementClass) {
+            return entitlementsByType.containsKey(entitlementClass);
+        }
+
+        public <E extends Entitlement> Stream<E> getEntitlements(Class<E> entitlementClass) {
+            return entitlementsByType.get(entitlementClass).stream().map(entitlementClass::cast);
+        }
+    }
+
+    final Map<Module, ModuleEntitlements> moduleEntitlementsMap = new HashMap<>();
+
     protected final Policy serverPolicy;
     protected final Map<String, Policy> pluginPolicies;
     private final Function<Class<?>, String> pluginResolver;
@@ -56,27 +84,110 @@ public PolicyManager(Policy defaultPolicy, Map<String, Policy> pluginPolicies, F
         this.pluginResolver = pluginResolver;
     }
 
-    public void checkFlagEntitlement(Class<?> callerClass, FlagEntitlementType type) {
+    private static List<Entitlement> lookupEntitlementsForModule(Policy policy, String moduleName) {
+        for (int i = 0; i < policy.scopes.size(); ++i) {
+            var scope = policy.scopes.get(i);
+            if (scope.name.equals(moduleName)) {
+                return scope.entitlements;
+            }
+        }
+        return null;
+    }
+
+    public void checkExitVM(Class<?> callerClass) {
+        checkEntitlementPresent(callerClass, ExitVMEntitlement.class);
+    }
+
+    public void checkCreateClassLoader(Class<?> callerClass) {
+        checkEntitlementPresent(callerClass, CreateClassLoaderEntitlement.class);
+    }
+
+    private void checkEntitlementPresent(Class<?> callerClass, Class<? extends Entitlement> entitlementClass) {
         var requestingModule = requestingModule(callerClass);
         if (isTriviallyAllowed(requestingModule)) {
             return;
         }
 
-        // TODO: real policy check. For now, we only allow our hardcoded System.exit policy for server.
-        // TODO: this will be checked using policies
-        if (requestingModule.isNamed()
-            && requestingModule.getName().equals("org.elasticsearch.server")
-            && (type == FlagEntitlementType.SYSTEM_EXIT || type == FlagEntitlementType.CREATE_CLASSLOADER)) {
-            logger.debug("Allowed: caller [{}] in module [{}] has entitlement [{}]", callerClass, requestingModule.getName(), type);
+        ModuleEntitlements entitlements = getEntitlementsOrThrow(callerClass, requestingModule);
+        if (entitlements.hasEntitlement(entitlementClass)) {
+            logger.debug(
+                () -> Strings.format(
+                    "Entitled: caller [%s], module [%s], type [%s]",
+                    callerClass,
+                    requestingModule.getName(),
+                    entitlementClass.getSimpleName()
+                )
+            );
             return;
         }
-
-        // TODO: plugins policy check using pluginResolver and pluginPolicies
         throw new NotEntitledException(
-            Strings.format("Missing entitlement [%s] for caller [%s] in module [%s]", type, callerClass, requestingModule.getName())
+            Strings.format(
+                "Missing entitlement: caller [%s], module [%s], type [%s]",
+                callerClass,
+                requestingModule.getName(),
+                entitlementClass.getSimpleName()
+            )
         );
     }
 
+    ModuleEntitlements getEntitlementsOrThrow(Class<?> callerClass, Module requestingModule) {
+        ModuleEntitlements cachedEntitlement = moduleEntitlementsMap.get(requestingModule);
+        if (cachedEntitlement != null) {
+            if (cachedEntitlement == ModuleEntitlements.NONE) {
+                throw new NotEntitledException(buildModuleNoPolicyMessage(callerClass, requestingModule) + "[CACHED]");
+            }
+            return cachedEntitlement;
+        }
+
+        if (isServerModule(requestingModule)) {
+            var scopeName = requestingModule.getName();
+            return getModuleEntitlementsOrThrow(callerClass, requestingModule, serverPolicy, scopeName);
+        }
+
+        // plugins
+        var pluginName = pluginResolver.apply(callerClass);
+        if (pluginName != null) {
+            var pluginPolicy = pluginPolicies.get(pluginName);
+            if (pluginPolicy != null) {
+                final String scopeName;
+                if (requestingModule.isNamed() == false) {
+                    scopeName = ALL_UNNAMED;
+                } else {
+                    scopeName = requestingModule.getName();
+                }
+                return getModuleEntitlementsOrThrow(callerClass, requestingModule, pluginPolicy, scopeName);
+            }
+        }
+
+        moduleEntitlementsMap.put(requestingModule, ModuleEntitlements.NONE);
+        throw new NotEntitledException(buildModuleNoPolicyMessage(callerClass, requestingModule));
+    }
+
+    private static String buildModuleNoPolicyMessage(Class<?> callerClass, Module requestingModule) {
+        return Strings.format("Missing entitlement policy: caller [%s], module [%s]", callerClass, requestingModule.getName());
+    }
+
+    private ModuleEntitlements getModuleEntitlementsOrThrow(Class<?> callerClass, Module module, Policy policy, String moduleName) {
+        var entitlements = lookupEntitlementsForModule(policy, moduleName);
+        if (entitlements == null) {
+            // Module without entitlements - remember we don't have any
+            moduleEntitlementsMap.put(module, ModuleEntitlements.NONE);
+            throw new NotEntitledException(buildModuleNoPolicyMessage(callerClass, module));
+        }
+        // We have a policy for this module
+        var classEntitlements = createClassEntitlements(entitlements);
+        moduleEntitlementsMap.put(module, classEntitlements);
+        return classEntitlements;
+    }
+
+    private static boolean isServerModule(Module requestingModule) {
+        return requestingModule.isNamed() && requestingModule.getLayer() == ModuleLayer.boot();
+    }
+
+    private ModuleEntitlements createClassEntitlements(List<Entitlement> entitlements) {
+        return new ModuleEntitlements(entitlements);
+    }
+
     private static Module requestingModule(Class<?> callerClass) {
         if (callerClass != null) {
             Module callerModule = callerClass.getModule();
@@ -102,10 +213,10 @@ private static Module requestingModule(Class<?> callerClass) {
 
     private static boolean isTriviallyAllowed(Module requestingModule) {
         if (requestingModule == null) {
-            logger.debug("Trivially allowed: entire call stack is in composed of classes in system modules");
+            logger.debug("Entitlement trivially allowed: entire call stack is in composed of classes in system modules");
             return true;
         }
-        logger.trace("Not trivially allowed");
+        logger.trace("Entitlement not trivially allowed");
         return false;
     }