Add TestEntitlementsRule with support for dynamic entitled node paths for testing (#132077) (#132636)

Use TestEntitlementsRule and closable entitled node paths for better isolation and management of entitlement related state during tests.

However, wipePendingDataDirectories in tests requires entitlement delegation to work as this uses server's FileSystemUtils. Until ES-10920 is solved, entitled node paths cannot be removed unless the test suite completes or explicitly removing all entitled node paths in-between tests (if cluster scope is TEST).

Relates to ES-12042

(cherry picked from commit 76dac08)

Author: mosche

server/src/internalClusterTest/java/org/elasticsearch/discovery/single/SingleNodeDiscoveryIT.java

@@ -87,7 +87,8 @@ public Path nodeConfigPath(int nodeOrdinal) {
             0,
             "other",
             Arrays.asList(getTestTransportPlugin(), MockHttpTransport.TestPlugin.class),
-            Function.identity()
+            Function.identity(),
+            TEST_ENTITLEMENTS::addEntitledNodePaths
         );
         try {
             other.beforeTest(random());
@@ -137,7 +138,8 @@ public Path nodeConfigPath(int nodeOrdinal) {
             0,
             "other",
             Arrays.asList(getTestTransportPlugin(), MockHttpTransport.TestPlugin.class),
-            Function.identity()
+            Function.identity(),
+            TEST_ENTITLEMENTS::addEntitledNodePaths
         );
         try (var mockLog = MockLog.capture(JoinHelper.class)) {
             mockLog.addExpectation(

server/src/internalClusterTest/java/org/elasticsearch/snapshots/MultiClusterRepoAccessIT.java

@@ -77,7 +77,8 @@ public Path nodeConfigPath(int nodeOrdinal) {
                 InternalSettingsPlugin.class,
                 getTestTransportPlugin()
             ),
-            Function.identity()
+            Function.identity(),
+            TEST_ENTITLEMENTS::addEntitledNodePaths
         );
         secondCluster.beforeTest(random());
     }

server/src/test/java/org/elasticsearch/bootstrap/EntitlementMetaTests.java

@@ -12,7 +12,6 @@
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.entitlement.bootstrap.TestEntitlementBootstrap;
 import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.test.ESTestCase.WithEntitlementsOnTestCode;
 
 import java.io.IOException;
 import java.nio.file.Path;
@@ -42,7 +41,7 @@
  */
 public class EntitlementMetaTests extends ESTestCase {
     public void testSelfTestPasses() {
-        assumeTrue("Not yet working in serverless", TestEntitlementBootstrap.isEnabledForTest());
+        assumeTrue("Not yet working in serverless", TestEntitlementBootstrap.isEnabledForTests());
         Elasticsearch.entitlementSelfTest();
     }
 

server/src/test/java/org/elasticsearch/bootstrap/WithEntitlementsOnTestCodeMetaTests.java

@@ -30,13 +30,13 @@ public class WithEntitlementsOnTestCodeMetaTests extends ESTestCase {
      * is called from server code. The self-test should pass as usual.
      */
     public void testSelfTestPasses() {
-        assumeTrue("Not yet working in serverless", TestEntitlementBootstrap.isEnabledForTest());
+        assumeTrue("Not yet working in serverless", TestEntitlementBootstrap.isEnabledForTests());
         Elasticsearch.entitlementSelfTest();
     }
 
     @SuppressForbidden(reason = "Testing that a forbidden API is disallowed")
     public void testForbiddenActionDenied() {
-        assumeTrue("Not yet working in serverless", TestEntitlementBootstrap.isEnabledForTest());
+        assumeTrue("Not yet working in serverless", TestEntitlementBootstrap.isEnabledForTests());
         assertThrows(NotEntitledException.class, () -> Path.of(".").toRealPath());
     }
 }

server/src/test/java/org/elasticsearch/indices/IndicesServiceCloseTests.java

@@ -77,7 +77,8 @@ private Node startNode() throws NodeValidationException {
         Node node = new MockNode(
             settings,
             Arrays.asList(getTestTransportPlugin(), MockHttpTransport.TestPlugin.class, InternalSettingsPlugin.class),
-            true
+            true,
+            () -> {}
         );
         node.start();
         return node;

test/framework/src/main/java/org/elasticsearch/entitlement/bootstrap/TestEntitlementBootstrap.java

@@ -9,22 +9,18 @@
 
 package org.elasticsearch.entitlement.bootstrap;
 
-import org.apache.lucene.tests.mockfile.FilterPath;
 import org.elasticsearch.bootstrap.TestBuildInfo;
 import org.elasticsearch.bootstrap.TestBuildInfoParser;
 import org.elasticsearch.bootstrap.TestScopeResolver;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.Booleans;
-import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.PathUtils;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.entitlement.initialization.EntitlementInitialization;
 import org.elasticsearch.entitlement.runtime.policy.PathLookup;
-import org.elasticsearch.entitlement.runtime.policy.PathLookup.BaseDir;
 import org.elasticsearch.entitlement.runtime.policy.Policy;
+import org.elasticsearch.entitlement.runtime.policy.PolicyManager;
 import org.elasticsearch.entitlement.runtime.policy.PolicyParser;
-import org.elasticsearch.entitlement.runtime.policy.TestPathLookup;
 import org.elasticsearch.entitlement.runtime.policy.TestPolicyManager;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
@@ -35,175 +31,52 @@
 import java.net.URI;
 import java.net.URL;
 import java.nio.file.Path;
-import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeSet;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.function.BiFunction;
-import java.util.function.Consumer;
 
 import static java.util.stream.Collectors.toCollection;
 import static java.util.stream.Collectors.toSet;
-import static org.elasticsearch.entitlement.runtime.policy.PathLookup.BaseDir.TEMP;
-import static org.elasticsearch.env.Environment.PATH_DATA_SETTING;
-import static org.elasticsearch.env.Environment.PATH_HOME_SETTING;
-import static org.elasticsearch.env.Environment.PATH_REPO_SETTING;
-import static org.elasticsearch.env.Environment.PATH_SHARED_DATA_SETTING;
 
 public class TestEntitlementBootstrap {
-
     private static final Logger logger = LogManager.getLogger(TestEntitlementBootstrap.class);
 
-    private static Map<BaseDir, Collection<Path>> baseDirPaths = new ConcurrentHashMap<>();
-    private static TestPolicyManager policyManager;
+    private static TestPathLookup TEST_PATH_LOOKUP;
+    private static TestPolicyManager POLICY_MANAGER;
 
     /**
      * Activates entitlement checking in tests.
      */
-    public static void bootstrap(@Nullable Path tempDir) throws IOException {
-        if (isEnabledForTest() == false) {
-            return;
-        }
-        var previousTempDir = baseDirPaths.put(TEMP, zeroOrOne(tempDir));
-        assert previousTempDir == null : "Test entitlement bootstrap called multiple times";
-        TestPathLookup pathLookup = new TestPathLookup(baseDirPaths);
-        policyManager = createPolicyManager(pathLookup);
-        EntitlementInitialization.initializeArgs = new EntitlementInitialization.InitializeArgs(pathLookup, Set.of(), policyManager);
-        logger.debug("Loading entitlement agent");
-        EntitlementBootstrap.loadAgent(EntitlementBootstrap.findAgentJar(), EntitlementInitialization.class.getName());
-    }
-
-    public static void registerNodeBaseDirs(Settings settings, Path configPath) {
-        if (policyManager == null) {
-            return;
-        }
-
-        Path homeDir = homeDir(settings);
-        Path configDir = configDir(configPath, homeDir);
-        Collection<Path> dataDirs = dataDirs(settings, homeDir);
-        Collection<Path> sharedDataDir = sharedDataDir(settings);
-        Collection<Path> repoDirs = repoDirs(settings);
-        logger.debug(
-            "Registering node dirs: config [{}], dataDirs [{}], sharedDataDir [{}], repoDirs [{}]",
-            configDir,
-            dataDirs,
-            sharedDataDir,
-            repoDirs
-        );
-        baseDirPaths.compute(BaseDir.CONFIG, baseDirModifier(paths -> paths.add(configDir)));
-        baseDirPaths.compute(BaseDir.DATA, baseDirModifier(paths -> paths.addAll(dataDirs)));
-        baseDirPaths.compute(BaseDir.SHARED_DATA, baseDirModifier(paths -> paths.addAll(sharedDataDir)));
-        baseDirPaths.compute(BaseDir.SHARED_REPO, baseDirModifier(paths -> paths.addAll(repoDirs)));
-        policyManager.clearModuleEntitlementsCache();
-    }
-
-    public static void unregisterNodeBaseDirs(Settings settings, Path configPath) {
-        if (policyManager == null) {
+    public static void bootstrap(Path tempDir) throws IOException {
+        if (isEnabledForTests() == false) {
             return;
         }
-
-        Path homeDir = homeDir(settings);
-        Path configDir = configDir(configPath, homeDir);
-        Collection<Path> dataDirs = dataDirs(settings, homeDir);
-        Collection<Path> sharedDataDir = sharedDataDir(settings);
-        Collection<Path> repoDirs = repoDirs(settings);
-        logger.debug(
-            "Unregistering node dirs: config [{}], dataDirs [{}], sharedDataDir [{}], repoDirs [{}]",
-            configDir,
-            dataDirs,
-            sharedDataDir,
-            repoDirs
-        );
-        baseDirPaths.compute(BaseDir.CONFIG, baseDirModifier(paths -> paths.remove(configDir)));
-        baseDirPaths.compute(BaseDir.DATA, baseDirModifier(paths -> paths.removeAll(dataDirs)));
-        baseDirPaths.compute(BaseDir.SHARED_DATA, baseDirModifier(paths -> paths.removeAll(sharedDataDir)));
-        baseDirPaths.compute(BaseDir.SHARED_REPO, baseDirModifier(paths -> paths.removeAll(repoDirs)));
-        policyManager.clearModuleEntitlementsCache();
-    }
-
-    private static Path homeDir(Settings settings) {
-        return absolutePath(PATH_HOME_SETTING.get(settings));
-    }
-
-    private static Path configDir(Path configDir, Path homeDir) {
-        return configDir != null ? unwrapFilterPath(configDir) : homeDir.resolve("config");
-    }
-
-    private static Collection<Path> dataDirs(Settings settings, Path homeDir) {
-        List<String> dataDirs = PATH_DATA_SETTING.get(settings);
-        return dataDirs.isEmpty()
-            ? List.of(homeDir.resolve("data"))
-            : dataDirs.stream().map(TestEntitlementBootstrap::absolutePath).toList();
+        assert POLICY_MANAGER == null && TEST_PATH_LOOKUP == null : "Test entitlement bootstrap called multiple times";
+        TEST_PATH_LOOKUP = new TestPathLookup(tempDir);
+        POLICY_MANAGER = createPolicyManager(TEST_PATH_LOOKUP);
+        loadAgent(POLICY_MANAGER, TEST_PATH_LOOKUP);
     }
 
-    private static Collection<Path> sharedDataDir(Settings settings) {
-        String sharedDataDir = PATH_SHARED_DATA_SETTING.get(settings);
-        return Strings.hasText(sharedDataDir) ? List.of(absolutePath(sharedDataDir)) : List.of();
-    }
-
-    private static Collection<Path> repoDirs(Settings settings) {
-        return PATH_REPO_SETTING.get(settings).stream().map(TestEntitlementBootstrap::absolutePath).toList();
-    }
-
-    private static BiFunction<BaseDir, Collection<Path>, Collection<Path>> baseDirModifier(Consumer<Collection<Path>> consumer) {
-        // always return a new unmodifiable copy
-        return (BaseDir baseDir, Collection<Path> paths) -> {
-            paths = paths == null ? new HashSet<>() : new HashSet<>(paths);
-            consumer.accept(paths);
-            return Collections.unmodifiableCollection(paths);
-        };
-    }
-
-    private static Path unwrapFilterPath(Path path) {
-        while (path instanceof FilterPath fPath) {
-            path = fPath.getDelegate();
-        }
-        return path;
-    }
-
-    @SuppressForbidden(reason = "must be resolved using the default file system, rather then the mocked test file system")
-    private static Path absolutePath(String path) {
-        return Paths.get(path).toAbsolutePath().normalize();
-    }
-
-    private static <T> List<T> zeroOrOne(T item) {
-        if (item == null) {
-            return List.of();
-        } else {
-            return List.of(item);
-        }
-    }
-
-    public static boolean isEnabledForTest() {
+    public static boolean isEnabledForTests() {
         return Booleans.parseBoolean(System.getProperty("es.entitlement.enableForTests", "false"));
     }
 
-    public static void setActive(boolean newValue) {
-        policyManager.setActive(newValue);
-    }
-
-    public static void setTriviallyAllowingTestCode(boolean newValue) {
-        policyManager.setTriviallyAllowingTestCode(newValue);
+    static TestPolicyManager testPolicyManager() {
+        return POLICY_MANAGER;
     }
 
-    public static void setEntitledTestPackages(String[] entitledTestPackages) {
-        policyManager.setEntitledTestPackages(entitledTestPackages);
+    static TestPathLookup testPathLookup() {
+        return TEST_PATH_LOOKUP;
     }
 
-    public static void resetAfterTest() {
-        // reset all base dirs except TEMP, which is initialized just once statically
-        baseDirPaths.keySet().retainAll(List.of(TEMP));
-        if (policyManager != null) {
-            policyManager.resetAfterTest();
-        }
+    private static void loadAgent(PolicyManager policyManager, PathLookup pathLookup) {
+        logger.debug("Loading entitlement agent");
+        EntitlementInitialization.initializeArgs = new EntitlementInitialization.InitializeArgs(pathLookup, Set.of(), policyManager);
+        EntitlementBootstrap.loadAgent(EntitlementBootstrap.findAgentJar(), EntitlementInitialization.class.getName());
     }
 
     private static TestPolicyManager createPolicyManager(PathLookup pathLookup) throws IOException {
@@ -224,7 +97,7 @@ private static TestPolicyManager createPolicyManager(PathLookup pathLookup) thro
 
         String separator = System.getProperty("path.separator");
 
-        // In productions, plugins would have access to their respective bundle directories,
+        // In production, plugins would have access to their respective bundle directories,
         // and so they'd be able to read from their jars. In testing, we approximate this
         // by considering the entire classpath to be "source paths" of all plugins. This
         // also has the effect of granting read access to everything on the test-only classpath,