Produce a PolicyScope for excluded system modules.

Fortunately, the JDK is still modular even if the tests are run sans modules,
so we can use the module API to identify the packages in every system module.

Author: prdoyle

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

@@ -54,7 +54,7 @@ public class PolicyManager {
      */
     static final Logger generalLogger = LogManager.getLogger(PolicyManager.class);
 
-    static final Set<String> MODULES_EXCLUDED_FROM_SYSTEM_MODULES = Set.of("java.desktop", "java.xml");
+    public static final Set<String> MODULES_EXCLUDED_FROM_SYSTEM_MODULES = Set.of("java.desktop", "java.xml");
 
     /**
      * Identifies a particular entitlement {@link Scope} within a {@link Policy}.

test/framework/src/main/java/org/elasticsearch/bootstrap/TestScopeResolver.java

@@ -9,11 +9,14 @@
 
 package org.elasticsearch.bootstrap;
 
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.SuppressForbidden;
-import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
+import org.elasticsearch.entitlement.runtime.policy.PolicyManager.PolicyScope;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 
+import java.lang.module.ModuleDescriptor;
+import java.lang.module.ModuleFinder;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.util.List;
@@ -22,49 +25,78 @@
 import java.util.TreeMap;
 import java.util.function.Function;
 
+import static java.util.Objects.requireNonNull;
+import static java.util.stream.Collectors.toSet;
 import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ALL_UNNAMED;
 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.MODULES_EXCLUDED_FROM_SYSTEM_MODULES;
 
-public record TestScopeResolver(Map<String, PolicyManager.PolicyScope> scopeMap) {
+public final class TestScopeResolver {
 
     private static final Logger logger = LogManager.getLogger(TestScopeResolver.class);
+    private final Map<String, PolicyScope> scopeMap;
+    private static final Map<String, PolicyScope> excludedSystemPackageScopes = computeExcludedSystemPackageScopes();
 
-    PolicyManager.PolicyScope getScope(Class<?> callerClass) {
+    public TestScopeResolver(Map<String, PolicyScope> scopeMap) {
+        this.scopeMap = scopeMap;
+    }
+
+    private static Map<String, PolicyScope> computeExcludedSystemPackageScopes() {
+        // Within any one module layer, module names are unique, so we just need the names
+        Set<String> systemModuleNames = ModuleFinder.ofSystem()
+            .findAll()
+            .stream()
+            .map(ref -> ref.descriptor().name())
+            .filter(MODULES_EXCLUDED_FROM_SYSTEM_MODULES::contains)
+            .collect(toSet());
+
+        Map<String, PolicyScope> result = new TreeMap<>();
+        ModuleLayer.boot().modules().stream().filter(m -> systemModuleNames.contains(m.getName())).forEach(m -> {
+            ModuleDescriptor desc = m.getDescriptor();
+            if (desc != null) {
+                desc.packages().forEach(pkg ->
+                // Our component identification logic returns SERVER for these
+                result.put(pkg, new PolicyScope(SERVER, SERVER.componentName, m.getName())));
+            }
+        });
+        return result;
+    }
+
+    public static @Nullable PolicyScope getExcludedSystemPackageScope(Class<?> callerClass) {
+        return excludedSystemPackageScopes.get(callerClass.getPackageName());
+    }
+
+    PolicyScope getScope(Class<?> callerClass) {
         var callerCodeSource = callerClass.getProtectionDomain().getCodeSource();
         if (callerCodeSource == null) {
-            // This case happens for JDK modules. Usually those are trivially allowed, but some are excluded,
-            // and those end up here.
-            // We have no test build info for those modules, so for now, let's just guess.
-            if (callerClass.getPackageName().equals("sun.java2d")) {
-                return new PolicyManager.PolicyScope(SERVER, SERVER.componentName, "java.desktop");
-            } else {
-                throw new IllegalArgumentException("Cannot identify scope for JDK class [" + callerClass + "]");
-            }
+            // This only happens for JDK classes. Furthermore, for trivially allowed modules, we shouldn't even get here.
+            // Hence, this must be an excluded system module, so check for that.
+            return requireNonNull(getExcludedSystemPackageScope(callerClass));
         }
 
         var location = callerCodeSource.getLocation().toString();
         var scope = scopeMap.get(location);
         if (scope == null) {
             // Special cases for libraries not handled by our automatically-generated scopeMap
             if (callerClass.getPackageName().startsWith("org.bouncycastle")) {
-                scope = new PolicyManager.PolicyScope(PLUGIN, "security", ALL_UNNAMED);
+                scope = new PolicyScope(PLUGIN, "security", ALL_UNNAMED);
                 logger.debug("Assuming bouncycastle is part of the security plugin");
             }
         }
         if (scope == null) {
             logger.warn("Cannot identify a scope for class [{}], location [{}]", callerClass.getName(), location);
-            return PolicyManager.PolicyScope.unknown(location);
+            return PolicyScope.unknown(location);
         }
         return scope;
     }
 
-    public static Function<Class<?>, PolicyManager.PolicyScope> createScopeResolver(
+    public static Function<Class<?>, PolicyScope> createScopeResolver(
         TestBuildInfo serverBuildInfo,
         List<TestBuildInfo> pluginsBuildInfo,
         Set<String> modularPlugins
     ) {
-        Map<String, PolicyManager.PolicyScope> scopeMap = new TreeMap<>(); // Sorted to make it easier to read during debugging
+        Map<String, PolicyScope> scopeMap = new TreeMap<>(); // Sorted to make it easier to read during debugging
         for (var pluginBuildInfo : pluginsBuildInfo) {
             boolean isModular = modularPlugins.contains(pluginBuildInfo.component());
             for (var location : pluginBuildInfo.locations()) {
@@ -76,7 +108,7 @@ public static Function<Class<?>, PolicyManager.PolicyScope> createScopeResolver(
                     String module = isModular ? location.module() : ALL_UNNAMED;
                     scopeMap.put(
                         getCodeSource(codeSource, location.representativeClass()),
-                        PolicyManager.PolicyScope.plugin(pluginBuildInfo.component(), module)
+                        PolicyScope.plugin(pluginBuildInfo.component(), module)
                     );
                 } catch (MalformedURLException e) {
                     throw new IllegalArgumentException("Cannot locate class [" + location.representativeClass() + "]", e);
@@ -91,7 +123,7 @@ public static Function<Class<?>, PolicyManager.PolicyScope> createScopeResolver(
                 continue;
             }
             try {
-                scopeMap.put(getCodeSource(classUrl, location.representativeClass()), PolicyManager.PolicyScope.server(location.module()));
+                scopeMap.put(getCodeSource(classUrl, location.representativeClass()), PolicyScope.server(location.module()));
             } catch (MalformedURLException e) {
                 throw new IllegalArgumentException("Cannot locate class [" + location.representativeClass() + "]", e);
             }

test/framework/src/main/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManager.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.entitlement.runtime.policy;
 
+import org.elasticsearch.bootstrap.TestScopeResolver;
 import org.elasticsearch.common.util.ArrayUtils;
 import org.elasticsearch.entitlement.runtime.policy.entitlements.Entitlement;
 import org.elasticsearch.test.ESTestCase;
@@ -97,8 +98,8 @@ public final void clearModuleEntitlementsCache() {
 
     @Override
     protected boolean isTrustedSystemClass(Class<?> requestingClass) {
-        if (requestingClass.getPackageName().startsWith("sun.java2d")) {
-            // This is part of the java.desktop module
+        if (TestScopeResolver.getExcludedSystemPackageScope(requestingClass) != null) {
+            // We don't trust the excluded packages even though they are in system modules
             return false;
         }
         ClassLoader loader = requestingClass.getClassLoader();