System properties: testOnlyClasspath and enableForTests

prdoyle · prdoyle · commit b303ab24302e · 2025-06-17T10:36:46.000-04:00
diff --git a/build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/ElasticsearchTestBasePlugin.java b/build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/ElasticsearchTestBasePlugin.java
@@ -37,6 +37,7 @@
 
 import javax.inject.Inject;
 
+import static java.util.Objects.requireNonNull;
 import static org.elasticsearch.gradle.internal.util.ParamsUtils.loadBuildParams;
 import static org.elasticsearch.gradle.util.FileUtils.mkdirs;
 import static org.elasticsearch.gradle.util.GradleUtils.maybeConfigure;
@@ -173,6 +174,16 @@ public void execute(Task t) {
             // we use 'temp' relative to CWD since this is per JVM and tests are forbidden from writing to CWD
             nonInputProperties.systemProperty("java.io.tmpdir", test.getWorkingDir().toPath().resolve("temp"));
 
+            SourceSetContainer sourceSets = project.getExtensions().getByType(SourceSetContainer.class);
+            SourceSet mainSourceSet = sourceSets.findByName(SourceSet.MAIN_SOURCE_SET_NAME);
+            SourceSet testSourceSet = sourceSets.findByName(SourceSet.TEST_SOURCE_SET_NAME);
+            if (mainSourceSet != null && testSourceSet != null) {
+                FileCollection mainRuntime = mainSourceSet.getRuntimeClasspath();
+                FileCollection testRuntime = testSourceSet.getRuntimeClasspath();
+                FileCollection testOnlyFiles = testRuntime.minus(mainRuntime);
+                nonInputProperties.systemProperty("es.entitlement.testOnlyPath", testOnlyFiles::getAsPath);
+            }
+
             test.systemProperties(getProviderFactory().systemPropertiesPrefixedBy("tests.").get());
             test.systemProperties(getProviderFactory().systemPropertiesPrefixedBy("es.").get());
 
@@ -211,13 +222,12 @@ public void execute(Task t) {
             project.getPluginManager().withPlugin("com.gradleup.shadow", p -> {
                 if (test.getName().equals(JavaPlugin.TEST_TASK_NAME)) {
                     // Remove output class files and any other dependencies from the test classpath, since the shadow JAR includes these
-                    SourceSetContainer sourceSets = project.getExtensions().getByType(SourceSetContainer.class);
-                    FileCollection mainRuntime = sourceSets.getByName(SourceSet.MAIN_SOURCE_SET_NAME).getRuntimeClasspath();
                     // Add any "shadow" dependencies. These are dependencies that are *not* bundled into the shadow JAR
                     Configuration shadowConfig = project.getConfigurations().getByName(ShadowBasePlugin.CONFIGURATION_NAME);
                     // Add the shadow JAR artifact itself
                     FileCollection shadowJar = project.files(project.getTasks().named("shadowJar"));
-                    FileCollection testRuntime = sourceSets.getByName(SourceSet.TEST_SOURCE_SET_NAME).getRuntimeClasspath();
+                    FileCollection mainRuntime = requireNonNull(mainSourceSet).getRuntimeClasspath();
+                    FileCollection testRuntime = requireNonNull(testSourceSet).getRuntimeClasspath();
                     test.setClasspath(testRuntime.minus(mainRuntime).plus(shadowConfig).plus(shadowJar));
                 }
             });
diff --git a/build-tools/src/main/java/org/elasticsearch/gradle/test/TestBuildInfoPlugin.java b/build-tools/src/main/java/org/elasticsearch/gradle/test/TestBuildInfoPlugin.java
@@ -18,8 +18,11 @@
 import org.gradle.api.provider.ProviderFactory;
 import org.gradle.api.tasks.SourceSet;
 import org.gradle.api.tasks.SourceSetContainer;
+import org.gradle.api.tasks.testing.Test;
 import org.gradle.language.jvm.tasks.ProcessResources;
 
+import java.util.List;
+
 import javax.inject.Inject;
 
 /**
@@ -53,5 +56,11 @@ public void apply(Project project) {
         project.getTasks().withType(ProcessResources.class).named("processResources").configure(task -> {
             task.into("META-INF", copy -> copy.from(testBuildInfoTask));
         });
+
+        project.getTasks().withType(Test.class).configureEach(test -> {
+            if (List.of("test", "internalClusterTest").contains(test.getName())) {
+                test.systemProperty("es.entitlement.enableForTests", "true");
+            }
+        });
     }
 }
diff --git a/test/framework/src/main/java/org/elasticsearch/entitlement/bootstrap/TestEntitlementBootstrap.java b/test/framework/src/main/java/org/elasticsearch/entitlement/bootstrap/TestEntitlementBootstrap.java
@@ -29,11 +29,14 @@
 import java.net.URL;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
 
 public class TestEntitlementBootstrap {
 
@@ -45,13 +48,20 @@ public class TestEntitlementBootstrap {
      * Activates entitlement checking in tests.
      */
     public static void bootstrap(Path tempDir) throws IOException {
+        if (isEnabledForTest() == false) {
+            return;
+        }
         TestPathLookup pathLookup = new TestPathLookup(List.of(tempDir));
         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 boolean isEnabledForTest() {
+        return Boolean.getBoolean("es.entitlement.enableForTests");
+    }
+
     public static void setActive(boolean newValue) {
         policyManager.setActive(newValue);
     }
@@ -61,7 +71,9 @@ public static void setTriviallyAllowingTestCode(boolean newValue) {
     }
 
     public static void reset() {
-        policyManager.reset();
+        if (policyManager != null) {
+            policyManager.reset();
+        }
     }
 
     private static TestPolicyManager createPolicyManager(PathLookup pathLookup) throws IOException {
@@ -79,13 +91,22 @@ private static TestPolicyManager createPolicyManager(PathLookup pathLookup) thro
 
         FilesEntitlementsValidation.validate(pluginPolicies, pathLookup);
 
+        String testOnlyPathProperty = System.getProperty("es.entitlement.testOnlyPath");
+        Set<String> testOnlyClassPath;
+        if (testOnlyPathProperty == null) {
+            testOnlyClassPath = Set.of();
+        } else {
+            testOnlyClassPath = Arrays.stream(testOnlyPathProperty.split(":")).collect(Collectors.toCollection(TreeSet::new));
+        }
+
         return new TestPolicyManager(
             HardcodedEntitlements.serverPolicy(null, null),
             HardcodedEntitlements.agentEntitlements(),
             pluginPolicies,
             scopeResolver,
             pluginSourcePaths,
-            pathLookup
+            pathLookup,
+            testOnlyClassPath
         );
     }
 
diff --git a/test/framework/src/main/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManager.java b/test/framework/src/main/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManager.java
@@ -13,13 +13,18 @@
 import org.elasticsearch.test.ESTestCase;
 
 import java.nio.file.Path;
+import java.security.CodeSource;
+import java.security.ProtectionDomain;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 
+import static java.util.Objects.requireNonNull;
+
 public class TestPolicyManager extends PolicyManager {
+
     boolean isActive;
     boolean isTriviallyAllowingTestCode;
 
@@ -29,15 +34,19 @@ public class TestPolicyManager extends PolicyManager {
      */
     final Map<Class<?>, ModuleEntitlements> classEntitlementsMap = new ConcurrentHashMap<>();
 
+    final Collection<String> testOnlyClasspath;
+
     public TestPolicyManager(
         Policy serverPolicy,
         List<Entitlement> apmAgentEntitlements,
         Map<String, Policy> pluginPolicies,
         Function<Class<?>, PolicyScope> scopeResolver,
         Map<String, Collection<Path>> pluginSourcePaths,
-        PathLookup pathLookup
+        PathLookup pathLookup,
+        Collection<String> testOnlyClasspath
     ) {
         super(serverPolicy, apmAgentEntitlements, pluginPolicies, scopeResolver, pluginSourcePaths, pathLookup);
+        this.testOnlyClasspath = testOnlyClasspath;
         reset();
     }
 
@@ -70,10 +79,13 @@ boolean isTriviallyAllowed(Class<?> requestingClass) {
         if (isActive == false) {
             return true;
         }
-        if (isTriviallyAllowingTestCode && isTestCaseClass(requestingClass)) {
+        if (isEntitlementClass(requestingClass)) {
+            return true;
+        }
+        if (isTestFrameworkClass(requestingClass)) {
             return true;
         }
-        if (isTestFrameworkClass(requestingClass) || isEntitlementClass(requestingClass)) {
+        if (isTriviallyAllowingTestCode && isTestCode(requestingClass)) {
             return true;
         }
         return super.isTriviallyAllowed(requestingClass);
@@ -84,6 +96,9 @@ private boolean isEntitlementClass(Class<?> requestingClass) {
             && (requestingClass.getName().contains("Test") == false);
     }
 
+    @Deprecated // TODO: reevaluate whether we want this.
+    // If we can simply check for dependencies the gradle worker has that aren't
+    // declared in the gradle config (namely org.gradle) that would be simpler.
     private boolean isTestFrameworkClass(Class<?> requestingClass) {
         String packageName = requestingClass.getPackageName();
         for (String prefix : TEST_FRAMEWORK_PACKAGE_PREFIXES) {
@@ -94,13 +109,25 @@ private boolean isTestFrameworkClass(Class<?> requestingClass) {
         return false;
     }
 
-    private boolean isTestCaseClass(Class<?> requestingClass) {
-        for (Class<?> candidate = requestingClass; candidate != null; candidate = candidate.getDeclaringClass()) {
+    private boolean isTestCode(Class<?> requestingClass) {
+        // TODO: Cache this? It's expensive
+        for (Class<?> candidate = requireNonNull(requestingClass); candidate != null; candidate = candidate.getDeclaringClass()) {
             if (ESTestCase.class.isAssignableFrom(candidate)) {
                 return true;
             }
         }
-        return false;
+        ProtectionDomain protectionDomain = requestingClass.getProtectionDomain();
+        CodeSource codeSource = protectionDomain.getCodeSource();
+        if (codeSource == null) {
+            // This can happen for JDK classes
+            return false;
+        }
+        String needle = codeSource.getLocation().getPath();
+        if (needle.endsWith("/")) {
+            needle = needle.substring(0, needle.length() - 1);
+        }
+        boolean result = testOnlyClasspath.contains(needle);
+        return result;
     }
 
     private static final String[] TEST_FRAMEWORK_PACKAGE_PREFIXES = {
diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
@@ -514,10 +514,16 @@ protected void afterIfSuccessful() throws Exception {}
 
     @BeforeClass
     public static void setupEntitlementsForClass() {
-        TestEntitlementBootstrap.setActive(false == getTestClass().isAnnotationPresent(WithoutEntitlements.class));
-        TestEntitlementBootstrap.setTriviallyAllowingTestCode(
-            false == getTestClass().isAnnotationPresent(WithEntitlementsOnTestCode.class)
-        );
+        boolean withoutEntitlements = getTestClass().isAnnotationPresent(WithoutEntitlements.class);
+        boolean withEntitlementsOnTestCode = getTestClass().isAnnotationPresent(WithEntitlementsOnTestCode.class);
+        if (TestEntitlementBootstrap.isEnabledForTest()) {
+            TestEntitlementBootstrap.setActive(false == withoutEntitlements);
+            TestEntitlementBootstrap.setTriviallyAllowingTestCode(false == withEntitlementsOnTestCode);
+        } else if (withEntitlementsOnTestCode) {
+            throw new AssertionError(
+                "Cannot use WithEntitlementsOnTestCode on tests that are not configured to use entitlements for testing"
+            );
+        }
     }
 
     @AfterClass
diff --git a/test/framework/src/test/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManagerTests.java b/test/framework/src/test/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManagerTests.java
@@ -31,7 +31,8 @@ public void setupPolicyManager() {
             Map.of(),
             c -> new PolicyScope(PLUGIN, "example-plugin" + scopeCounter.incrementAndGet(), "org.example.module"),
             Map.of(),
-            new TestPathLookup(List.of())
+            new TestPathLookup(List.of()),
+            List.of()
         );
         policyManager.setActive(true);
     }
