Test isTestFrameworkClass and prevent adding redundant entries

Author: mosche

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

@@ -87,8 +87,8 @@ public static void setTriviallyAllowingTestCode(boolean newValue) {
         policyManager.setTriviallyAllowingTestCode(newValue);
     }
 
-    public static void addEntitledTestPackages(String[] entitledTestPackages) {
-        policyManager.addEntitledTestPackages(entitledTestPackages);
+    public static void setEntitledTestPackages(String[] entitledTestPackages) {
+        policyManager.setEntitledTestPackages(entitledTestPackages);
     }
 
     public static void reset() {

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

@@ -63,7 +63,11 @@ public void setTriviallyAllowingTestCode(boolean newValue) {
         this.isTriviallyAllowingTestCode = newValue;
     }
 
-    public void addEntitledTestPackages(String[] entitledTestPackages) {
+    public void setEntitledTestPackages(String... entitledTestPackages) {
+        assertNoRedundantPrefixes(TEST_FRAMEWORK_PACKAGE_PREFIXES, entitledTestPackages, false);
+        if (entitledTestPackages.length > 1) {
+            assertNoRedundantPrefixes(entitledTestPackages, entitledTestPackages, true);
+        }
         String[] packages = ArrayUtils.concat(this.entitledTestPackages, entitledTestPackages);
         Arrays.sort(packages);
         this.entitledTestPackages = packages;
@@ -120,13 +124,44 @@ private boolean isEntitlementClass(Class<?> requestingClass) {
     }
 
     private boolean isTestFrameworkClass(Class<?> requestingClass) {
-        String packageName = requestingClass.getPackageName();
-        int idx = Arrays.binarySearch(entitledTestPackages, packageName);
+        return isTestFrameworkClass(entitledTestPackages, requestingClass.getPackageName());
+    }
+
+    // no redundant entries allowed, see assertNoRedundantPrefixes
+    static boolean isTestFrameworkClass(String[] sortedPrefixes, String packageName) {
+        int idx = Arrays.binarySearch(sortedPrefixes, packageName);
         if (idx >= 0) {
             return true;
         }
-        idx = -idx - 2; // candidate package (insertion point - 1)
-        return idx >= 0 && idx < entitledTestPackages.length && packageName.startsWith(entitledTestPackages[idx]);
+        idx = -idx - 2; // candidate package index (insertion point - 1)
+        if (idx >= 0 && idx < sortedPrefixes.length) {
+            String candidate = sortedPrefixes[idx];
+            if (packageName.startsWith(candidate)
+                && (packageName.length() == candidate.length() || packageName.charAt(candidate.length()) == '.')) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private static boolean isNotPrefixMatch(String name, String prefix, boolean discardExactMatch) {
+        assert prefix.endsWith(".") == false : "Invalid package prefix ending with '.' [" + prefix + "]";
+        if (name == prefix || name.startsWith(prefix)) {
+            if (name.length() == prefix.length()) {
+                return discardExactMatch;
+            }
+            return false == (name.length() > prefix.length() && name.charAt(prefix.length()) == '.');
+        }
+        return true;
+    }
+
+    static void assertNoRedundantPrefixes(String[] setA, String[] setB, boolean discardExactMatch) {
+        for (String a : setA) {
+            for (String b : setB) {
+                assert isNotPrefixMatch(a, b, discardExactMatch) && isNotPrefixMatch(b, a, discardExactMatch)
+                    : "Redundant prefix entries: [" + a + ", " + b + "]";
+            }
+        }
     }
 
     private boolean isTestCode(Class<?> requestingClass) {

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

@@ -536,7 +536,7 @@ public static void setupEntitlementsForClass() {
             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());
+                TestEntitlementBootstrap.setEntitledTestPackages(entitledPackages.value());
             }
         } else if (withEntitlementsOnTestCode) {
             throw new AssertionError(

test/framework/src/test/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManagerTests.java

@@ -13,11 +13,15 @@
 import org.elasticsearch.test.ESTestCase;
 import org.junit.Before;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ComponentKind.PLUGIN;
+import static org.hamcrest.Matchers.both;
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
 
 public class TestPolicyManagerTests extends ESTestCase {
     TestPolicyManager policyManager;
@@ -60,4 +64,43 @@ public void testIsTriviallyAllowed() {
         policyManager.setTriviallyAllowingTestCode(false);
         assertFalse(policyManager.isTriviallyAllowed(getClass()));
     }
+
+    public void testDefaultEntitledTestPackages() {
+        String[] testPackages = policyManager.entitledTestPackages.clone();
+        TestPolicyManager.assertNoRedundantPrefixes(testPackages, testPackages, true);
+
+        Arrays.sort(testPackages);
+        assertThat("Entitled test framework packages are not sorted", policyManager.entitledTestPackages, equalTo(testPackages));
+    }
+
+    public void testRejectSetRedundantEntitledTestPackages() {
+        var throwable = expectThrows(AssertionError.class, () -> policyManager.setEntitledTestPackages("org.apache.lucene.tests"));
+        var baseMatcher = both(containsString("Redundant prefix entries"));
+        assertThat(throwable.getMessage(), baseMatcher.and(containsString("org.apache.lucene.tests, org.apache.lucene.tests")));
+
+        throwable = expectThrows(AssertionError.class, () -> policyManager.setEntitledTestPackages("org.apache.lucene"));
+        assertThat(throwable.getMessage(), baseMatcher.and(containsString("org.apache.lucene.tests, org.apache.lucene")));
+
+        throwable = expectThrows(AssertionError.class, () -> policyManager.setEntitledTestPackages("org.apache.lucene.tests.whatever"));
+        assertThat(throwable.getMessage(), baseMatcher.and(containsString("org.apache.lucene.tests, org.apache.lucene.tests.whatever")));
+
+        throwable = expectThrows(AssertionError.class, () -> policyManager.setEntitledTestPackages("my.package", "my.package.sub"));
+        assertThat(throwable.getMessage(), baseMatcher.and(containsString("my.package, my.package.sub")));
+
+        throwable = expectThrows(AssertionError.class, () -> policyManager.setEntitledTestPackages("trailing.dot."));
+        assertThat(throwable.getMessage(), containsString("Invalid package prefix ending with '.' [trailing.dot.]"));
+    }
+
+    public void testIsTestFrameworkClass() {
+        String[] sortedPrefixes = { "a.b", "a.bc", "a.c" };
+
+        assertTrue(TestPolicyManager.isTestFrameworkClass(sortedPrefixes, "a.b"));
+        assertTrue(TestPolicyManager.isTestFrameworkClass(sortedPrefixes, "a.b.c"));
+        assertTrue(TestPolicyManager.isTestFrameworkClass(sortedPrefixes, "a.bc"));
+        assertTrue(TestPolicyManager.isTestFrameworkClass(sortedPrefixes, "a.bc.a"));
+
+        assertFalse(TestPolicyManager.isTestFrameworkClass(sortedPrefixes, "a"));
+        assertFalse(TestPolicyManager.isTestFrameworkClass(sortedPrefixes, "a.ba"));
+        assertFalse(TestPolicyManager.isTestFrameworkClass(sortedPrefixes, "a.bcc"));
+    }
 }