Initial TestPolicyManager implementation (elastic#128700)

* Initial TestPolicyManager implementation

* The forbidden APIs check is not messing around

Author: prdoyle

libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java

@@ -376,14 +376,21 @@ boolean isTriviallyAllowed(Class<?> requestingClass) {
             generalLogger.debug("Entitlement trivially allowed from outermost frame");
             return true;
         }
-        if (SYSTEM_LAYER_MODULES.contains(requestingClass.getModule())) {
+        if (isTrustedSystemClass(requestingClass)) {
             generalLogger.debug("Entitlement trivially allowed from system module [{}]", requestingClass.getModule().getName());
             return true;
         }
         generalLogger.trace("Entitlement not trivially allowed");
         return false;
     }
 
+    /**
+     * The main decision point for what counts as a trusted built-in JDK class.
+     */
+    protected boolean isTrustedSystemClass(Class<?> requestingClass) {
+        return SYSTEM_LAYER_MODULES.contains(requestingClass.getModule());
+    }
+
     @Override
     public String toString() {
         return "PolicyManager{" + "serverEntitlements=" + serverEntitlements + ", pluginsEntitlements=" + pluginsEntitlements + '}';

test/framework/src/main/java/org/elasticsearch/entitlement/initialization/TestEntitlementInitialization.java

@@ -20,6 +20,7 @@
 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.TestPolicyManager;
 import org.elasticsearch.plugins.PluginDescriptor;
 
 import java.io.IOException;
@@ -109,15 +110,14 @@ private static PolicyManager createPolicyManager(PathLookup pathLookup) throws I
 
         FilesEntitlementsValidation.validate(pluginPolicies, pathLookup);
 
-        PolicyManager policyManager = new PolicyManager(
+        return new TestPolicyManager(
             HardcodedEntitlements.serverPolicy(null, null),
             HardcodedEntitlements.agentEntitlements(),
             pluginPolicies,
             scopeResolver,
             Map.of(),
             pathLookup
         );
-        throw new IllegalStateException("Not yet implemented!");
     }
 
 }

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

@@ -0,0 +1,36 @@
+/*
+ * 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.nio.file.Path;
+import java.util.stream.Stream;
+
+public class TestPathLookup implements PathLookup {
+    @Override
+    public Path pidFile() {
+        return null;
+    }
+
+    @Override
+    public Stream<Path> getBaseDirPaths(BaseDir baseDir) {
+        return Stream.empty();
+    }
+
+    @Override
+    public Stream<Path> resolveRelativePaths(BaseDir baseDir, Path relativePath) {
+        return Stream.empty();
+    }
+
+    @Override
+    public Stream<Path> resolveSettingPaths(BaseDir baseDir, String settingName) {
+        return Stream.empty();
+    }
+
+}

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

@@ -0,0 +1,58 @@
+/*
+ * 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 java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+
+public class TestPolicyManager extends PolicyManager {
+    public TestPolicyManager(
+        Policy serverPolicy,
+        List<Entitlement> apmAgentEntitlements,
+        Map<String, Policy> pluginPolicies,
+        Function<Class<?>, PolicyScope> scopeResolver,
+        Map<String, Path> sourcePaths,
+        PathLookup pathLookup
+    ) {
+        super(serverPolicy, apmAgentEntitlements, pluginPolicies, scopeResolver, sourcePaths, pathLookup);
+    }
+
+    /**
+     * Called between tests so each test is not affected by prior tests
+     */
+    public void reset() {
+        super.moduleEntitlementsMap.clear();
+    }
+
+    @Override
+    protected boolean isTrustedSystemClass(Class<?> requestingClass) {
+        ClassLoader loader = requestingClass.getClassLoader();
+        return loader == null || loader == ClassLoader.getPlatformClassLoader();
+    }
+
+    @Override
+    boolean isTriviallyAllowed(Class<?> requestingClass) {
+        return isTestFrameworkClass(requestingClass) || isEntitlementClass(requestingClass) || super.isTriviallyAllowed(requestingClass);
+    }
+
+    private boolean isEntitlementClass(Class<?> requestingClass) {
+        return requestingClass.getPackageName().startsWith("org.elasticsearch.entitlement")
+            && (requestingClass.getName().contains("Test") == false);
+    }
+
+    private boolean isTestFrameworkClass(Class<?> requestingClass) {
+        String packageName = requestingClass.getPackageName();
+        return packageName.startsWith("org.junit") || packageName.startsWith("org.gradle");
+    }
+}

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

@@ -0,0 +1,59 @@
+/*
+ * 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.PolicyManager.PolicyScope;
+import org.elasticsearch.test.ESTestCase;
+import org.junit.Before;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ComponentKind.PLUGIN;
+
+public class TestPolicyManagerTests extends ESTestCase {
+    TestPolicyManager policyManager;
+
+    @Before
+    public void setupPolicyManager() {
+        AtomicInteger scopeCounter = new AtomicInteger(0);
+        policyManager = new TestPolicyManager(
+            new Policy("empty", List.of()),
+            List.of(),
+            Map.of(),
+            c -> new PolicyScope(PLUGIN, "example-plugin" + scopeCounter.incrementAndGet(), "org.example.module"),
+            Map.of(),
+            new TestPathLookup()
+        );
+    }
+
+    public void testReset() {
+        assertTrue(policyManager.moduleEntitlementsMap.isEmpty());
+        assertEquals("example-plugin1", policyManager.getEntitlements(getClass()).componentName());
+        assertEquals("example-plugin1", policyManager.getEntitlements(getClass()).componentName());
+        assertFalse(policyManager.moduleEntitlementsMap.isEmpty());
+
+        policyManager.reset();
+
+        assertTrue(policyManager.moduleEntitlementsMap.isEmpty());
+        assertEquals("example-plugin2", policyManager.getEntitlements(getClass()).componentName());
+        assertEquals("example-plugin2", policyManager.getEntitlements(getClass()).componentName());
+        assertFalse(policyManager.moduleEntitlementsMap.isEmpty());
+    }
+
+    public void testIsTriviallyAllowed() {
+        assertTrue(policyManager.isTriviallyAllowed(String.class));
+        assertTrue(policyManager.isTriviallyAllowed(org.junit.Before.class));
+        assertTrue(policyManager.isTriviallyAllowed(PolicyManager.class));
+
+        assertFalse(policyManager.isTriviallyAllowed(getClass()));
+    }
+}