[8.x] [Entitlelments] Server policy patching via system property (#124904) (#125008)

* [Entitlelments] Server policy patching via system property (#124904)

This PR adds a mechanism to patch the server layer policy by merging entitlements from a policy provided via a system property (merge, not replace).

* change to support 8.x language level

Author: ldematte

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

@@ -14,6 +14,7 @@
 import com.sun.tools.attach.AttachNotSupportedException;
 import com.sun.tools.attach.VirtualMachine;
 
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.entitlement.initialization.EntitlementInitialization;
 import org.elasticsearch.entitlement.runtime.policy.Policy;
@@ -33,6 +34,7 @@
 public class EntitlementBootstrap {
 
     public record BootstrapArgs(
+        @Nullable Policy serverPolicyPatch,
         Map<String, Policy> pluginPolicies,
         Function<Class<?>, String> pluginResolver,
         Function<String, Stream<String>> settingResolver,
@@ -78,6 +80,7 @@ public static BootstrapArgs bootstrapArgs() {
      * Activates entitlement checking. Once this method returns, calls to methods protected by Entitlements from classes without a valid
      * policy will throw {@link org.elasticsearch.entitlement.runtime.api.NotEntitledException}.
      *
+     * @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 settingResolver a functor to resolve a setting name pattern for one or more Elasticsearch settings.
@@ -94,6 +97,7 @@ public static BootstrapArgs bootstrapArgs() {
      * @param suppressFailureLogClasses   classes for which we do not need or want to log Entitlements failures
      */
     public static void bootstrap(
+        Policy serverPolicyPatch,
         Map<String, Policy> pluginPolicies,
         Function<Class<?>, String> pluginResolver,
         Function<String, Stream<String>> settingResolver,
@@ -114,6 +118,7 @@ public static void bootstrap(
             throw new IllegalStateException("plugin data is already set");
         }
         EntitlementBootstrap.bootstrapArgs = new BootstrapArgs(
+            serverPolicyPatch,
             pluginPolicies,
             pluginResolver,
             settingResolver,

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

@@ -23,6 +23,7 @@
 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.PolicyUtils;
 import org.elasticsearch.entitlement.runtime.policy.Scope;
 import org.elasticsearch.entitlement.runtime.policy.entitlements.CreateClassLoaderEntitlement;
 import org.elasticsearch.entitlement.runtime.policy.entitlements.Entitlement;
@@ -269,8 +270,13 @@ private static PolicyManager createPolicyManager() {
             );
         }
 
-        // TODO(ES-10031): Decide what goes in the elasticsearch default policy and extend it
-        var serverPolicy = new Policy("server", serverScopes);
+        var serverPolicy = new Policy(
+            "server",
+            bootstrapArgs.serverPolicyPatch() == null
+                ? serverScopes
+                : PolicyUtils.mergeScopes(serverScopes, bootstrapArgs.serverPolicyPatch().scopes())
+        );
+
         // 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

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

@@ -141,9 +141,9 @@ ModuleEntitlements policyEntitlements(String componentName, Path componentPath,
 
     public static final String ALL_UNNAMED = "ALL-UNNAMED";
 
-    private static final Set<Module> systemModules = findSystemModules();
+    private static final Set<Module> SYSTEM_LAYER_MODULES = findSystemLayerModules();
 
-    private static Set<Module> findSystemModules() {
+    private static Set<Module> findSystemLayerModules() {
         var systemModulesDescriptors = ModuleFinder.ofSystem()
             .findAll()
             .stream()
@@ -163,6 +163,13 @@ private static Set<Module> findSystemModules() {
         ).collect(Collectors.toUnmodifiableSet());
     }
 
+    // Anything in the boot layer that is not in the system layer, is in the server layer
+    public static final Set<Module> SERVER_LAYER_MODULES = ModuleLayer.boot()
+        .modules()
+        .stream()
+        .filter(m -> SYSTEM_LAYER_MODULES.contains(m) == false)
+        .collect(Collectors.toUnmodifiableSet());
+
     private final Map<String, Path> sourcePaths;
     /**
      * The package name containing classes from the APM agent.
@@ -725,7 +732,7 @@ private static boolean isTriviallyAllowed(Class<?> requestingClass) {
             generalLogger.debug("Entitlement trivially allowed: no caller frames outside the entitlement library");
             return true;
         }
-        if (systemModules.contains(requestingClass.getModule())) {
+        if (SYSTEM_LAYER_MODULES.contains(requestingClass.getModule())) {
             generalLogger.debug("Entitlement trivially allowed from system module [{}]", requestingClass.getModule().getName());
             return true;
         }

libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyParserUtils.java renamed to libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyUtils.java

@@ -10,6 +10,9 @@
 package org.elasticsearch.entitlement.runtime.policy;
 
 import org.elasticsearch.core.Strings;
+import org.elasticsearch.entitlement.runtime.policy.entitlements.Entitlement;
+import org.elasticsearch.entitlement.runtime.policy.entitlements.FilesEntitlement;
+import org.elasticsearch.entitlement.runtime.policy.entitlements.WriteSystemPropertiesEntitlement;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 
@@ -20,21 +23,23 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardOpenOption;
+import java.util.ArrayList;
 import java.util.Base64;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static java.util.Objects.requireNonNull;
 import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ALL_UNNAMED;
 
-public class PolicyParserUtils {
+public class PolicyUtils {
 
-    private static final Logger logger = LogManager.getLogger(PolicyParserUtils.class);
+    private static final Logger logger = LogManager.getLogger(PolicyUtils.class);
 
     public record PluginData(Path pluginPath, boolean isModular, boolean isExternalPlugin) {
         public PluginData {
@@ -44,8 +49,6 @@ public record PluginData(Path pluginPath, boolean isModular, boolean isExternalP
 
     private static final String POLICY_FILE_NAME = "entitlement-policy.yaml";
 
-    public static final String POLICY_OVERRIDE_PREFIX = "es.entitlements.policy.";
-
     public static Map<String, Policy> createPluginPolicies(Collection<PluginData> pluginData, Map<String, String> overrides, String version)
         throws IOException {
         Map<String, Policy> pluginPolicies = new HashMap<>(pluginData.size());
@@ -54,9 +57,15 @@ public static Map<String, Policy> createPluginPolicies(Collection<PluginData> pl
             String pluginName = pluginRoot.getFileName().toString();
             final Set<String> moduleNames = getModuleNames(pluginRoot, entry.isModular());
 
-            var overriddenPolicy = parsePolicyOverrideIfExists(overrides, version, entry.isExternalPlugin(), pluginName, moduleNames);
-            if (overriddenPolicy.isPresent()) {
-                pluginPolicies.put(pluginName, overriddenPolicy.get());
+            var overriddenPolicy = parseEncodedPolicyIfExists(
+                overrides.get(pluginName),
+                version,
+                entry.isExternalPlugin(),
+                pluginName,
+                moduleNames
+            );
+            if (overriddenPolicy != null) {
+                pluginPolicies.put(pluginName, overriddenPolicy);
             } else {
                 Path policyFile = pluginRoot.resolve(POLICY_FILE_NAME);
                 var policy = parsePolicyIfExists(pluginName, policyFile, entry.isExternalPlugin());
@@ -67,59 +76,54 @@ public static Map<String, Policy> createPluginPolicies(Collection<PluginData> pl
         return pluginPolicies;
     }
 
-    static Optional<Policy> parsePolicyOverrideIfExists(
-        Map<String, String> overrides,
+    public static Policy parseEncodedPolicyIfExists(
+        String encodedPolicy,
         String version,
         boolean externalPlugin,
-        String pluginName,
+        String layerName,
         Set<String> moduleNames
     ) {
-        var policyOverride = overrides.get(pluginName);
-        if (policyOverride != null) {
+        if (encodedPolicy != null) {
             try {
-                var versionedPolicy = decodeOverriddenPluginPolicy(policyOverride, pluginName, externalPlugin);
-                validatePolicyScopes(pluginName, versionedPolicy.policy(), moduleNames, "<override>");
+                var versionedPolicy = decodeEncodedPolicy(encodedPolicy, layerName, externalPlugin);
+                validatePolicyScopes(layerName, versionedPolicy.policy(), moduleNames, "<patch>");
 
                 // Empty versions defaults to "any"
                 if (versionedPolicy.versions().isEmpty() || versionedPolicy.versions().contains(version)) {
-                    logger.info("Using policy override for plugin [{}]", pluginName);
-                    return Optional.of(versionedPolicy.policy());
+                    logger.info("Using policy patch for layer [{}]", layerName);
+                    return versionedPolicy.policy();
                 } else {
                     logger.warn(
-                        "Found a policy override with version mismatch. The override will not be applied. "
-                            + "Plugin [{}]; policy versions [{}]; current version [{}]",
-                        pluginName,
+                        "Found a policy patch with version mismatch. The patch will not be applied. "
+                            + "Layer [{}]; policy versions [{}]; current version [{}]",
+                        layerName,
                         String.join(",", versionedPolicy.versions()),
                         version
                     );
                 }
             } catch (Exception ex) {
                 logger.warn(
-                    Strings.format(
-                        "Found a policy override with invalid content. The override will not be applied. Plugin [%s]",
-                        pluginName
-                    ),
+                    Strings.format("Found a policy patch with invalid content. The patch will not be applied. Layer [%s]", layerName),
                     ex
                 );
             }
         }
-        return Optional.empty();
+        return null;
     }
 
-    static VersionedPolicy decodeOverriddenPluginPolicy(String base64String, String pluginName, boolean isExternalPlugin)
-        throws IOException {
+    static VersionedPolicy decodeEncodedPolicy(String base64String, String layerName, boolean isExternalPlugin) throws IOException {
         byte[] policyDefinition = Base64.getDecoder().decode(base64String);
-        return new PolicyParser(new ByteArrayInputStream(policyDefinition), pluginName, isExternalPlugin).parseVersionedPolicy();
+        return new PolicyParser(new ByteArrayInputStream(policyDefinition), layerName, isExternalPlugin).parseVersionedPolicy();
     }
 
-    private static void validatePolicyScopes(String pluginName, Policy policy, Set<String> moduleNames, String policyLocation) {
+    private static void validatePolicyScopes(String layerName, Policy policy, Set<String> moduleNames, String policyLocation) {
         // TODO: should this check actually be part of the parser?
         for (Scope scope : policy.scopes()) {
             if (moduleNames.contains(scope.moduleName()) == false) {
                 throw new IllegalStateException(
                     Strings.format(
-                        "Invalid module name in policy: plugin [%s] does not have module [%s]; available modules [%s]; policy path [%s]",
-                        pluginName,
+                        "Invalid module name in policy: layer [%s] does not have module [%s]; available modules [%s]; policy path [%s]",
+                        layerName,
                         scope.moduleName(),
                         String.join(", ", moduleNames),
                         policyLocation
@@ -147,4 +151,51 @@ private static Set<String> getModuleNames(Path pluginRoot, boolean isModular) {
         return Set.of(ALL_UNNAMED);
     }
 
+    public static List<Scope> mergeScopes(List<Scope> mainScopes, List<Scope> additionalScopes) {
+        var result = new ArrayList<Scope>();
+        var additionalScopesMap = additionalScopes.stream().collect(Collectors.toMap(Scope::moduleName, Scope::entitlements));
+        for (var mainScope : mainScopes) {
+            List<Entitlement> additionalEntitlements = additionalScopesMap.remove(mainScope.moduleName());
+            if (additionalEntitlements == null) {
+                result.add(mainScope);
+            } else {
+                result.add(new Scope(mainScope.moduleName(), mergeEntitlements(mainScope.entitlements(), additionalEntitlements)));
+            }
+        }
+
+        for (var remainingEntry : additionalScopesMap.entrySet()) {
+            result.add(new Scope(remainingEntry.getKey(), remainingEntry.getValue()));
+        }
+        return result;
+    }
+
+    static List<Entitlement> mergeEntitlements(List<Entitlement> a, List<Entitlement> b) {
+        Map<Class<? extends Entitlement>, Entitlement> entitlementMap = a.stream()
+            .collect(Collectors.toMap(Entitlement::getClass, Function.identity()));
+
+        for (var entitlement : b) {
+            entitlementMap.merge(entitlement.getClass(), entitlement, PolicyUtils::mergeEntitlement);
+        }
+        return entitlementMap.values().stream().toList();
+    }
+
+    static Entitlement mergeEntitlement(Entitlement entitlement1, Entitlement entitlement2) {
+        if (entitlement1 instanceof FilesEntitlement e) {
+            return merge(e, (FilesEntitlement) entitlement2);
+        }
+        if (entitlement1 instanceof WriteSystemPropertiesEntitlement e) {
+            return merge(e, (WriteSystemPropertiesEntitlement) entitlement2);
+        }
+        return entitlement1;
+    }
+
+    private static FilesEntitlement merge(FilesEntitlement a, FilesEntitlement b) {
+        return new FilesEntitlement(Stream.concat(a.filesData().stream(), b.filesData().stream()).distinct().toList());
+    }
+
+    private static WriteSystemPropertiesEntitlement merge(WriteSystemPropertiesEntitlement a, WriteSystemPropertiesEntitlement b) {
+        return new WriteSystemPropertiesEntitlement(
+            Stream.concat(a.properties().stream(), b.properties().stream()).collect(Collectors.toUnmodifiableSet())
+        );
+    }
 }