Painless: Add whitelist extensions

modules/lang-painless/src/main/java/org/elasticsearch/painless/Definition.java

@@ -46,29 +46,6 @@ public final class Definition {
 
     private static final Pattern TYPE_NAME_PATTERN = Pattern.compile("^[_a-zA-Z][._a-zA-Z0-9]*$");
 
-    public static final String[] DEFINITION_FILES = new String[] {
-            "org.elasticsearch.txt",
-            "java.lang.txt",
-            "java.math.txt",
-            "java.text.txt",
-            "java.time.txt",
-            "java.time.chrono.txt",
-            "java.time.format.txt",
-            "java.time.temporal.txt",
-            "java.time.zone.txt",
-            "java.util.txt",
-            "java.util.function.txt",
-            "java.util.regex.txt",
-            "java.util.stream.txt",
-            "joda.time.txt"
-    };
-
-    /**
-     * Whitelist that is "built in" to Painless and required by all scripts.
-     */
-    public static final Definition DEFINITION = new Definition(
-        Collections.singletonList(WhitelistLoader.loadFromResourceFiles(Definition.class, DEFINITION_FILES)));
-
     /** Some native types as constants: */
     public final Type voidType;
     public final Type booleanType;

modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessExtension.java

@@ -0,0 +1,30 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless;
+
+import java.util.List;
+import java.util.Map;
+
+import org.elasticsearch.script.ScriptContext;
+
+public interface PainlessExtension {
+
+    Map<ScriptContext<?>, List<Whitelist>> getContextWhitelists();
+}

modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessPlugin.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.painless;
 
 
+import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.plugins.ExtensiblePlugin;
@@ -28,22 +29,49 @@
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.ServiceLoader;
 
 /**
  * Registers Painless as a plugin.
  */
 public final class PainlessPlugin extends Plugin implements ScriptPlugin, ExtensiblePlugin {
 
+    private final Map<ScriptContext<?>, List<Whitelist>> extendedWhitelists = new HashMap<>();
+
     @Override
     public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
-        return new PainlessScriptEngine(settings, contexts);
+        Map<ScriptContext<?>, List<Whitelist>> contextsWithWhitelists = new HashMap<>();
+        for (ScriptContext<?> context : contexts) {
+            // we might have a context that only uses the base whitelists, so would not have been filled in by reloadSPI
+            List<Whitelist> whitelists = extendedWhitelists.get(context);
+            if (whitelists == null) {
+                whitelists = new ArrayList<>(Whitelist.BASE_WHITELISTS);
+            }
+            contextsWithWhitelists.put(context, whitelists);
+        }
+        return new PainlessScriptEngine(settings, contextsWithWhitelists);
     }
 
     @Override
     public List<Setting<?>> getSettings() {
         return Arrays.asList(CompilerSettings.REGEX_ENABLED);
     }
+
+    @Override
+    public void reloadSPI(ClassLoader loader) {
+        for (PainlessExtension extension : ServiceLoader.load(PainlessExtension.class, loader)) {
+            for (Map.Entry<ScriptContext<?>, List<Whitelist>> entry : extension.getContextWhitelists().entrySet()) {
+                List<Whitelist> existing = extendedWhitelists.computeIfAbsent(entry.getKey(),
+                    c -> new ArrayList<>(Whitelist.BASE_WHITELISTS));
+                existing.addAll(entry.getValue());
+            }
+        }
+    }
 }

modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessScriptEngine.java

@@ -82,7 +82,7 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
 
     /**
      * Default compiler settings to be used. Note that {@link CompilerSettings} is mutable but this instance shouldn't be mutated outside
-     * of {@link PainlessScriptEngine#PainlessScriptEngine(Settings, Collection)}.
+     * of {@link PainlessScriptEngine#PainlessScriptEngine(Settings, Map)}.
      */
     private final CompilerSettings defaultCompilerSettings = new CompilerSettings();
 
@@ -92,23 +92,19 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
      * Constructor.
      * @param settings The settings to initialize the engine with.
      */
