Clean up PolicyManager and ScopeResolver tests (#127115) (#127238)

* Simplify PolicyManagerTests

* Clean and simplify ScopeResolverTests

Author: prdoyle

libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.entitlement.runtime.policy.entitlements.CreateClassLoaderEntitlement;
 import org.elasticsearch.entitlement.runtime.policy.entitlements.ExitVMEntitlement;
 import org.elasticsearch.entitlement.runtime.policy.entitlements.FilesEntitlement;
+import org.elasticsearch.entitlement.runtime.policy.entitlements.OutboundNetworkEntitlement;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.compiler.InMemoryJavaCompiler;
 import org.elasticsearch.test.jar.JarUtils;
@@ -31,20 +32,17 @@
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.nio.file.Path;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
 import static java.util.Map.entry;
-import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ALL_UNNAMED;
 import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ComponentKind.SERVER;
-import static org.hamcrest.Matchers.aMapWithSize;
 import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.is;
-import static org.hamcrest.Matchers.sameInstance;
 
 @ESTestCase.WithoutSecurityManager
 public class PolicyManagerTests extends ESTestCase {
@@ -89,223 +87,81 @@ public static void beforeClass() {
         }
     }
 
-    public void testGetEntitlementsThrowsOnMissingPluginUnnamedModule() {
-        var plugin1SourcePath = Path.of("modules", "plugin1");
-        var policyManager = new PolicyManager(
-            createEmptyTestServerPolicy(),
-            List.of(),
-            Map.of("plugin1", createPluginPolicy("plugin.module")),
-            c -> PolicyScope.plugin("plugin1", moduleName(c)),
-            Map.of("plugin1", plugin1SourcePath),
-            NO_ENTITLEMENTS_MODULE,
-            TEST_PATH_LOOKUP,
-            Set.of()
-        );
-
-        // Any class from the current module (unnamed) will do
-        var callerClass = this.getClass();
-        var requestingModule = callerClass.getModule();
-
-        assertEquals(
-            "No policy for the unnamed module",
-            policyManager.defaultEntitlements("plugin1", plugin1SourcePath, requestingModule.getName()),
-            policyManager.getEntitlements(callerClass)
-        );
-
-        assertEquals(
-            Map.of(requestingModule, policyManager.defaultEntitlements("plugin1", plugin1SourcePath, requestingModule.getName())),
-            policyManager.moduleEntitlementsMap
-        );
-    }
-
-    public void testGetEntitlementsThrowsOnMissingPolicyForPlugin() {
-        var plugin1SourcePath = Path.of("modules", "plugin1");
-        var policyManager = new PolicyManager(
-            createEmptyTestServerPolicy(),
-            List.of(),
-            Map.of(),
-            c -> PolicyScope.plugin("plugin1", moduleName(c)),
-            Map.of("plugin1", plugin1SourcePath),
-            NO_ENTITLEMENTS_MODULE,
-            TEST_PATH_LOOKUP,
-            Set.of()
-        );
-
-        // Any class from the current module (unnamed) will do
-        var callerClass = this.getClass();
-        var requestingModule = callerClass.getModule();
-
-        assertEquals(
-            "No policy for this plugin",
-            policyManager.defaultEntitlements("plugin1", plugin1SourcePath, requestingModule.getName()),
-            policyManager.getEntitlements(callerClass)
-        );
+    public void testGetEntitlements() {
+        // A mutable policyScope we can use to program specific replies
+        AtomicReference<PolicyScope> policyScope = new AtomicReference<>();
 
-        assertEquals(
-            Map.of(requestingModule, policyManager.defaultEntitlements("plugin1", plugin1SourcePath, requestingModule.getName())),
-            policyManager.moduleEntitlementsMap
-        );
-    }
-
-    public void testGetEntitlementsFailureIsCached() {
+        // A common policy with a variety of entitlements to test
+        Path thisSourcePath = PolicyManager.getComponentPathFromClass(getClass());
         var plugin1SourcePath = Path.of("modules", "plugin1");
         var policyManager = new PolicyManager(
-            createEmptyTestServerPolicy(),
+            new Policy("server", List.of(new Scope("org.example.httpclient", List.of(new OutboundNetworkEntitlement())))),
             List.of(),
-            Map.of(),
-            c -> PolicyScope.plugin("plugin1", moduleName(c)),
+            Map.of("plugin1", new Policy("plugin1", List.of(new Scope("plugin.module1", List.of(new ExitVMEntitlement()))))),
+            c -> policyScope.get(),
             Map.of("plugin1", plugin1SourcePath),
             NO_ENTITLEMENTS_MODULE,
             TEST_PATH_LOOKUP,
             Set.of()
         );
 
-        // Any class from the current module (unnamed) will do
-        var callerClass = this.getClass();
-        var requestingModule = callerClass.getModule();
+        // "Unspecified" below means that the module is not named in the policy
 
-        assertEquals(
-            policyManager.defaultEntitlements("plugin1", plugin1SourcePath, requestingModule.getName()),
-            policyManager.getEntitlements(callerClass)
-        );
-        assertEquals(
-            Map.of(requestingModule, policyManager.defaultEntitlements("plugin1", plugin1SourcePath, requestingModule.getName())),
-            policyManager.moduleEntitlementsMap
+        policyScope.set(PolicyScope.server("org.example.httpclient"));
+        resetAndCheckEntitlements(
+            "Specified entitlements for server",
+            getClass(),
+            policyManager.policyEntitlements(
+                SERVER.componentName,
+                thisSourcePath,
+                "org.example.httpclient",
+                List.of(new OutboundNetworkEntitlement())
+            ),
+            policyManager
         );
 
-        // A second time
-        assertEquals(
-            policyManager.defaultEntitlements("plugin1", plugin1SourcePath, requestingModule.getName()),
-            policyManager.getEntitlements(callerClass)
+        policyScope.set(PolicyScope.server("plugin.unspecifiedModule"));
+        resetAndCheckEntitlements(
+            "Default entitlements for unspecified module",
+            getClass(),
+            policyManager.defaultEntitlements(SERVER.componentName, thisSourcePath, "plugin.unspecifiedModule"),
+            policyManager
         );
 
-        // Nothing new in the map
-        assertEquals(
-            Map.of(requestingModule, policyManager.defaultEntitlements("plugin1", plugin1SourcePath, requestingModule.getName())),
-            policyManager.moduleEntitlementsMap
+        policyScope.set(PolicyScope.plugin("plugin1", "plugin.module1"));
+        resetAndCheckEntitlements(
+            "Specified entitlements for plugin",
+            getClass(),
+            policyManager.policyEntitlements("plugin1", plugin1SourcePath, "plugin.module1", List.of(new ExitVMEntitlement())),
+            policyManager
         );
-    }
 
-    public void testGetEntitlementsReturnsEntitlementsForPluginUnnamedModule() {
-        var policyManager = new PolicyManager(
-            createEmptyTestServerPolicy(),
-            List.of(),
-            Map.ofEntries(entry("plugin2", createPluginPolicy(ALL_UNNAMED))),
-            c -> PolicyScope.plugin("plugin2", moduleName(c)),
-            Map.of("plugin2", Path.of("modules", "plugin2")),
-            NO_ENTITLEMENTS_MODULE,
-            TEST_PATH_LOOKUP,
-            Set.of()
+        policyScope.set(PolicyScope.plugin("plugin1", "plugin.unspecifiedModule"));
+        resetAndCheckEntitlements(
+            "Default entitlements for plugin",
+            getClass(),
+            policyManager.defaultEntitlements("plugin1", plugin1SourcePath, "plugin.unspecifiedModule"),
+            policyManager
         );
-
-        // Any class from the current module (unnamed) will do
-        var callerClass = this.getClass();
-
-        var entitlements = policyManager.getEntitlements(callerClass);
-        assertThat(entitlements.hasEntitlement(CreateClassLoaderEntitlement.class), is(true));
     }
 
-    public void testGetEntitlementsReturnsDefaultOnMissingPolicyForServer() throws ClassNotFoundException {
-        var policyManager = new PolicyManager(
-            createTestServerPolicy("example"),
-            List.of(),
-            Map.of(),
-            c -> PolicyScope.server(moduleName(c)),
-            Map.of(),
-            NO_ENTITLEMENTS_MODULE,
-            TEST_PATH_LOOKUP,
-            Set.of()
-        );
-
-        // Any class will do, since our resolver is hardcoded to use SERVER_COMPONENT_NAME.
-        // Let's pick one with a known module name.
-        String httpserverModuleName = "jdk.httpserver";
-        var mockServerClass = ModuleLayer.boot().findLoader(httpserverModuleName).loadClass("com.sun.net.httpserver.HttpServer");
-        var mockServerSourcePath = PolicyManager.getComponentPathFromClass(mockServerClass);
-        var requestingModule = mockServerClass.getModule();
-
+    private void resetAndCheckEntitlements(
+        String message,
+        Class<?> requestingClass,
+        ModuleEntitlements expectedEntitlements,
+        PolicyManager policyManager
+    ) {
+        policyManager.moduleEntitlementsMap.clear();
+        assertEquals(message, expectedEntitlements, policyManager.getEntitlements(requestingClass));
         assertEquals(
-            "No policy for this module in server",
-            policyManager.defaultEntitlements(SERVER.componentName, mockServerSourcePath, httpserverModuleName),
-            policyManager.getEntitlements(mockServerClass)
-        );
-
-        assertEquals(
-            Map.of(requestingModule, policyManager.defaultEntitlements(SERVER.componentName, mockServerSourcePath, httpserverModuleName)),
+            "Map has precisely the one expected entry",
+            Map.of(requestingClass.getModule(), expectedEntitlements),
             policyManager.moduleEntitlementsMap
         );
-    }
 
-    public void testGetEntitlementsReturnsEntitlementsForServerModule() throws ClassNotFoundException {
-        String httpserverModuleName = "jdk.httpserver";
-        var policyManager = new PolicyManager(
-            createTestServerPolicy(httpserverModuleName),
-            List.of(),
-            Map.of(),
-            c -> PolicyScope.server(moduleName(c)),
-            Map.of(),
-            NO_ENTITLEMENTS_MODULE,
-            TEST_PATH_LOOKUP,
-            Set.of()
-        );
-
-        // Any class will do, since our resolver is hardcoded to use SERVER_COMPONENT_NAME.
-        // Let's pick one with a known module name.
-        var mockServerClass = ModuleLayer.boot().findLoader(httpserverModuleName).loadClass("com.sun.net.httpserver.HttpServer");
-
-        var entitlements = policyManager.getEntitlements(mockServerClass);
-        assertThat(entitlements.hasEntitlement(CreateClassLoaderEntitlement.class), is(true));
-        assertThat(entitlements.hasEntitlement(ExitVMEntitlement.class), is(true));
-    }
-
-    public void testGetEntitlementsReturnsEntitlementsForPluginModule() throws IOException, ClassNotFoundException {
-        final Path home = createTempDir();
-
-        Path jar = createMockPluginJar(home);
-
-        var policyManager = new PolicyManager(
-            createEmptyTestServerPolicy(),
-            List.of(),
-            Map.of("mock-plugin", createPluginPolicy("org.example.plugin")),
-            c -> PolicyScope.plugin("mock-plugin", moduleName(c)),
-            Map.of("mock-plugin", Path.of("modules", "mock-plugin")),
-            NO_ENTITLEMENTS_MODULE,
-            TEST_PATH_LOOKUP,
-            Set.of()
-        );
-
-        var layer = createLayerForJar(jar, "org.example.plugin");
-        var mockPluginClass = layer.findLoader("org.example.plugin").loadClass("q.B");
-
-        var entitlements = policyManager.getEntitlements(mockPluginClass);
-        assertThat(entitlements.hasEntitlement(CreateClassLoaderEntitlement.class), is(true));
-        assertThat(entitlements.fileAccess().canRead(TEST_BASE_DIR), is(true));
-    }
-
-    public void testGetEntitlementsResultIsCached() {
-        var policyManager = new PolicyManager(
-            createEmptyTestServerPolicy(),
-            List.of(),
-            Map.ofEntries(entry("plugin2", createPluginPolicy(ALL_UNNAMED))),
-            c -> PolicyScope.plugin("plugin2", moduleName(c)),
-            Map.of("plugin2", Path.of("modules", "plugin2")),
-            NO_ENTITLEMENTS_MODULE,
-            TEST_PATH_LOOKUP,
-            Set.of()
-        );
-
-        // Any class from the current module (unnamed) will do
-        var callerClass = this.getClass();
-
-        var entitlements = policyManager.getEntitlements(callerClass);
-        assertThat(entitlements.hasEntitlement(CreateClassLoaderEntitlement.class), is(true));
-        assertThat(policyManager.moduleEntitlementsMap, aMapWithSize(1));
-        var cachedResult = policyManager.moduleEntitlementsMap.values().stream().findFirst().orElseThrow();
-        var entitlementsAgain = policyManager.getEntitlements(callerClass);
-
-        // Nothing new in the map
-        assertThat(policyManager.moduleEntitlementsMap, aMapWithSize(1));
-        assertThat(entitlementsAgain, sameInstance(cachedResult));
+        // Fetch a second time and verify the map is unchanged
+        policyManager.getEntitlements(requestingClass);
+        assertEquals("Map is unchanged", Map.of(requestingClass.getModule(), expectedEntitlements), policyManager.moduleEntitlementsMap);
     }
 
     public void testRequestingClassFastPath() throws IOException, ClassNotFoundException {
@@ -560,24 +416,6 @@ public void testFilesEntitlementsWithExclusive() {
         );
     }
 
-    /**
-     * If the plugin resolver tells us a class is in a plugin, don't conclude that it's in an agent.
-     */
-    public void testPluginResolverOverridesAgents() {
-        var policyManager = new PolicyManager(
-            createEmptyTestServerPolicy(),
-            List.of(new CreateClassLoaderEntitlement()),
-            Map.of(),
-            c -> PolicyScope.plugin("test", moduleName(c)), // Insist that the class is in a plugin
-            Map.of(),
-            NO_ENTITLEMENTS_MODULE,
-            TEST_PATH_LOOKUP,
-            Set.of()
-        );
-        ModuleEntitlements notAgentsEntitlements = policyManager.getEntitlements(TestAgent.class);
-        assertThat(notAgentsEntitlements.hasEntitlement(CreateClassLoaderEntitlement.class), is(false));
-    }
-
     private static Class<?> makeClassInItsOwnModule() throws IOException, ClassNotFoundException {
         final Path home = createTempDir();
         Path jar = createMockPluginJar(home);
@@ -602,27 +440,6 @@ private static Policy createEmptyTestServerPolicy() {
         return new Policy("server", List.of());
     }
 
-    private static Policy createTestServerPolicy(String scopeName) {
-        return new Policy("server", List.of(new Scope(scopeName, List.of(new ExitVMEntitlement(), new CreateClassLoaderEntitlement()))));
-    }
-
-    private static Policy createPluginPolicy(String... pluginModules) {
-        return new Policy(
-            "plugin",
-            Arrays.stream(pluginModules)
-                .map(
-                    name -> new Scope(
-                        name,
-                        List.of(
-                            new FilesEntitlement(List.of(FilesEntitlement.FileData.ofPath(TEST_BASE_DIR, FilesEntitlement.Mode.READ))),
-                            new CreateClassLoaderEntitlement()
-                        )
-                    )
-                )
-                .toList()
-        );
-    }
-
     private static Path createMockPluginJarForUnnamedModule(Path home) throws IOException {
         Path jar = home.resolve("unnamed-mock-plugin.jar");