From b7d91197bd705c1aa0b34bfd5d77d8a1135a62eb Mon Sep 17 00:00:00 2001
From: Patrick Doyle
Date: Tue, 6 May 2025 13:55:40 -0400
Subject: [PATCH 1/3] Split out EntitlementsCache
---
.../EntitlementInitialization.java | 3 +-
.../initialization/EntitlementsCacheImpl.java | 28 +++++++
.../runtime/policy/EntitlementsCache.java | 76 +++++++++++++++++++
.../runtime/policy/PolicyManager.java | 62 ++++-----------
.../policy/EntitlementsCacheForTesting.java | 24 ++++++
.../runtime/policy/PolicyManagerTests.java | 38 +++++++---
6 files changed, 169 insertions(+), 62 deletions(-)
create mode 100644 libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementsCacheImpl.java
create mode 100644 libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/EntitlementsCache.java
create mode 100644 libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/EntitlementsCacheForTesting.java
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java
index 53bfb5c57b23c..82fbd2dfd6c1c 100644
--- a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java
+++ b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java
@@ -86,7 +86,8 @@ private static PolicyManager createPolicyManager() {
EntitlementBootstrap.bootstrapArgs().sourcePaths(),
ENTITLEMENTS_MODULE,
pathLookup,
- bootstrapArgs.suppressFailureLogClasses()
+ bootstrapArgs.suppressFailureLogClasses(),
+ new EntitlementsCacheImpl()
);
}
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementsCacheImpl.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementsCacheImpl.java
new file mode 100644
index 0000000000000..a56dae88af4ce
--- /dev/null
+++ b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementsCacheImpl.java
@@ -0,0 +1,28 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.entitlement.initialization;
+
+import org.elasticsearch.entitlement.runtime.policy.EntitlementsCache;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * The production {@link EntitlementsCache}. (Tests use {@code EntitlementCacheForTesting}.)
+ */
+final class EntitlementsCacheImpl extends ConcurrentHashMap implements EntitlementsCache {
+ @Override
+ public ModuleEntitlements computeIfAbsent(Class> key, Function super Class>, ? extends ModuleEntitlements> mappingFunction) {
+ // We cache per module rather than per class to make the cache smaller and increase the hit ratio
+ return computeIfAbsent(key.getModule(), m -> requireNonNull(mappingFunction.apply(key)));
+ }
+}
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/EntitlementsCache.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/EntitlementsCache.java
new file mode 100644
index 0000000000000..e3eaa310d32ab
--- /dev/null
+++ b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/EntitlementsCache.java
@@ -0,0 +1,76 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.entitlement.runtime.policy;
+
+import org.elasticsearch.entitlement.runtime.policy.entitlements.Entitlement;
+import org.elasticsearch.logging.Logger;
+
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Stream;
+
+/**
+ * Quickly provides entitlement info for a given class. Performance-critical.
+ */
+public interface EntitlementsCache {
+ /**
+ * @return true if entitlement checking should be bypassed for the given class;
+ * false if the normal built-in "trivially allowed" rules apply.
+ */
+ default boolean isAlwaysAllowed(Class> key) {
+ return false;
+ }
+
+ ModuleEntitlements computeIfAbsent(Class> key, Function super Class>, ? extends ModuleEntitlements> mappingFunction);
+
+ /**
+ * This class contains all the entitlements by type, plus the {@link FileAccessTree} for the special case of filesystem entitlements.
+ *
+ * We use layers when computing {@link ModuleEntitlements}; first, we check whether the module we are building it for is in the
+ * server layer ({@link PolicyManager#SERVER_LAYER_MODULES}) (*).
+ * If it is, we use the server policy, using the same caller class module name as the scope, and read the entitlements for that scope.
+ * Otherwise, we use the {@code PluginResolver} to identify the correct plugin layer and find the policy for it (if any).
+ * If the plugin is modular, we again use the same caller class module name as the scope, and read the entitlements for that scope.
+ * If it's not, we use the single {@code ALL-UNNAMED} scope – in this case there is one scope and all entitlements apply
+ * to all the plugin code.
+ *
+ *
+ * (*) implementation detail: this is currently done in an indirect way: we know the module is not in the system layer
+ * (otherwise the check would have been already trivially allowed), so we just check that the module is named, and it belongs to the
+ * boot {@link ModuleLayer}. We might want to change this in the future to make it more consistent/easier to maintain.
+ *
+ *
+ * @param componentName the plugin name or else one of the special component names like "(server)".
+ */
+ record ModuleEntitlements(
+ String componentName,
+ Map, List> entitlementsByType,
+ FileAccessTree fileAccess,
+ Logger logger
+ ) {
+
+ public ModuleEntitlements {
+ entitlementsByType = Map.copyOf(entitlementsByType);
+ }
+
+ public boolean hasEntitlement(Class extends Entitlement> entitlementClass) {
+ return entitlementsByType.containsKey(entitlementClass);
+ }
+
+ public Stream getEntitlements(Class entitlementClass) {
+ var entitlements = entitlementsByType.get(entitlementClass);
+ if (entitlements == null) {
+ return Stream.empty();
+ }
+ return entitlements.stream().map(entitlementClass::cast);
+ }
+ }
+}
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java
index aa901e9900d17..08bbd3ba8598f 100644
--- a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java
+++ b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java
@@ -14,6 +14,7 @@
import org.elasticsearch.core.SuppressForbidden;
import org.elasticsearch.entitlement.instrumentation.InstrumentationService;
import org.elasticsearch.entitlement.runtime.api.NotEntitledException;
+import org.elasticsearch.entitlement.runtime.policy.EntitlementsCache.ModuleEntitlements;
import org.elasticsearch.entitlement.runtime.policy.FileAccessTree.ExclusiveFileEntitlement;
import org.elasticsearch.entitlement.runtime.policy.FileAccessTree.ExclusivePath;
import org.elasticsearch.entitlement.runtime.policy.entitlements.CreateClassLoaderEntitlement;
@@ -121,7 +122,7 @@
* All these methods start in the same way: the components identified in the previous section are used to establish if and how to check:
* If the caller class belongs to {@link PolicyManager#SYSTEM_LAYER_MODULES}, no check is performed (the call is trivially allowed, see
* {@link PolicyManager#isTriviallyAllowed}).
- * Otherwise, we lazily compute and create a {@link PolicyManager.ModuleEntitlements} record (see
+ * Otherwise, we lazily compute and create a {@link ModuleEntitlements} record (see
* {@link PolicyManager#computeEntitlements}). The record is cached so it can be used in following checks, stored in a
* {@code Module -> ModuleEntitlement} map.
*
@@ -181,49 +182,6 @@ public enum ComponentKind {
}
}
- /**
- * This class contains all the entitlements by type, plus the {@link FileAccessTree} for the special case of filesystem entitlements.
- *
- * We use layers when computing {@link ModuleEntitlements}; first, we check whether the module we are building it for is in the
- * server layer ({@link PolicyManager#SERVER_LAYER_MODULES}) (*).
- * If it is, we use the server policy, using the same caller class module name as the scope, and read the entitlements for that scope.
- * Otherwise, we use the {@code PluginResolver} to identify the correct plugin layer and find the policy for it (if any).
- * If the plugin is modular, we again use the same caller class module name as the scope, and read the entitlements for that scope.
- * If it's not, we use the single {@code ALL-UNNAMED} scope – in this case there is one scope and all entitlements apply
- * to all the plugin code.
- *
- *
- * (*) implementation detail: this is currently done in an indirect way: we know the module is not in the system layer
- * (otherwise the check would have been already trivially allowed), so we just check that the module is named, and it belongs to the
- * boot {@link ModuleLayer}. We might want to change this in the future to make it more consistent/easier to maintain.
- *
- *
- * @param componentName the plugin name or else one of the special component names like "(server)".
- */
- record ModuleEntitlements(
- String componentName,
- Map, List> entitlementsByType,
- FileAccessTree fileAccess,
- Logger logger
- ) {
-
- ModuleEntitlements {
- entitlementsByType = Map.copyOf(entitlementsByType);
- }
-
- public boolean hasEntitlement(Class extends Entitlement> entitlementClass) {
- return entitlementsByType.containsKey(entitlementClass);
- }
-
- public Stream getEntitlements(Class entitlementClass) {
- var entitlements = entitlementsByType.get(entitlementClass);
- if (entitlements == null) {
- return Stream.empty();
- }
- return entitlements.stream().map(entitlementClass::cast);
- }
- }
-
private FileAccessTree getDefaultFileAccess(Path componentPath) {
return FileAccessTree.withoutExclusivePaths(FilesEntitlement.EMPTY, pathLookup, componentPath);
}
@@ -249,8 +207,6 @@ ModuleEntitlements policyEntitlements(String componentName, Path componentPath,
);
}
- final Map moduleEntitlementsMap = new ConcurrentHashMap<>();
-
private final Map> serverEntitlements;
private final List apmAgentEntitlements;
private final Map>> pluginsEntitlements;
@@ -303,6 +259,8 @@ private static Set findSystemLayerModules() {
*/
private final List exclusivePaths;
+ final EntitlementsCache moduleEntitlementsCache;
+
public PolicyManager(
Policy serverPolicy,
List apmAgentEntitlements,
@@ -311,7 +269,8 @@ public PolicyManager(
Map sourcePaths,
Module entitlementsModule,
PathLookup pathLookup,
- Set> suppressFailureLogClasses
+ Set> suppressFailureLogClasses,
+ EntitlementsCache moduleEntitlementsCache
) {
this.serverEntitlements = buildScopeEntitlementsMap(requireNonNull(serverPolicy));
this.apmAgentEntitlements = apmAgentEntitlements;
@@ -323,6 +282,7 @@ public PolicyManager(
this.entitlementsModule = entitlementsModule;
this.pathLookup = requireNonNull(pathLookup);
this.mutedClasses = suppressFailureLogClasses;
+ this.moduleEntitlementsCache = moduleEntitlementsCache;
List exclusiveFileEntitlements = new ArrayList<>();
for (var e : serverEntitlements.entrySet()) {
@@ -723,7 +683,7 @@ private void checkEntitlementPresent(Class> callerClass, Class extends Entit
}
ModuleEntitlements getEntitlements(Class> requestingClass) {
- return moduleEntitlementsMap.computeIfAbsent(requestingClass.getModule(), m -> computeEntitlements(requestingClass));
+ return moduleEntitlementsCache.computeIfAbsent(requestingClass, this::computeEntitlements);
}
private ModuleEntitlements computeEntitlements(Class> requestingClass) {
@@ -827,7 +787,7 @@ Optional findRequestingFrame(Stream frames) {
/**
* @return true if permission is granted regardless of the entitlement
*/
- private static boolean isTriviallyAllowed(Class> requestingClass) {
+ private boolean isTriviallyAllowed(Class> requestingClass) {
if (generalLogger.isTraceEnabled()) {
generalLogger.trace("Stack trace for upcoming trivially-allowed check", new Exception());
}
@@ -843,6 +803,10 @@ private static boolean isTriviallyAllowed(Class> requestingClass) {
generalLogger.debug("Entitlement trivially allowed from system module [{}]", requestingClass.getModule().getName());
return true;
}
+ if (moduleEntitlementsCache.isAlwaysAllowed(requestingClass)) {
+ generalLogger.debug("Entitlement trivially allowed by the EntitlementsCache");
+ return true;
+ }
generalLogger.trace("Entitlement not trivially allowed");
return false;
}
diff --git a/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/EntitlementsCacheForTesting.java b/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/EntitlementsCacheForTesting.java
new file mode 100644
index 0000000000000..75173fadf9fff
--- /dev/null
+++ b/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/EntitlementsCacheForTesting.java
@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.entitlement.runtime.policy;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * When testing, we don't use modules, so we cache per-class.
+ */
+public class EntitlementsCacheForTesting extends ConcurrentHashMap, EntitlementsCache.ModuleEntitlements>
+ implements
+ EntitlementsCache {
+ @Override
+ public boolean isAlwaysAllowed(Class> key) {
+ return key.getPackageName().startsWith("org.gradle") || key.getPackageName().startsWith("org.junit");
+ }
+}
diff --git a/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java b/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java
index dcb05e8983f05..36e943596272e 100644
--- a/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java
+++ b/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java
@@ -12,7 +12,7 @@
import org.elasticsearch.bootstrap.ScopeResolver;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.core.Strings;
-import org.elasticsearch.entitlement.runtime.policy.PolicyManager.ModuleEntitlements;
+import org.elasticsearch.entitlement.runtime.policy.EntitlementsCache.ModuleEntitlements;
import org.elasticsearch.entitlement.runtime.policy.PolicyManager.PolicyScope;
import org.elasticsearch.entitlement.runtime.policy.agent.TestAgent;
import org.elasticsearch.entitlement.runtime.policy.agent.inner.TestInnerAgent;
@@ -93,6 +93,7 @@ public void testGetEntitlements() {
// A common policy with a variety of entitlements to test
Path thisSourcePath = PolicyManager.getComponentPathFromClass(getClass());
var plugin1SourcePath = Path.of("modules", "plugin1");
+ var cache = new EntitlementsCacheForTesting();
var policyManager = new PolicyManager(
new Policy("server", List.of(new Scope("org.example.httpclient", List.of(new OutboundNetworkEntitlement())))),
List.of(),
@@ -101,7 +102,8 @@ public void testGetEntitlements() {
Map.of("plugin1", plugin1SourcePath),
NO_ENTITLEMENTS_MODULE,
TEST_PATH_LOOKUP,
- Set.of()
+ Set.of(),
+ cache
);
// "Unspecified" below means that the module is not named in the policy
@@ -109,6 +111,7 @@ public void testGetEntitlements() {
policyScope.set(PolicyScope.server("org.example.httpclient"));
resetAndCheckEntitlements(
"Specified entitlements for server",
+ cache,
getClass(),
policyManager.policyEntitlements(
SERVER.componentName,
@@ -122,6 +125,7 @@ public void testGetEntitlements() {
policyScope.set(PolicyScope.server("plugin.unspecifiedModule"));
resetAndCheckEntitlements(
"Default entitlements for unspecified module",
+ cache,
getClass(),
policyManager.defaultEntitlements(SERVER.componentName, thisSourcePath, "plugin.unspecifiedModule"),
policyManager
@@ -130,6 +134,7 @@ public void testGetEntitlements() {
policyScope.set(PolicyScope.plugin("plugin1", "plugin.module1"));
resetAndCheckEntitlements(
"Specified entitlements for plugin",
+ cache,
getClass(),
policyManager.policyEntitlements("plugin1", plugin1SourcePath, "plugin.module1", List.of(new ExitVMEntitlement())),
policyManager
@@ -138,6 +143,7 @@ public void testGetEntitlements() {
policyScope.set(PolicyScope.plugin("plugin1", "plugin.unspecifiedModule"));
resetAndCheckEntitlements(
"Default entitlements for plugin",
+ cache,
getClass(),
policyManager.defaultEntitlements("plugin1", plugin1SourcePath, "plugin.unspecifiedModule"),
policyManager
@@ -146,21 +152,22 @@ public void testGetEntitlements() {
private void resetAndCheckEntitlements(
String message,
+ EntitlementsCacheForTesting cache,
Class> requestingClass,
ModuleEntitlements expectedEntitlements,
PolicyManager policyManager
) {
- policyManager.moduleEntitlementsMap.clear();
+ cache.clear();
assertEquals(message, expectedEntitlements, policyManager.getEntitlements(requestingClass));
assertEquals(
"Map has precisely the one expected entry",
Map.of(requestingClass.getModule(), expectedEntitlements),
- policyManager.moduleEntitlementsMap
+ policyManager.moduleEntitlementsCache
);
// Fetch a second time and verify the map is unchanged
policyManager.getEntitlements(requestingClass);
- assertEquals("Map is unchanged", Map.of(requestingClass.getModule(), expectedEntitlements), policyManager.moduleEntitlementsMap);
+ assertEquals("Map is unchanged", Map.of(requestingClass.getModule(), expectedEntitlements), policyManager.moduleEntitlementsCache);
}
public void testRequestingClassFastPath() throws IOException, ClassNotFoundException {
@@ -211,7 +218,8 @@ public void testAgentsEntitlements() throws IOException, ClassNotFoundException
Map.of(),
NO_ENTITLEMENTS_MODULE,
TEST_PATH_LOOKUP,
- Set.of()
+ Set.of(),
+ new EntitlementsCacheForTesting()
);
ModuleEntitlements agentsEntitlements = policyManager.getEntitlements(TestAgent.class);
assertThat(agentsEntitlements.hasEntitlement(CreateClassLoaderEntitlement.class), is(true));
@@ -240,7 +248,8 @@ public void testDuplicateEntitlements() {
Map.of(),
NO_ENTITLEMENTS_MODULE,
TEST_PATH_LOOKUP,
- Set.of()
+ Set.of(),
+ new EntitlementsCacheForTesting()
)
);
assertEquals(
@@ -258,7 +267,8 @@ public void testDuplicateEntitlements() {
Map.of(),
NO_ENTITLEMENTS_MODULE,
TEST_PATH_LOOKUP,
- Set.of()
+ Set.of(),
+ new EntitlementsCacheForTesting()
)
);
assertEquals(
@@ -296,7 +306,8 @@ public void testDuplicateEntitlements() {
Map.of("plugin1", Path.of("modules", "plugin1")),
NO_ENTITLEMENTS_MODULE,
TEST_PATH_LOOKUP,
- Set.of()
+ Set.of(),
+ new EntitlementsCacheForTesting()
)
);
assertEquals(
@@ -348,7 +359,8 @@ public void testFilesEntitlementsWithExclusive() {
Map.of("plugin1", Path.of("modules", "plugin1"), "plugin2", Path.of("modules", "plugin2")),
NO_ENTITLEMENTS_MODULE,
TEST_PATH_LOOKUP,
- Set.of()
+ Set.of(),
+ new EntitlementsCacheForTesting()
)
);
assertThat(
@@ -401,7 +413,8 @@ public void testFilesEntitlementsWithExclusive() {
Map.of(),
NO_ENTITLEMENTS_MODULE,
TEST_PATH_LOOKUP,
- Set.of()
+ Set.of(),
+ new EntitlementsCacheForTesting()
)
);
assertEquals(
@@ -431,7 +444,8 @@ private static PolicyManager policyManager(Module entitlementsModule) {
Map.of(),
entitlementsModule,
TEST_PATH_LOOKUP,
- Set.of()
+ Set.of(),
+ new EntitlementsCacheForTesting()
);
}
From f2d9e853fc3510f1baa0581d6565a625b6924fea Mon Sep 17 00:00:00 2001
From: Patrick Doyle
Date: Tue, 6 May 2025 14:34:06 -0400
Subject: [PATCH 2/3] isTriviallyAllowed unconditionally delegate to
EntitlementsCache
---
.../entitlement/runtime/policy/PolicyManager.java | 10 ++++------
1 file changed, 4 insertions(+), 6 deletions(-)
diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java
index 08bbd3ba8598f..21b7fabe38a8d 100644
--- a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java
+++ b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java
@@ -803,12 +803,10 @@ private boolean isTriviallyAllowed(Class> requestingClass) {
generalLogger.debug("Entitlement trivially allowed from system module [{}]", requestingClass.getModule().getName());
return true;
}
- if (moduleEntitlementsCache.isAlwaysAllowed(requestingClass)) {
- generalLogger.debug("Entitlement trivially allowed by the EntitlementsCache");
- return true;
- }
- generalLogger.trace("Entitlement not trivially allowed");
- return false;
+
+ boolean result = moduleEntitlementsCache.isAlwaysAllowed(requestingClass);
+ generalLogger.trace("Returning isTriviallyAllowed=[{}]", result);
+ return result;
}
/**
From 7faa7961f63ccfeed4cbb539a1a56c8300aebae4 Mon Sep 17 00:00:00 2001
From: Patrick Doyle
Date: Tue, 6 May 2025 14:43:08 -0400
Subject: [PATCH 3/3] Fix unit test
---
.../entitlement/runtime/policy/PolicyManagerTests.java | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java b/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java
index 36e943596272e..68c346c18d453 100644
--- a/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java
+++ b/libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java
@@ -161,13 +161,13 @@ private void resetAndCheckEntitlements(
assertEquals(message, expectedEntitlements, policyManager.getEntitlements(requestingClass));
assertEquals(
"Map has precisely the one expected entry",
- Map.of(requestingClass.getModule(), expectedEntitlements),
+ Map.of(requestingClass, expectedEntitlements),
policyManager.moduleEntitlementsCache
);
// Fetch a second time and verify the map is unchanged
policyManager.getEntitlements(requestingClass);
- assertEquals("Map is unchanged", Map.of(requestingClass.getModule(), expectedEntitlements), policyManager.moduleEntitlementsCache);
+ assertEquals("Map is unchanged", Map.of(requestingClass, expectedEntitlements), policyManager.moduleEntitlementsCache);
}
public void testRequestingClassFastPath() throws IOException, ClassNotFoundException {