Entitle com.unboundid.ldap.listener as test package

Author: mosche

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

@@ -87,6 +87,10 @@ public static void setTriviallyAllowingTestCode(boolean newValue) {
         policyManager.setTriviallyAllowingTestCode(newValue);
     }
 
+    public static void addEntitledTestPackages(String[] entitledTestPackages) {
+        policyManager.addEntitledTestPackages(entitledTestPackages);
+    }
+
     public static void reset() {
         if (policyManager != null) {
             policyManager.reset();

test/framework/src/main/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManager.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.entitlement.runtime.policy;
 
+import org.elasticsearch.common.util.ArrayUtils;
 import org.elasticsearch.entitlement.runtime.policy.entitlements.Entitlement;
 import org.elasticsearch.test.ESTestCase;
 
@@ -17,6 +18,7 @@
 import java.nio.file.Path;
 import java.security.CodeSource;
 import java.security.ProtectionDomain;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -29,6 +31,7 @@ public class TestPolicyManager extends PolicyManager {
 
     boolean isActive;
     boolean isTriviallyAllowingTestCode;
+    String[] entitledTestPackages = TEST_FRAMEWORK_PACKAGE_PREFIXES;
 
     /**
      * We don't have modules in tests, so we can't use the inherited map of entitlements per module.
@@ -60,6 +63,12 @@ public void setTriviallyAllowingTestCode(boolean newValue) {
         this.isTriviallyAllowingTestCode = newValue;
     }
 
+    public void addEntitledTestPackages(String[] entitledTestPackages) {
+        String[] packages = ArrayUtils.concat(this.entitledTestPackages, entitledTestPackages);
+        Arrays.sort(packages);
+        this.entitledTestPackages = packages;
+    }
+
     /**
      * Called between tests so each test is not affected by prior tests
      */
@@ -110,17 +119,14 @@ 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) {
-            if (packageName.startsWith(prefix)) {
-                return true;
-            }
+        int idx = Arrays.binarySearch(entitledTestPackages, packageName);
+        if (idx >= 0) {
+            return true;
         }
-        return false;
+        idx = -idx - 2; // candidate package (insertion point - 1)
+        return idx >= 0 && idx < entitledTestPackages.length && packageName.startsWith(entitledTestPackages[idx]);
     }
 
     private boolean isTestCode(Class<?> requestingClass) {
@@ -163,6 +169,10 @@ private boolean isTestCode(Class<?> requestingClass) {
         "org.bouncycastle.jsse.provider" // Used in test code if FIPS is enabled, support more fine-grained config in ES-12128
     };
 
+    static {
+        Arrays.sort(TEST_FRAMEWORK_PACKAGE_PREFIXES);
+    }
+
     @Override
     protected ModuleEntitlements getEntitlements(Class<?> requestingClass) {
         return classEntitlementsMap.computeIfAbsent(requestingClass, this::computeEntitlements);

test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java

@@ -517,16 +517,30 @@ protected void afterIfSuccessful() throws Exception {}
     public @interface WithEntitlementsOnTestCode {
     }
 
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.TYPE)
+    @Inherited
+    public @interface EntitledTestPackages {
+        String[] value();
+    }
+
     @BeforeClass
     public static void setupEntitlementsForClass() {
         boolean withoutEntitlements = getTestClass().isAnnotationPresent(WithoutEntitlements.class);
         boolean withEntitlementsOnTestCode = getTestClass().isAnnotationPresent(WithEntitlementsOnTestCode.class);
+        EntitledTestPackages entitledPackages = getTestClass().getAnnotation(EntitledTestPackages.class);
+
         if (TestEntitlementBootstrap.isEnabledForTest()) {
             TestEntitlementBootstrap.setActive(false == withoutEntitlements);
             TestEntitlementBootstrap.setTriviallyAllowingTestCode(false == withEntitlementsOnTestCode);
+            if (entitledPackages != null) {
+                assert withEntitlementsOnTestCode == false : "Cannot use @WithEntitlementsOnTestCode together with @EntitledTestPackages";
+                assert entitledPackages.value().length > 0 : "No test packages specified in @EntitledTestPackages";
+                TestEntitlementBootstrap.addEntitledTestPackages(entitledPackages.value());
+            }
         } else if (withEntitlementsOnTestCode) {
             throw new AssertionError(
-                "Cannot use WithEntitlementsOnTestCode on tests that are not configured to use entitlements for testing"
+                "Cannot use @WithEntitlementsOnTestCode on tests that are not configured to use entitlements for testing"
             );
         }
     }

x-pack/plugin/core/src/main/plugin-metadata/entitlement-policy.yaml

@@ -18,7 +18,6 @@ org.apache.httpcomponents.httpasyncclient:
   - manage_threads
 unboundid.ldapsdk:
   - set_https_connection_properties # TODO: review if we need this once we have proper test coverage
-  - inbound_network # For com.unboundid.ldap.listener.LDAPListener
   - outbound_network
   - manage_threads
   - write_system_properties:

x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/authc/ldap/ActiveDirectoryRealmTests.java

@@ -33,6 +33,7 @@
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.mustache.MustacheScriptEngine;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.ESTestCase.EntitledTestPackages;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.watcher.ResourceWatcherService;
@@ -106,6 +107,7 @@
  * The username used to authenticate then has to be in the form of CN=user. Finally the username needs to be added as an
  * additional bind DN with a password in the test setup since it really is not a DN in the ldif file
  */
+@EntitledTestPackages(value = { "com.unboundid.ldap.listener" }) // tests start LDAP server that listens for incoming connections
 public class ActiveDirectoryRealmTests extends ESTestCase {
 
     private static final String PASSWORD = "password";

x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/authc/ldap/support/LdapTestCase.java

@@ -29,6 +29,7 @@
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.env.TestEnvironment;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.ESTestCase.EntitledTestPackages;
 import org.elasticsearch.watcher.ResourceWatcherService;
 import org.elasticsearch.xpack.core.XPackSettings;
 import org.elasticsearch.xpack.core.security.authc.RealmConfig;
@@ -73,6 +74,7 @@
 import static org.elasticsearch.xpack.core.security.authc.ldap.support.SessionFactorySettings.URLS_SETTING;
 import static org.hamcrest.Matchers.is;
 
+@EntitledTestPackages(value = { "com.unboundid.ldap.listener" }) // tests start LDAP server that listens for incoming connections
 public abstract class LdapTestCase extends ESTestCase {
 
     protected static final RealmConfig.RealmIdentifier REALM_IDENTIFIER = new RealmConfig.RealmIdentifier("ldap", "ldap1");