diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java index a77c86d5ffd0..8d3efe4eb98e 100644 --- a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java +++ b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java @@ -18,7 +18,6 @@ import org.elasticsearch.logging.Logger; 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; @@ -56,8 +55,8 @@ public class PolicyManager { final Map moduleEntitlementsMap = new HashMap<>(); - protected final Policy serverPolicy; - protected final Map pluginPolicies; + protected final Map> serverEntitlements; + protected final Map>> pluginsEntitlements; private final Function, String> pluginResolver; public static final String ALL_UNNAMED = "ALL-UNNAMED"; @@ -79,19 +78,16 @@ public class PolicyManager { } public PolicyManager(Policy defaultPolicy, Map pluginPolicies, Function, String> pluginResolver) { - this.serverPolicy = Objects.requireNonNull(defaultPolicy); - this.pluginPolicies = Collections.unmodifiableMap(Objects.requireNonNull(pluginPolicies)); + this.serverEntitlements = buildScopeEntitlementsMap(Objects.requireNonNull(defaultPolicy)); + this.pluginsEntitlements = Objects.requireNonNull(pluginPolicies) + .entrySet() + .stream() + .collect(Collectors.toUnmodifiableMap(Map.Entry::getKey, e -> buildScopeEntitlementsMap(e.getValue()))); this.pluginResolver = pluginResolver; } - private static List 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; + private static Map> buildScopeEntitlementsMap(Policy policy) { + return policy.scopes.stream().collect(Collectors.toUnmodifiableMap(scope -> scope.name, scope -> scope.entitlements)); } public void checkExitVM(Class callerClass) { @@ -141,21 +137,21 @@ public class PolicyManager { if (isServerModule(requestingModule)) { var scopeName = requestingModule.getName(); - return getModuleEntitlementsOrThrow(callerClass, requestingModule, serverPolicy, scopeName); + return getModuleEntitlementsOrThrow(callerClass, requestingModule, serverEntitlements, scopeName); } // plugins var pluginName = pluginResolver.apply(callerClass); if (pluginName != null) { - var pluginPolicy = pluginPolicies.get(pluginName); - if (pluginPolicy != null) { + var pluginEntitlements = pluginsEntitlements.get(pluginName); + if (pluginEntitlements != null) { final String scopeName; if (requestingModule.isNamed() == false) { scopeName = ALL_UNNAMED; } else { scopeName = requestingModule.getName(); } - return getModuleEntitlementsOrThrow(callerClass, requestingModule, pluginPolicy, scopeName); + return getModuleEntitlementsOrThrow(callerClass, requestingModule, pluginEntitlements, scopeName); } } @@ -167,15 +163,20 @@ public class PolicyManager { 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); + private ModuleEntitlements getModuleEntitlementsOrThrow( + Class callerClass, + Module module, + Map> scopeEntitlements, + String moduleName + ) { + var entitlements = scopeEntitlements.get(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); + var classEntitlements = new ModuleEntitlements(entitlements); moduleEntitlementsMap.put(module, classEntitlements); return classEntitlements; } @@ -184,10 +185,6 @@ public class PolicyManager { return requestingModule.isNamed() && requestingModule.getLayer() == ModuleLayer.boot(); } - private ModuleEntitlements createClassEntitlements(List entitlements) { - return new ModuleEntitlements(entitlements); - } - private static Module requestingModule(Class callerClass) { if (callerClass != null) { Module callerModule = callerClass.getModule(); @@ -222,6 +219,6 @@ public class PolicyManager { @Override public String toString() { - return "PolicyManager{" + "serverPolicy=" + serverPolicy + ", pluginPolicies=" + pluginPolicies + '}'; + return "PolicyManager{" + "serverEntitlements=" + serverEntitlements + ", pluginsEntitlements=" + pluginsEntitlements + '}'; } }