-    public PainlessScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
+    public PainlessScriptEngine(Settings settings, Map<ScriptContext<?>, List<Whitelist>> contexts) {
         super(settings);
 
         defaultCompilerSettings.setRegexesEnabled(CompilerSettings.REGEX_ENABLED.get(settings));
 
         Map<ScriptContext<?>, Compiler> contextsToCompilers = new HashMap<>();
 
-        // Placeholder definition used for all contexts until SPI is fully integrated.  Reduces memory foot print
-        // by re-using the same definition since caching isn't implemented at this time.
-        Definition definition = new Definition(
-            Collections.singletonList(WhitelistLoader.loadFromResourceFiles(Definition.class, Definition.DEFINITION_FILES)));
-
-        for (ScriptContext<?> context : contexts) {
+        for (Map.Entry<ScriptContext<?>, List<Whitelist>> entry : contexts.entrySet()) {
+            ScriptContext<?> context = entry.getKey();
             if (context.instanceClazz.equals(SearchScript.class) || context.instanceClazz.equals(ExecutableScript.class)) {
-                contextsToCompilers.put(context, new Compiler(GenericElasticsearchScript.class, definition));
+                contextsToCompilers.put(context, new Compiler(GenericElasticsearchScript.class, new Definition(entry.getValue())));
             } else {
-                contextsToCompilers.put(context, new Compiler(context.instanceClazz, definition));
+                contextsToCompilers.put(context, new Compiler(context.instanceClazz, new Definition(entry.getValue())));
             }
         }
 

modules/lang-painless/src/main/java/org/elasticsearch/painless/Whitelist.java

@@ -34,6 +34,26 @@
  */
 public final class Whitelist {
 
+    private static final String[] BASE_WHITELIST_FILES = new String[] {
+        "org.elasticsearch.txt",
+        "java.lang.txt",
+        "java.math.txt",
+        "java.text.txt",
+        "java.time.txt",
+        "java.time.chrono.txt",
+        "java.time.format.txt",
+        "java.time.temporal.txt",
+        "java.time.zone.txt",
+        "java.util.txt",
+        "java.util.function.txt",
+        "java.util.regex.txt",
+        "java.util.stream.txt",
+        "joda.time.txt"
+    };
+
+    public static final List<Whitelist> BASE_WHITELISTS =
+        Collections.singletonList(WhitelistLoader.loadFromResourceFiles(Whitelist.class, BASE_WHITELIST_FILES));
+
     /**
      * Struct represents the equivalent of a Java class in Painless complete with super classes,
      * constructors, methods, and fields.  In Painless a class is known as a struct primarily to avoid

modules/lang-painless/src/main/java/org/elasticsearch/painless/WhitelistLoader.java

@@ -25,6 +25,8 @@
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.nio.charset.StandardCharsets;
+import java.security.AccessController;
+import java.security.PrivilegedAction;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -296,8 +298,9 @@ public static Whitelist loadFromResourceFiles(Class<?> resource, String... filep
                 throw new RuntimeException("error in [" + filepath + "] at line [" + number + "]", exception);
             }
         }
+        ClassLoader loader = AccessController.doPrivileged((PrivilegedAction<ClassLoader>)resource::getClassLoader);
 
-        return new Whitelist(resource.getClassLoader(), whitelistStructs);
+        return new Whitelist(loader, whitelistStructs);
     }
 
     private WhitelistLoader() {}

modules/lang-painless/src/main/plugin-metadata/plugin-security.policy

@@ -20,4 +20,7 @@
 grant {
   // needed to generate runtime classes
   permission java.lang.RuntimePermission "createClassLoader";
+
+  // needed to find the classloader to load whitelisted classes from
+  permission java.lang.RuntimePermission "getClassLoader";
 };

modules/lang-painless/src/test/java/org/elasticsearch/painless/AnalyzerCasterTests.java

@@ -23,14 +23,9 @@
 import org.elasticsearch.painless.Definition.Type;
 import org.elasticsearch.test.ESTestCase;
 
-import java.util.Collections;
-
-import static org.elasticsearch.painless.Definition.DEFINITION_FILES;
-
 public class AnalyzerCasterTests extends ESTestCase {
 
-    private static final Definition definition = new Definition(
-        Collections.singletonList(WhitelistLoader.loadFromResourceFiles(Definition.class, DEFINITION_FILES)));
+    private static final Definition definition = new Definition(Whitelist.BASE_WHITELISTS);
 
     private static void assertCast(Type actual, Type expected, boolean mustBeExplicit) {
         Location location = new Location("dummy", 0);

modules/lang-painless/src/test/java/org/elasticsearch/painless/BaseClassTests.java

@@ -19,9 +19,6 @@
 
 package org.elasticsearch.painless;
 
-import org.elasticsearch.script.ScriptContext;
-
-import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
@@ -37,8 +34,7 @@
  */
 public class BaseClassTests extends ScriptTestCase {
 
-    private final Definition definition = new Definition(
-        Collections.singletonList(WhitelistLoader.loadFromResourceFiles(Definition.class, Definition.DEFINITION_FILES)));
+    private final Definition definition = new Definition(Whitelist.BASE_WHITELISTS);
 
     public abstract static class Gets {
 

modules/lang-painless/src/test/java/org/elasticsearch/painless/DebugTests.java

@@ -25,7 +25,6 @@
 import org.elasticsearch.script.ScriptException;
 
 import java.io.IOException;
-import java.util.Collections;
 import java.util.Map;
 
 import static java.util.Collections.singletonList;
@@ -35,8 +34,7 @@
 import static org.hamcrest.Matchers.not;
 
 public class DebugTests extends ScriptTestCase {
-    private final Definition definition = new Definition(
-        Collections.singletonList(WhitelistLoader.loadFromResourceFiles(Definition.class, Definition.DEFINITION_FILES)));
+    private final Definition definition = new Definition(Whitelist.BASE_WHITELISTS);
 
     public void testExplain() {
         // Debug.explain can explain an object

modules/lang-painless/src/test/java/org/elasticsearch/painless/Debugger.java

@@ -24,7 +24,6 @@
 
 import java.io.PrintWriter;
 import java.io.StringWriter;
-import java.util.Collections;
 
 /** quick and dirty tools for debugging */
 final class Debugger {
@@ -40,8 +39,7 @@ static String toString(Class<?> iface, String source, CompilerSettings settings)
         PrintWriter outputWriter = new PrintWriter(output);
         Textifier textifier = new Textifier();
         try {
-            new Compiler(iface, new Definition(
-                    Collections.singletonList(WhitelistLoader.loadFromResourceFiles(Definition.class, Definition.DEFINITION_FILES))))
+            new Compiler(iface, new Definition(Whitelist.BASE_WHITELISTS))
                     .compile("<debugging>", source, settings, textifier);
         } catch (Exception e) {
             textifier.print(outputWriter);

modules/lang-painless/src/test/java/org/elasticsearch/painless/DefBootstrapTests.java

@@ -30,8 +30,7 @@
 import org.elasticsearch.test.ESTestCase;
 
 public class DefBootstrapTests extends ESTestCase {
-    private final Definition definition = new Definition(
-        Collections.singletonList(WhitelistLoader.loadFromResourceFiles(Definition.class, Definition.DEFINITION_FILES)));
+    private final Definition definition = new Definition(Whitelist.BASE_WHITELISTS);
 
     /** calls toString() on integers, twice */
     public void testOneType() throws Throwable {

modules/lang-painless/src/test/java/org/elasticsearch/painless/FactoryTests.java

@@ -24,16 +24,18 @@
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 
 public class FactoryTests extends ScriptTestCase {
 
-    protected Collection<ScriptContext<?>> scriptContexts() {
-        Collection<ScriptContext<?>> contexts = super.scriptContexts();
-        contexts.add(StatefulFactoryTestScript.CONTEXT);
-        contexts.add(FactoryTestScript.CONTEXT);
-        contexts.add(EmptyTestScript.CONTEXT);
-        contexts.add(TemplateScript.CONTEXT);
+    @Override
+    protected Map<ScriptContext<?>, List<Whitelist>> scriptContexts() {
+        Map<ScriptContext<?>, List<Whitelist>> contexts = super.scriptContexts();
+        contexts.put(StatefulFactoryTestScript.CONTEXT, Whitelist.BASE_WHITELISTS);
+        contexts.put(FactoryTestScript.CONTEXT, Whitelist.BASE_WHITELISTS);
+        contexts.put(EmptyTestScript.CONTEXT, Whitelist.BASE_WHITELISTS);
+        contexts.put(TemplateScript.CONTEXT, Whitelist.BASE_WHITELISTS);
 
         return contexts;
     }

modules/lang-painless/src/test/java/org/elasticsearch/painless/NeedsScoreTests.java

@@ -24,12 +24,16 @@
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.script.ExecutableScript;
+import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 /**
  * Test that needsScores() is reported correctly depending on whether _score is used
@@ -40,8 +44,10 @@ public class NeedsScoreTests extends ESSingleNodeTestCase {
     public void testNeedsScores() {
         IndexService index = createIndex("test", Settings.EMPTY, "type", "d", "type=double");
 
-        PainlessScriptEngine service = new PainlessScriptEngine(Settings.EMPTY,
-            Arrays.asList(SearchScript.CONTEXT, ExecutableScript.CONTEXT));
+        Map<ScriptContext<?>, List<Whitelist>> contexts = new HashMap<>();
+        contexts.put(SearchScript.CONTEXT, Whitelist.BASE_WHITELISTS);
+        contexts.put(ExecutableScript.CONTEXT, Whitelist.BASE_WHITELISTS);
+        PainlessScriptEngine service = new PainlessScriptEngine(Settings.EMPTY, contexts);
 
         QueryShardContext shardContext = index.newQueryShardContext(0, null, () -> 0, null);
         SearchLookup lookup = new SearchLookup(index.mapperService(), shardContext::getForField, null);

modules/lang-painless/src/test/java/org/elasticsearch/painless/PainlessDocGenerator.java

@@ -27,7 +27,6 @@
 import org.elasticsearch.painless.Definition.Method;
 import org.elasticsearch.painless.Definition.Struct;
 import org.elasticsearch.painless.Definition.Type;
-import org.elasticsearch.painless.api.Augmentation;
 
 import java.io.IOException;
 import java.io.PrintStream;
@@ -36,7 +35,6 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardOpenOption;
-import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
@@ -68,8 +66,7 @@ public static void main(String[] args) throws IOException {
                 Files.newOutputStream(indexPath, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE),
                 false, StandardCharsets.UTF_8.name())) {
             emitGeneratedWarning(indexStream);
-            List<Type> types = new Definition(Collections.singletonList(
-                    WhitelistLoader.loadFromResourceFiles(Definition.class, Definition.DEFINITION_FILES))).
+            List<Type> types = new Definition(Whitelist.BASE_WHITELISTS).
                     allSimpleTypes().stream().sorted(comparing(t -> t.name)).collect(toList());
             for (Type type : types) {
                 if (type.clazz.isPrimitive()) {