Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 11, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,8 @@ public Path nodeConfigPath(int nodeOrdinal) {
InternalSettingsPlugin.class,
getTestTransportPlugin()
),
Function.identity()
Function.identity(),
TEST_ENTITLEMENTS::addEntitledNodePaths
);
secondCluster.beforeTest(random());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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 {
Expand All @@ -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,
Expand Down
Loading