Painless: Add whitelist extensions (#28161)

This commit adds a PainlessExtension which may be plugged in via SPI to
add additional classes, methods and members to the painless whitelist on
a per context basis. An example plugin adding and using a whitelist is
also added.

Author: rjernst

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

@@ -22,6 +22,7 @@
 import org.elasticsearch.bootstrap.BootstrapInfo;
 import org.elasticsearch.painless.antlr.Walker;
 import org.elasticsearch.painless.node.SSource;
+import org.elasticsearch.painless.spi.Whitelist;
 import org.objectweb.asm.util.Printer;
 
 import java.lang.reflect.Constructor;

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

@@ -20,6 +20,7 @@
 package org.elasticsearch.painless;
 
 import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.painless.spi.Whitelist;
 
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
@@ -46,29 +47,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/PainlessPlugin.java

@@ -22,28 +22,56 @@
 
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.painless.spi.PainlessExtension;
+import org.elasticsearch.painless.spi.Whitelist;
 import org.elasticsearch.plugins.ExtensiblePlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.ScriptPlugin;
 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.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

@@ -19,12 +19,12 @@
 
 package org.elasticsearch.painless;
 
-import org.apache.logging.log4j.core.tools.Generate;
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.SpecialPermission;
 import org.elasticsearch.common.component.AbstractComponent;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.painless.Compiler.Loader;
+import org.elasticsearch.painless.spi.Whitelist;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
@@ -45,7 +45,6 @@
 import java.security.ProtectionDomain;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -82,7 +81,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 +91,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/spi/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.spi;
+
+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/Whitelist.java renamed to modules/lang-painless/src/main/java/org/elasticsearch/painless/spi/Whitelist.java

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-package org.elasticsearch.painless;
+package org.elasticsearch.painless.spi;
 
 import java.util.Collections;
 import java.util.List;
@@ -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 renamed to modules/lang-painless/src/main/java/org/elasticsearch/painless/spi/WhitelistLoader.java

@@ -17,14 +17,16 @@
  * under the License.
  */
 
-package org.elasticsearch.painless;
+package org.elasticsearch.painless.spi;
 
 import java.io.InputStreamReader;
 import java.io.LineNumberReader;
 import java.lang.reflect.Constructor;
 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";
 };